diff --git a/apps/shared/OpenClawKit/Sources/OpenClawProtocol/GatewayModels.swift b/apps/shared/OpenClawKit/Sources/OpenClawProtocol/GatewayModels.swift index 9d80f3529648..cb604d845df9 100644 --- a/apps/shared/OpenClawKit/Sources/OpenClawProtocol/GatewayModels.swift +++ b/apps/shared/OpenClawKit/Sources/OpenClawProtocol/GatewayModels.swift @@ -425,6 +425,222 @@ public struct ErrorShape: Codable, Sendable { } } +public struct GatewaySuspendTaskBlocker: Codable, Sendable { + public let taskid: String + public let status: String + public let runtime: AnyCodable + public let runid: String? + public let label: String? + public let title: String? + + public init( + taskid: String, + status: String, + runtime: AnyCodable, + runid: String? = nil, + label: String? = nil, + title: String? = nil) + { + self.taskid = taskid + self.status = status + self.runtime = runtime + self.runid = runid + self.label = label + self.title = title + } + + private enum CodingKeys: String, CodingKey { + case taskid = "taskId" + case status + case runtime + case runid = "runId" + case label + case title + } +} + +public struct GatewaySuspendBlocker: Codable, Sendable { + public let kind: AnyCodable + public let count: Int + public let message: String + public let task: GatewaySuspendTaskBlocker? + + public init( + kind: AnyCodable, + count: Int, + message: String, + task: GatewaySuspendTaskBlocker? = nil) + { + self.kind = kind + self.count = count + self.message = message + self.task = task + } + + private enum CodingKeys: String, CodingKey { + case kind + case count + case message + case task + } +} + +public struct GatewaySuspendPrepareParams: Codable, Sendable { + public let requestid: String + + public init( + requestid: String) + { + self.requestid = requestid + } + + private enum CodingKeys: String, CodingKey { + case requestid = "requestId" + } +} + +public struct GatewaySuspendPrepareBusyResult: Codable, Sendable { + public let status: String + public let reason: AnyCodable + public let retryafterms: Int + public let activecount: Int + public let blockers: [GatewaySuspendBlocker] + + public init( + status: String, + reason: AnyCodable, + retryafterms: Int, + activecount: Int, + blockers: [GatewaySuspendBlocker]) + { + self.status = status + self.reason = reason + self.retryafterms = retryafterms + self.activecount = activecount + self.blockers = blockers + } + + private enum CodingKeys: String, CodingKey { + case status + case reason + case retryafterms = "retryAfterMs" + case activecount = "activeCount" + case blockers + } +} + +public struct GatewaySuspendPrepareReadyResult: Codable, Sendable { + public let status: String + public let suspensionid: String + public let expiresatms: Int + public let activecount: Int + public let blockers: [GatewaySuspendBlocker] + + public init( + status: String, + suspensionid: String, + expiresatms: Int, + activecount: Int, + blockers: [GatewaySuspendBlocker]) + { + self.status = status + self.suspensionid = suspensionid + self.expiresatms = expiresatms + self.activecount = activecount + self.blockers = blockers + } + + private enum CodingKeys: String, CodingKey { + case status + case suspensionid = "suspensionId" + case expiresatms = "expiresAtMs" + case activecount = "activeCount" + case blockers + } +} + +public struct GatewaySuspendStatusParams: Codable, Sendable { + public let suspensionid: String + + public init( + suspensionid: String) + { + self.suspensionid = suspensionid + } + + private enum CodingKeys: String, CodingKey { + case suspensionid = "suspensionId" + } +} + +public struct GatewaySuspendStatusRunningResult: Codable, Sendable { + public let status: String + + public init( + status: String) + { + self.status = status + } + + private enum CodingKeys: String, CodingKey { + case status + } +} + +public struct GatewaySuspendStatusReadyResult: Codable, Sendable { + public let status: String + public let expiresatms: Int + + public init( + status: String, + expiresatms: Int) + { + self.status = status + self.expiresatms = expiresatms + } + + private enum CodingKeys: String, CodingKey { + case status + case expiresatms = "expiresAtMs" + } +} + +public struct GatewaySuspendResumeParams: Codable, Sendable { + public let suspensionid: String + + public init( + suspensionid: String) + { + self.suspensionid = suspensionid + } + + private enum CodingKeys: String, CodingKey { + case suspensionid = "suspensionId" + } +} + +public struct GatewaySuspendResumeResult: Codable, Sendable { + public let ok: Bool + public let status: String + public let resumed: Bool + + public init( + ok: Bool, + status: String, + resumed: Bool) + { + self.ok = ok + self.status = status + self.resumed = resumed + } + + private enum CodingKeys: String, CodingKey { + case ok + case status + case resumed + } +} + public struct EnvironmentSummary: Codable, Sendable { public let id: String public let type: String @@ -9874,6 +10090,68 @@ public struct ShutdownEvent: Codable, Sendable { } } +public enum GatewaySuspendPrepareResult: Codable, Sendable { + case busy(GatewaySuspendPrepareBusyResult) + case ready(GatewaySuspendPrepareReadyResult) + + private enum CodingKeys: String, CodingKey { + case discriminator = "status" + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + let discriminator = try container.decode(String.self, forKey: .discriminator) + switch discriminator { + case "busy": self = try .busy(GatewaySuspendPrepareBusyResult(from: decoder)) + case "ready": self = try .ready(GatewaySuspendPrepareReadyResult(from: decoder)) + default: + throw DecodingError.dataCorruptedError( + forKey: .discriminator, + in: container, + debugDescription: "Unknown GatewaySuspendPrepareResult discriminator value" + ) + } + } + + public func encode(to encoder: Encoder) throws { + switch self { + case .busy(let value): try value.encode(to: encoder) + case .ready(let value): try value.encode(to: encoder) + } + } +} + +public enum GatewaySuspendStatusResult: Codable, Sendable { + case running(GatewaySuspendStatusRunningResult) + case ready(GatewaySuspendStatusReadyResult) + + private enum CodingKeys: String, CodingKey { + case discriminator = "status" + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + let discriminator = try container.decode(String.self, forKey: .discriminator) + switch discriminator { + case "running": self = try .running(GatewaySuspendStatusRunningResult(from: decoder)) + case "ready": self = try .ready(GatewaySuspendStatusReadyResult(from: decoder)) + default: + throw DecodingError.dataCorruptedError( + forKey: .discriminator, + in: container, + debugDescription: "Unknown GatewaySuspendStatusResult discriminator value" + ) + } + } + + public func encode(to encoder: Encoder) throws { + switch self { + case .running(let value): try value.encode(to: encoder) + case .ready(let value): try value.encode(to: encoder) + } + } +} + public enum PluginCatalogInstallAction: Codable, Sendable { case clawhub(PluginCatalogClawHubInstall) case official(PluginCatalogOfficialInstall) diff --git a/docs/docs_map.md b/docs/docs_map.md index 63e846e4c41a..65d36fa127d4 100644 --- a/docs/docs_map.md +++ b/docs/docs_map.md @@ -3339,6 +3339,7 @@ Do not edit it by hand; run `pnpm docs:map:gen`. - Headings: - H2: What is available today - H2: Recommended path + - H2: Cooperative host suspension - H2: App code vs plugin code - H2: Related diff --git a/docs/gateway/external-apps.md b/docs/gateway/external-apps.md index 59832d73a054..115419b3250b 100644 --- a/docs/gateway/external-apps.md +++ b/docs/gateway/external-apps.md @@ -50,6 +50,94 @@ terminal result. For durable conversation state, use the `sessions.*` methods. For UI integrations, subscribe to Gateway events and render only the event families your app understands. +## Cooperative host suspension + +Hosting controllers that freeze or snapshot a running process can use the +host-neutral suspension handshake: + +1. Stop admitting external ingress controlled by the host. +2. Call `gateway.suspend.prepare` with a stable, unique `requestId`. +3. If the response is `busy`, keep the process running and retry later. +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. + +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. + +The RPC contract is: + +- `gateway.suspend.prepare` — `operator.admin`; params + `{ "requestId": "stable-host-operation-id" }` +- `gateway.suspend.status` — `operator.read`; params + `{ "suspensionId": "id-from-prepare" }` +- `gateway.suspend.resume` — `operator.admin`; params + `{ "suspensionId": "id-from-prepare" }` + +IDs are trimmed, must contain a non-whitespace character, and are limited to +128 characters. A busy prepare result has `status: "busy"`, `reason`, +`retryAfterMs`, `activeCount`, and `blockers`. A ready result has this shape: + +```json +{ + "status": "ready", + "suspensionId": "2c3f...", + "expiresAtMs": 1770000000000, + "activeCount": 0, + "blockers": [] +} +``` + +Status returns `{"status":"running"}` or a ready result with `expiresAtMs`. +Resume returns `{"ok":true,"status":"running","resumed":true}`; repeating it +after a successful resume returns `resumed: false`. + +A competing request ID or transient scheduler-resume failure returns retryable +`UNAVAILABLE` with `retryAfterMs`. During scheduler recovery, prepare, status, +and resume all return that error, the Gateway remains not-ready and +fail-closed, and the host must not freeze or snapshot it. OpenClaw retries the +scheduler automatically and reopens admission only after recovery succeeds. A +mismatched resume ID returns `INVALID_REQUEST`. Prepare shares the Gateway's +control-plane write budget of three attempts per minute; honor the returned +retry delay. WebSocket clients are bucketed by device and IP. Admin HTTP +controllers are bucketed by resolved client IP, so controllers behind one +proxy can share a budget. + +Preparation is refuse-only: OpenClaw closes new root/session/command admission, +pauses automatic cron ticks, and inspects work synchronously. If anything is +active, it resumes the scheduler and reopens admission before returning +`busy`; it does not interrupt or drain that work. A ready lease lasts two +minutes. Repeating `prepare` with the same `requestId` renews it; expiry resumes +the scheduler before reopening admission. +Restart emission that becomes due during a ready lease waits until the lease +resumes; an in-flight restart makes preparation return `busy`. + +While ready, `/healthz` remains live and `/readyz` returns `503`. Local or +authenticated readiness responses include `gateway-draining`; unauthenticated +remote probes receive only `{ "ready": false }`. The HTTP health probe, +suspension methods on existing WebSocket connections, and an already-enabled +Admin HTTP RPC route remain available. Other RPCs return retryable +`UNAVAILABLE`. Built-in HTTP user-work routes, including OpenAI-compatible +APIs, tool/session operations, node watches, and configured hooks, return +`503` with `error.code: "gateway_unavailable"`. + +This handshake does not persist incoming messages, stop third-party channel +transports, or control the hosting platform. The host must fence its ingress +before preparation and remains responsible for wake, snapshot/freeze, and +stop. `activeCount` is the aggregate tracked-work count, while `blockers` +contains the non-zero category counts and bounded task details. This is not a +general process-quiescence barrier. Channel health, maintenance, cache refresh, +plugin-owned HTTP routes, and plugin-owned background work can remain active. +The hosting platform must freeze or snapshot the full process tree and its +filesystem consistently; unregistered work cannot be proven idle by this first +contract. + ## App code vs plugin code Use Gateway RPC when code lives outside OpenClaw: diff --git a/docs/gateway/protocol.md b/docs/gateway/protocol.md index 948815d0a9ad..25ac45af4b92 100644 --- a/docs/gateway/protocol.md +++ b/docs/gateway/protocol.md @@ -348,6 +348,7 @@ methods. Treat this as feature discovery, not a full enumeration of - `system-event` appends a system event and can update/broadcast presence context. - `last-heartbeat` returns the latest persisted heartbeat event. - `set-heartbeats` toggles heartbeat processing on the gateway. + - `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. diff --git a/docs/plugins/admin-http-rpc.md b/docs/plugins/admin-http-rpc.md index ee6299819194..ab82eac297f1 100644 --- a/docs/plugins/admin-http-rpc.md +++ b/docs/plugins/admin-http-rpc.md @@ -168,7 +168,7 @@ HTTP status follows the error code: - discovery: `commands.list` Returns the HTTP RPC method names allowed by this plugin. -- gateway: `health`, `status`, `logs.tail`, `usage.status`, `usage.cost`, `gateway.restart.request` +- gateway: `health`, `status`, `logs.tail`, `usage.status`, `usage.cost`, `gateway.restart.request`, `gateway.suspend.prepare`, `gateway.suspend.status`, `gateway.suspend.resume` - config: `config.get`, `config.schema`, `config.schema.lookup`, `config.set`, `config.patch`, `config.apply` - channels: `channels.status`, `channels.start`, `channels.stop`, `channels.logout` - web: `web.login.start`, `web.login.wait` @@ -209,11 +209,11 @@ Shared-token WebSocket clients without a trusted device identity cannot self-dec `400 INVALID_REQUEST` -: The request body is not valid JSON, the `method` field is missing, or the method is not in the plugin allowlist. +: The request body is not valid JSON, the `method` field is missing, the method is not in the plugin allowlist, or a suspension resume ID does not match the active lease. `503 UNAVAILABLE` -: The Gateway method handler is unavailable. Check Gateway logs and retry after the Gateway finishes startup. +: The Gateway method is starting, rate-limited, suspended, or waiting on a competing suspension/resume operation. Inspect `error.details` when present and honor `error.retryAfterMs` before retrying. ## Related diff --git a/extensions/admin-http-rpc/src/handler.test.ts b/extensions/admin-http-rpc/src/handler.test.ts index dadd0edc1a92..272e0e39bd83 100644 --- a/extensions/admin-http-rpc/src/handler.test.ts +++ b/extensions/admin-http-rpc/src/handler.test.ts @@ -133,6 +133,27 @@ describe("admin-http-rpc plugin handler", () => { }, ); + it.each([ + ["gateway.suspend.prepare", { requestId: "host-request-1" }], + ["gateway.suspend.status", { suspensionId: "suspension-1" }], + ["gateway.suspend.resume", { suspensionId: "suspension-1" }], + ] as const)("dispatches suspension method %s through Admin HTTP", async (method, params) => { + dispatchGatewayMethod.mockResolvedValueOnce({ + ok: true, + payload: { status: "ok" }, + }); + + const result = await invoke({ id: "suspension", method, params }); + + expect(dispatchGatewayMethod).toHaveBeenCalledWith(method, params); + expect(result.captured.statusCode).toBe(200); + expect(result.json).toEqual({ + id: "suspension", + ok: true, + payload: { status: "ok" }, + }); + }); + it("rejects methods outside the admin HTTP RPC allowlist", async () => { const result = await invoke({ id: "bad", method: "sessions.send" }); diff --git a/extensions/admin-http-rpc/src/methods.ts b/extensions/admin-http-rpc/src/methods.ts index 45e8d400cb0e..b675f94bd619 100644 --- a/extensions/admin-http-rpc/src/methods.ts +++ b/extensions/admin-http-rpc/src/methods.ts @@ -10,6 +10,9 @@ const ADMIN_HTTP_RPC_ALLOWED_METHOD_GROUPS = { "usage.status", "usage.cost", "gateway.restart.request", + "gateway.suspend.prepare", + "gateway.suspend.status", + "gateway.suspend.resume", ], discovery: ["commands.list"], config: [ diff --git a/packages/gateway-client/src/client.ts b/packages/gateway-client/src/client.ts index 90fed8911326..7f214cf6124a 100644 --- a/packages/gateway-client/src/client.ts +++ b/packages/gateway-client/src/client.ts @@ -1,12 +1,5 @@ // Gateway Client module implements client behavior. import { randomUUID } from "node:crypto"; -import type { - ConnectParams, - ErrorShape, - EventFrame, - HelloOk, - RequestFrame, -} from "@openclaw/gateway-protocol"; import { GATEWAY_CLIENT_MODES, GATEWAY_CLIENT_NAMES, @@ -22,8 +15,13 @@ import { type ConnectErrorRecoveryAdvice, } from "@openclaw/gateway-protocol/connect-error-details"; import { + type ConnectParams, + type ErrorShape, + type EventFrame, + type HelloOk, isGatewayEventFrame, isGatewayResponseFrame, + type RequestFrame, } from "@openclaw/gateway-protocol/frame-guards"; import { resolveGatewayStartupRetryAfterMs } from "@openclaw/gateway-protocol/startup-unavailable"; import { MIN_CLIENT_PROTOCOL_VERSION, PROTOCOL_VERSION } from "@openclaw/gateway-protocol/version"; diff --git a/packages/gateway-protocol/src/frame-guards.ts b/packages/gateway-protocol/src/frame-guards.ts index 8df10e77dcc4..4ee846673f1c 100644 --- a/packages/gateway-protocol/src/frame-guards.ts +++ b/packages/gateway-protocol/src/frame-guards.ts @@ -1,4 +1,13 @@ -import type { EventFrame, ResponseFrame } from "./schema/types.js"; +import type { EventFrame, ResponseFrame } from "./schema/frames.js"; +export type { + ConnectParams, + ErrorShape, + EventFrame, + GatewayFrame, + HelloOk, + RequestFrame, + ResponseFrame, +} from "./schema/frames.js"; function isRecord(value: unknown): value is Record { return Boolean(value) && typeof value === "object" && !Array.isArray(value); diff --git a/packages/gateway-protocol/src/gateway-suspend.test.ts b/packages/gateway-protocol/src/gateway-suspend.test.ts new file mode 100644 index 000000000000..9e6b07ce62bf --- /dev/null +++ b/packages/gateway-protocol/src/gateway-suspend.test.ts @@ -0,0 +1,73 @@ +import { describe, expect, it } from "vitest"; +import { + validateGatewaySuspendPrepareParams, + validateGatewaySuspendPrepareResult, + validateGatewaySuspendResumeResult, + validateGatewaySuspendStatusResult, +} from "./index.js"; + +describe("gateway suspension protocol", () => { + it("keeps prepare params closed and bounded", () => { + expect(validateGatewaySuspendPrepareParams({ requestId: "host-request" })).toBe(true); + expect(validateGatewaySuspendPrepareParams({ requestId: " " })).toBe(false); + expect(validateGatewaySuspendPrepareParams({ requestId: "host-request", extra: true })).toBe( + false, + ); + }); + + it("validates busy and ready prepare results", () => { + expect( + validateGatewaySuspendPrepareResult({ + status: "busy", + reason: "active-work", + retryAfterMs: 20_000, + activeCount: 2, + blockers: [ + { kind: "queue", count: 1, message: "one queued operation" }, + { + kind: "task", + count: 1, + message: "one active task", + task: { taskId: "task-1", status: "running", runtime: "subagent" }, + }, + ], + }), + ).toBe(true); + expect( + validateGatewaySuspendPrepareResult({ + status: "ready", + suspensionId: "suspension-id", + expiresAtMs: 123, + activeCount: 0, + blockers: [], + }), + ).toBe(true); + }); + + it("validates status and resume results", () => { + expect(validateGatewaySuspendStatusResult({ status: "running" })).toBe(true); + expect(validateGatewaySuspendStatusResult({ status: "ready", expiresAtMs: 123 })).toBe(true); + expect( + validateGatewaySuspendResumeResult({ ok: true, status: "running", resumed: false }), + ).toBe(true); + expect( + validateGatewaySuspendResumeResult({ + ok: true, + status: "running", + resumed: false, + warnings: [], + }), + ).toBe(false); + }); + + it("keeps scheduler recovery on the error frame instead of success results", () => { + const recovering = { + status: "recovering", + reason: "scheduler-resume-failed", + retryAfterMs: 1_000, + }; + + expect(validateGatewaySuspendPrepareResult(recovering)).toBe(false); + expect(validateGatewaySuspendStatusResult(recovering)).toBe(false); + }); +}); diff --git a/packages/gateway-protocol/src/index.ts b/packages/gateway-protocol/src/index.ts index 63552bd38fdd..ea66305a8570 100644 --- a/packages/gateway-protocol/src/index.ts +++ b/packages/gateway-protocol/src/index.ts @@ -9,6 +9,8 @@ export { type ClawHubTrustErrorDetails, } from "./clawhub-trust-error-details.js"; import { Compile, type Validator as TypeBoxValidator } from "typebox/compile"; +import type { ValidationError } from "./validation-errors.js"; +export { formatValidationErrors, type ValidationError } from "./validation-errors.js"; import { type AgentEvent, AgentEventSchema, @@ -185,6 +187,26 @@ import { UpdateStatusParamsSchema, type ConnectParams, ConnectParamsSchema, + type GatewaySuspendBlocker, + GatewaySuspendBlockerSchema, + type GatewaySuspendPrepareParams, + GatewaySuspendPrepareBusyResultSchema, + GatewaySuspendPrepareParamsSchema, + GatewaySuspendPrepareReadyResultSchema, + type GatewaySuspendPrepareResult, + GatewaySuspendPrepareResultSchema, + type GatewaySuspendResumeParams, + GatewaySuspendResumeParamsSchema, + type GatewaySuspendResumeResult, + GatewaySuspendResumeResultSchema, + type GatewaySuspendStatusParams, + GatewaySuspendStatusParamsSchema, + GatewaySuspendStatusReadyResultSchema, + type GatewaySuspendStatusResult, + GatewaySuspendStatusResultSchema, + GatewaySuspendStatusRunningResultSchema, + type GatewaySuspendTaskBlocker, + GatewaySuspendTaskBlockerSchema, type CronAddParams, CronAddParamsSchema, type CronAddResult, @@ -664,20 +686,6 @@ import { WorktreesBranchesResultSchema, } from "./schema.js"; -/** Normalized validation error shape exposed by every protocol validator. */ -export type ValidationError = { - /** Failed schema keyword, when the validator can report one. */ - keyword?: string; - /** JSON-pointer path to the failing data location. */ - instancePath?: string; - /** JSON-pointer path to the failing schema location. */ - schemaPath?: string; - /** Validator-specific keyword parameters for richer diagnostics. */ - params?: Record; - /** Human-readable validation message. */ - message?: string; -}; - /** Runtime validator shape shared by gateway clients and server handlers. */ export type ProtocolValidator = ((data: unknown) => data is T) & { /** Last validation errors, matching Ajv-style caller expectations. */ @@ -728,6 +736,24 @@ function lazyCompile(schema: unknown): ProtocolValidator { // constants so call sites can pair validation with the wire contract directly. export const validateCommandsListParams = lazyCompile(CommandsListParamsSchema); export const validateConnectParams = lazyCompile(ConnectParamsSchema); +export const validateGatewaySuspendPrepareParams = lazyCompile( + GatewaySuspendPrepareParamsSchema, +); +export const validateGatewaySuspendPrepareResult = lazyCompile( + GatewaySuspendPrepareResultSchema, +); +export const validateGatewaySuspendStatusParams = lazyCompile( + GatewaySuspendStatusParamsSchema, +); +export const validateGatewaySuspendStatusResult = lazyCompile( + GatewaySuspendStatusResultSchema, +); +export const validateGatewaySuspendResumeParams = lazyCompile( + GatewaySuspendResumeParamsSchema, +); +export const validateGatewaySuspendResumeResult = lazyCompile( + GatewaySuspendResumeResultSchema, +); export const validateRequestFrame = lazyCompile(RequestFrameSchema); export const validateResponseFrame = lazyCompile(ResponseFrameSchema); export const validateEventFrame = lazyCompile(EventFrameSchema); @@ -1199,77 +1225,22 @@ export const validateWebLoginStartParams = lazyCompile(WebLoginStartParamsSchema); export const validateWebLoginWaitParams = lazyCompile(WebLoginWaitParamsSchema); -function firstStringParam(value: unknown): string | undefined { - if (typeof value === "string" && value.trim()) { - return value; - } - if (Array.isArray(value)) { - return value.find( - (entry): entry is string => typeof entry === "string" && entry.trim().length > 0, - ); - } - return undefined; -} - -/** Convert validator errors into compact operator-facing failure text. */ -export function formatValidationErrors(errors: ValidationError[] | null | undefined) { - if (!errors?.length) { - return "unknown validation error"; - } - - const parts: string[] = []; - - for (const err of errors) { - const keyword = typeof err?.keyword === "string" ? err.keyword : ""; - const instancePath = typeof err?.instancePath === "string" ? err.instancePath : ""; - - if (keyword === "additionalProperties") { - const additionalProperty = - firstStringParam(err?.params?.additionalProperty) ?? - firstStringParam(err?.params?.additionalProperties); - if (additionalProperty) { - const where = instancePath ? `at ${instancePath}` : "at root"; - parts.push(`${where}: unexpected property '${additionalProperty}'`); - continue; - } - } - if (keyword === "required") { - const missingProperty = - firstStringParam(err?.params?.missingProperty) ?? - firstStringParam(err?.params?.requiredProperties); - if (missingProperty) { - const where = instancePath ? `at ${instancePath}: ` : ""; - parts.push(`${where}must have required property '${missingProperty}'`); - continue; - } - } - - const failingKeyword = - typeof err?.params?.failingKeyword === "string" ? err.params.failingKeyword : ""; - // TypeBox reports conditional required-property misses through if/then - // keywords, which otherwise hide the actionable missing-property context. - const message = - keyword === "then" || (keyword === "if" && failingKeyword === "then") - ? "must have required conditional properties" - : typeof err?.message === "string" && err.message.trim() - ? err.message - : "validation error"; - const where = instancePath ? `at ${instancePath}: ` : ""; - parts.push(`${where}${message}`); - } - - // De-dupe while preserving order. - const unique = uniqueStrings(parts.filter((part) => part.trim())); - if (!unique.length) { - return "unknown validation error"; - } - return unique.join("; "); -} - // Schema exports stay explicit to make additions/removals reviewable as public // protocol surface changes. export { ConnectParamsSchema, + GatewaySuspendTaskBlockerSchema, + GatewaySuspendBlockerSchema, + GatewaySuspendPrepareBusyResultSchema, + GatewaySuspendPrepareParamsSchema, + GatewaySuspendPrepareReadyResultSchema, + GatewaySuspendPrepareResultSchema, + GatewaySuspendStatusReadyResultSchema, + GatewaySuspendStatusRunningResultSchema, + GatewaySuspendStatusParamsSchema, + GatewaySuspendStatusResultSchema, + GatewaySuspendResumeParamsSchema, + GatewaySuspendResumeResultSchema, GATEWAY_SERVER_CAPS, HelloOkSchema, RequestFrameSchema, @@ -1577,6 +1548,14 @@ export { export type { GatewayFrame, ConnectParams, + GatewaySuspendTaskBlocker, + GatewaySuspendBlocker, + GatewaySuspendPrepareParams, + GatewaySuspendPrepareResult, + GatewaySuspendStatusParams, + GatewaySuspendStatusResult, + GatewaySuspendResumeParams, + GatewaySuspendResumeResult, HelloOk, RequestFrame, ResponseFrame, @@ -1869,9 +1848,6 @@ export type { SessionsGroupsDeleteParams, SessionsGroupsMutationResult, }; -function uniqueStrings(values: string[]): string[] { - return [...new Set(values)]; -} // The protocol package cannot import core session types. This local structural // result mirrors the wire contract and keeps the package independent of src/. diff --git a/packages/gateway-protocol/src/schema.ts b/packages/gateway-protocol/src/schema.ts index 0acbc3ade890..5bf85248c734 100644 --- a/packages/gateway-protocol/src/schema.ts +++ b/packages/gateway-protocol/src/schema.ts @@ -20,6 +20,7 @@ export * from "./schema/environments.js"; export * from "./schema/exec-approvals.js"; export * from "./schema/devices.js"; export * from "./schema/frames.js"; +export * from "./schema/gateway-suspend.js"; export * from "./schema/logs-chat.js"; export * from "./schema/nodes.js"; export * from "./schema/protocol-schemas.js"; diff --git a/packages/gateway-protocol/src/schema/error-codes.ts b/packages/gateway-protocol/src/schema/error-codes.ts index d93647bf8652..080d7d976640 100644 --- a/packages/gateway-protocol/src/schema/error-codes.ts +++ b/packages/gateway-protocol/src/schema/error-codes.ts @@ -1,5 +1,5 @@ // Gateway Protocol schema module defines protocol validation shapes. -import type { ErrorShape } from "./types.js"; +import type { ErrorShape } from "./frames.js"; /** Gateway JSON-RPC style error codes shared by clients and server handlers. */ export const ErrorCodes = { diff --git a/packages/gateway-protocol/src/schema/frames.ts b/packages/gateway-protocol/src/schema/frames.ts index ea431909dde2..2f3a1561755a 100644 --- a/packages/gateway-protocol/src/schema/frames.ts +++ b/packages/gateway-protocol/src/schema/frames.ts @@ -1,4 +1,5 @@ // Gateway Protocol schema module defines protocol validation shapes. +import type { Static } from "typebox"; import { Type } from "typebox"; import { GatewayClientIdSchema, GatewayClientModeSchema, NonEmptyString } from "./primitives.js"; import { SnapshotSchema, StateVersionSchema } from "./snapshot.js"; @@ -213,3 +214,13 @@ export const GatewayFrameSchema = Type.Union( [RequestFrameSchema, ResponseFrameSchema, EventFrameSchema], { discriminator: "type" }, ); + +// Frame types are owner-local because they cross the public client/plugin SDK. +// Keeping them off the aggregate registry avoids retaining every RPC schema. +export type ConnectParams = Static; +export type HelloOk = Static; +export type ErrorShape = Static; +export type RequestFrame = Static; +export type ResponseFrame = Static; +export type EventFrame = Static; +export type GatewayFrame = Static; diff --git a/packages/gateway-protocol/src/schema/gateway-suspend.ts b/packages/gateway-protocol/src/schema/gateway-suspend.ts new file mode 100644 index 000000000000..b8b37569391b --- /dev/null +++ b/packages/gateway-protocol/src/schema/gateway-suspend.ts @@ -0,0 +1,108 @@ +// Gateway Protocol schemas for cooperative host suspension. +import { Type } from "typebox"; + +const SuspensionTokenSchema = Type.String({ minLength: 1, maxLength: 128, pattern: "\\S" }); +const CountSchema = Type.Integer({ minimum: 0 }); + +export const GatewaySuspendTaskBlockerSchema = Type.Object( + { + taskId: Type.String(), + status: Type.Literal("running"), + runtime: Type.Union([ + Type.Literal("subagent"), + Type.Literal("acp"), + Type.Literal("cli"), + Type.Literal("cron"), + ]), + runId: Type.Optional(Type.String()), + label: Type.Optional(Type.String()), + title: Type.Optional(Type.String()), + }, + { additionalProperties: false }, +); + +export const GatewaySuspendBlockerSchema = Type.Object( + { + kind: Type.Union([ + Type.Literal("queue"), + Type.Literal("reply"), + Type.Literal("embedded-run"), + Type.Literal("cron-run"), + Type.Literal("task"), + Type.Literal("root-request"), + Type.Literal("session-admission"), + Type.Literal("session-mutation"), + Type.Literal("chat-run"), + Type.Literal("queued-turn"), + Type.Literal("terminal-persistence"), + Type.Literal("terminal-session"), + ]), + count: CountSchema, + message: Type.String(), + task: Type.Optional(GatewaySuspendTaskBlockerSchema), + }, + { additionalProperties: false }, +); + +export const GatewaySuspendPrepareParamsSchema = Type.Object( + { requestId: SuspensionTokenSchema }, + { additionalProperties: false }, +); + +export const GatewaySuspendPrepareBusyResultSchema = Type.Object( + { + status: Type.Literal("busy"), + reason: Type.Union([Type.Literal("active-work"), Type.Literal("gateway-draining")]), + retryAfterMs: CountSchema, + activeCount: CountSchema, + blockers: Type.Array(GatewaySuspendBlockerSchema), + }, + { additionalProperties: false }, +); + +export const GatewaySuspendPrepareReadyResultSchema = Type.Object( + { + status: Type.Literal("ready"), + suspensionId: SuspensionTokenSchema, + expiresAtMs: CountSchema, + activeCount: CountSchema, + blockers: Type.Array(GatewaySuspendBlockerSchema), + }, + { additionalProperties: false }, +); + +export const GatewaySuspendPrepareResultSchema = Type.Union([ + GatewaySuspendPrepareBusyResultSchema, + GatewaySuspendPrepareReadyResultSchema, +]); + +export const GatewaySuspendStatusParamsSchema = Type.Object( + { suspensionId: SuspensionTokenSchema }, + { additionalProperties: false }, +); + +export const GatewaySuspendStatusRunningResultSchema = Type.Object( + { status: Type.Literal("running") }, + { additionalProperties: false }, +); + +export const GatewaySuspendStatusReadyResultSchema = Type.Object( + { status: Type.Literal("ready"), expiresAtMs: CountSchema }, + { additionalProperties: false }, +); + +export const GatewaySuspendStatusResultSchema = Type.Union([ + GatewaySuspendStatusRunningResultSchema, + GatewaySuspendStatusReadyResultSchema, +]); + +export const GatewaySuspendResumeParamsSchema = GatewaySuspendStatusParamsSchema; + +export const GatewaySuspendResumeResultSchema = Type.Object( + { + ok: Type.Literal(true), + status: Type.Literal("running"), + resumed: Type.Boolean(), + }, + { additionalProperties: false }, +); diff --git a/packages/gateway-protocol/src/schema/protocol-schemas.ts b/packages/gateway-protocol/src/schema/protocol-schemas.ts index 1746f7ae43b5..47d171c6e64c 100644 --- a/packages/gateway-protocol/src/schema/protocol-schemas.ts +++ b/packages/gateway-protocol/src/schema/protocol-schemas.ts @@ -221,6 +221,20 @@ import { ShutdownEventSchema, TickEventSchema, } from "./frames.js"; +import { + GatewaySuspendBlockerSchema, + GatewaySuspendPrepareBusyResultSchema, + GatewaySuspendPrepareParamsSchema, + GatewaySuspendPrepareReadyResultSchema, + GatewaySuspendPrepareResultSchema, + GatewaySuspendResumeParamsSchema, + GatewaySuspendResumeResultSchema, + GatewaySuspendStatusReadyResultSchema, + GatewaySuspendStatusRunningResultSchema, + GatewaySuspendStatusParamsSchema, + GatewaySuspendStatusResultSchema, + GatewaySuspendTaskBlockerSchema, +} from "./gateway-suspend.js"; import { ChatAbortedEventSchema, ChatAbortParamsSchema, @@ -420,6 +434,18 @@ export const ProtocolSchemas = { StateVersion: StateVersionSchema, Snapshot: SnapshotSchema, ErrorShape: ErrorShapeSchema, + GatewaySuspendTaskBlocker: GatewaySuspendTaskBlockerSchema, + GatewaySuspendBlocker: GatewaySuspendBlockerSchema, + GatewaySuspendPrepareParams: GatewaySuspendPrepareParamsSchema, + GatewaySuspendPrepareBusyResult: GatewaySuspendPrepareBusyResultSchema, + GatewaySuspendPrepareReadyResult: GatewaySuspendPrepareReadyResultSchema, + GatewaySuspendPrepareResult: GatewaySuspendPrepareResultSchema, + GatewaySuspendStatusParams: GatewaySuspendStatusParamsSchema, + GatewaySuspendStatusRunningResult: GatewaySuspendStatusRunningResultSchema, + GatewaySuspendStatusReadyResult: GatewaySuspendStatusReadyResultSchema, + GatewaySuspendStatusResult: GatewaySuspendStatusResultSchema, + GatewaySuspendResumeParams: GatewaySuspendResumeParamsSchema, + GatewaySuspendResumeResult: GatewaySuspendResumeResultSchema, // Environment and agent-facing control RPC payloads. EnvironmentStatus: EnvironmentStatusSchema, diff --git a/packages/gateway-protocol/src/schema/sessions.ts b/packages/gateway-protocol/src/schema/sessions.ts index a7df1908ed4f..9e404afeecbf 100644 --- a/packages/gateway-protocol/src/schema/sessions.ts +++ b/packages/gateway-protocol/src/schema/sessions.ts @@ -1,4 +1,5 @@ // Gateway Protocol schema module defines protocol validation shapes. +import type { Static } from "typebox"; import { Type } from "typebox"; import { PluginJsonValueSchema } from "./plugins.js"; import { NonEmptyString, SessionLabelString } from "./primitives.js"; @@ -406,6 +407,7 @@ export const SessionsPatchParamsSchema = Type.Object( }, { additionalProperties: false }, ); +export type SessionsPatchParams = Static; /** Updates or clears one plugin namespace value on a session record. */ export const SessionsPluginPatchParamsSchema = Type.Object( diff --git a/packages/gateway-protocol/src/schema/types.ts b/packages/gateway-protocol/src/schema/types.ts index 12f71e9a54dd..ef054fb28a9d 100644 --- a/packages/gateway-protocol/src/schema/types.ts +++ b/packages/gateway-protocol/src/schema/types.ts @@ -1,28 +1,30 @@ /** * Static TypeScript types derived from the canonical gateway protocol schemas. * - * Keep aliases wired through `ProtocolSchemas` so validators, runtime schemas, - * and exported compile-time types cannot drift apart. + * Owner-local schema modules export hot public types directly. The remaining + * aliases stay wired through `ProtocolSchemas` so validators, runtime schemas, + * and compile-time types cannot drift apart. */ import type { Static } from "typebox"; -import { ProtocolSchemas } from "./protocol-schemas.js"; +import type { ProtocolSchemas } from "./protocol-schemas.js"; /** Stable schema names registered in the protocol schema registry. */ type ProtocolSchemaName = keyof typeof ProtocolSchemas; /** Inferred TypeScript type for a named TypeBox protocol schema. */ type SchemaType = Static<(typeof ProtocolSchemas)[TName]>; -/** Connection handshake, envelope, snapshot, and shared error wire types. */ -export type ConnectParams = SchemaType<"ConnectParams">; -export type HelloOk = SchemaType<"HelloOk">; -export type RequestFrame = SchemaType<"RequestFrame">; -export type ResponseFrame = SchemaType<"ResponseFrame">; -export type EventFrame = SchemaType<"EventFrame">; -export type GatewayFrame = SchemaType<"GatewayFrame">; +/** Snapshot and shared state wire types. */ export type Snapshot = SchemaType<"Snapshot">; export type PresenceEntry = SchemaType<"PresenceEntry">; -export type ErrorShape = SchemaType<"ErrorShape">; export type StateVersion = SchemaType<"StateVersion">; +export type GatewaySuspendTaskBlocker = SchemaType<"GatewaySuspendTaskBlocker">; +export type GatewaySuspendBlocker = SchemaType<"GatewaySuspendBlocker">; +export type GatewaySuspendPrepareParams = SchemaType<"GatewaySuspendPrepareParams">; +export type GatewaySuspendPrepareResult = SchemaType<"GatewaySuspendPrepareResult">; +export type GatewaySuspendStatusParams = SchemaType<"GatewaySuspendStatusParams">; +export type GatewaySuspendStatusResult = SchemaType<"GatewaySuspendStatusResult">; +export type GatewaySuspendResumeParams = SchemaType<"GatewaySuspendResumeParams">; +export type GatewaySuspendResumeResult = SchemaType<"GatewaySuspendResumeResult">; /** Environment status RPC payloads used by CLI and Control UI surfaces. */ export type EnvironmentStatus = SchemaType<"EnvironmentStatus">; @@ -113,7 +115,6 @@ export type SessionsSendParams = SchemaType<"SessionsSendParams">; export type SessionsMessagesSubscribeParams = SchemaType<"SessionsMessagesSubscribeParams">; export type SessionsMessagesUnsubscribeParams = SchemaType<"SessionsMessagesUnsubscribeParams">; export type SessionsAbortParams = SchemaType<"SessionsAbortParams">; -export type SessionsPatchParams = SchemaType<"SessionsPatchParams">; export type SessionsPluginPatchParams = SchemaType<"SessionsPluginPatchParams">; export type SessionsPluginPatchResult = SchemaType<"SessionsPluginPatchResult">; export type SessionsResetParams = SchemaType<"SessionsResetParams">; diff --git a/packages/gateway-protocol/src/validation-errors.ts b/packages/gateway-protocol/src/validation-errors.ts new file mode 100644 index 000000000000..ab017f5af881 --- /dev/null +++ b/packages/gateway-protocol/src/validation-errors.ts @@ -0,0 +1,76 @@ +/** Normalized validation error shape exposed by every protocol validator. */ +export type ValidationError = { + /** Failed schema keyword, when the validator can report one. */ + keyword?: string; + /** JSON-pointer path to the failing data location. */ + instancePath?: string; + /** JSON-pointer path to the failing schema location. */ + schemaPath?: string; + /** Validator-specific keyword parameters for richer diagnostics. */ + params?: Record; + /** Human-readable validation message. */ + message?: string; +}; + +function firstStringParam(value: unknown): string | undefined { + if (typeof value === "string" && value.trim()) { + return value; + } + if (Array.isArray(value)) { + return value.find( + (entry): entry is string => typeof entry === "string" && entry.trim().length > 0, + ); + } + return undefined; +} + +/** Convert validator errors into compact operator-facing failure text. */ +export function formatValidationErrors(errors: ValidationError[] | null | undefined) { + if (!errors?.length) { + return "unknown validation error"; + } + + const parts: string[] = []; + + for (const err of errors) { + const keyword = typeof err?.keyword === "string" ? err.keyword : ""; + const instancePath = typeof err?.instancePath === "string" ? err.instancePath : ""; + + if (keyword === "additionalProperties") { + const additionalProperty = + firstStringParam(err?.params?.additionalProperty) ?? + firstStringParam(err?.params?.additionalProperties); + if (additionalProperty) { + const where = instancePath ? `at ${instancePath}` : "at root"; + parts.push(`${where}: unexpected property '${additionalProperty}'`); + continue; + } + } + if (keyword === "required") { + const missingProperty = + firstStringParam(err?.params?.missingProperty) ?? + firstStringParam(err?.params?.requiredProperties); + if (missingProperty) { + const where = instancePath ? `at ${instancePath}: ` : ""; + parts.push(`${where}must have required property '${missingProperty}'`); + continue; + } + } + + const failingKeyword = + typeof err?.params?.failingKeyword === "string" ? err.params.failingKeyword : ""; + // TypeBox reports conditional required-property misses through if/then + // keywords, which otherwise hide the actionable missing-property context. + const message = + keyword === "then" || (keyword === "if" && failingKeyword === "then") + ? "must have required conditional properties" + : typeof err?.message === "string" && err.message.trim() + ? err.message + : "validation error"; + const where = instancePath ? `at ${instancePath}: ` : ""; + parts.push(`${where}${message}`); + } + + const unique = [...new Set(parts.filter((part) => part.trim()))]; + return unique.length > 0 ? unique.join("; ") : "unknown validation error"; +} diff --git a/scripts/check-plugin-sdk-exports.mjs b/scripts/check-plugin-sdk-exports.mjs index c0f8d1aba7c4..11af131c1fb2 100755 --- a/scripts/check-plugin-sdk-exports.mjs +++ b/scripts/check-plugin-sdk-exports.mjs @@ -49,6 +49,7 @@ const exportSet = new Set(exportedNames); const requiredRuntimeShimEntries = ["compat.js", "root-alias.cjs"]; const forbiddenPublicDeclarationSpecifiers = ["@openclaw/llm-core"]; +const FORBIDDEN_PUBLIC_PROTOCOL_REGISTRY_RE = /\bdeclare\s+const\s+ProtocolSchemas(?:\$\d+)?\b/u; const RELATIVE_DECLARATION_SPECIFIER_RE = /\b(?:from|import)\s*(?:\(\s*)?["']([^"']+)["']/gu; const requiredSubpathExports = { "secret-input-runtime": [ @@ -140,6 +141,12 @@ while (declarationQueue.length > 0) { } declarationPaths.add(dtsPath); const dtsContent = readFileSync(dtsPath, "utf8"); + if (FORBIDDEN_PUBLIC_PROTOCOL_REGISTRY_RE.test(dtsContent)) { + console.error( + `FORBIDDEN PUBLIC DTS REGISTRY: ${relative(resolve(scriptDir, ".."), dtsPath)} retains ProtocolSchemas`, + ); + missing += 1; + } for (const match of dtsContent.matchAll(RELATIVE_DECLARATION_SPECIFIER_RE)) { const specifier = match[1]; if (!specifier?.startsWith(".")) { diff --git a/src/agents/main-session-restart-recovery.test.ts b/src/agents/main-session-restart-recovery.test.ts index ba9f919eec32..6572d2e232a6 100644 --- a/src/agents/main-session-restart-recovery.test.ts +++ b/src/agents/main-session-restart-recovery.test.ts @@ -14,6 +14,11 @@ import { registerAgentRunContext, resetAgentRunContextForTest, } from "../infra/agent-events.js"; +import { + getActiveGatewayRootWorkCount, + resetGatewayWorkAdmission, + tryBeginGatewaySuspendAdmission, +} from "../process/gateway-work-admission.js"; import { INTERNAL_RUNTIME_CONTEXT_BEGIN, INTERNAL_RUNTIME_CONTEXT_END, @@ -24,6 +29,7 @@ import { markStartupOrphanedMainSessionsForRecovery, recoverStartupOrphanedMainSessions, recoverRestartAbortedMainSessions, + scheduleRestartAbortedMainSessionRecovery, } from "./main-session-restart-recovery.js"; import type { SessionLockInspection } from "./session-write-lock.js"; @@ -36,10 +42,12 @@ let tmpDir: string; beforeEach(async () => { vi.clearAllMocks(); resetAgentRunContextForTest(); + resetGatewayWorkAdmission(); tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-main-restart-recovery-")); }); afterEach(async () => { + resetGatewayWorkAdmission(); await fs.rm(tmpDir, { recursive: true, force: true }); }); @@ -1282,6 +1290,54 @@ describe("main-session-restart-recovery", () => { expect(customStore["agent:main:main"]?.abortedLastRun).toBe(false); }); + it("admits each scheduled recovery attempt as independent root work", async () => { + const sessionsDir = await makeSessionsDir(); + await writeStore(sessionsDir, { + "agent:main:main": { + sessionId: "main-session", + updatedAt: Date.now() - 10_000, + status: "running", + abortedLastRun: true, + pendingFinalDelivery: true, + pendingFinalDeliveryText: "interrupted response", + }, + }); + + const suspensionRef: { + current: ReturnType; + } = { current: null }; + vi.mocked(callGateway) + .mockImplementationOnce(async () => { + expect(getActiveGatewayRootWorkCount()).toBe(1); + suspensionRef.current = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspensionRef.current?.commit()).toBe(true); + throw new Error("retry after suspension"); + }) + .mockImplementationOnce(async () => { + expect(getActiveGatewayRootWorkCount()).toBe(1); + return { runId: "run-resumed" }; + }); + + scheduleRestartAbortedMainSessionRecovery({ + delayMs: 0, + maxRetries: 2, + stateDir: tmpDir, + }); + + await vi.waitFor(() => { + expect(callGateway).toHaveBeenCalledOnce(); + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); + expect(suspensionRef.current?.release()).toBe(true); + + await vi.waitFor(() => { + expect(callGateway).toHaveBeenCalledTimes(2); + const store = readSessionStoreForTest(path.join(sessionsDir, "sessions.json")); + expect(store["agent:main:main"]?.abortedLastRun).toBe(false); + }); + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); + it("fails marked sessions whose transcript tail cannot be resumed", async () => { const sessionsDir = await makeSessionsDir(); await writeStore(sessionsDir, { diff --git a/src/agents/main-session-restart-recovery.ts b/src/agents/main-session-restart-recovery.ts index 73e75389b2f1..2eec32ff8983 100644 --- a/src/agents/main-session-restart-recovery.ts +++ b/src/agents/main-session-restart-recovery.ts @@ -26,6 +26,7 @@ import { listAgentRunsForSession, } from "../infra/agent-events.js"; import { createSubsystemLogger } from "../logging/subsystem.js"; +import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js"; import { CommandLane } from "../process/lanes.js"; import { isAcpSessionKey, @@ -956,12 +957,17 @@ export function scheduleRestartAbortedMainSessionRecovery( const startupRecoveryCutoffMs = Date.now(); const runRecoveryAttempt = (attempt: number, delay: number) => { - void recoverStartupOrphanedMainSessions({ - cfg: params.cfg, - stateDir: params.stateDir, - resumedSessionKeys, - updatedBeforeMs: startupRecoveryCutoffMs, - }) + // Delayed retries outlive startup; each attempt must independently block + // host suspension while it reads and rewrites recovery session state. + void runWithGatewayIndependentRootWorkAdmission( + async () => + await recoverStartupOrphanedMainSessions({ + cfg: params.cfg, + stateDir: params.stateDir, + resumedSessionKeys, + updatedBeforeMs: startupRecoveryCutoffMs, + }), + ) .then((result) => { if (result.failed > 0 && attempt < maxRetries) { scheduleAttempt(attempt + 1, delay * RETRY_BACKOFF_MULTIPLIER); diff --git a/src/agents/subagent-orphan-recovery.test.ts b/src/agents/subagent-orphan-recovery.test.ts index 7f71ad86d75d..56bf01289bae 100644 --- a/src/agents/subagent-orphan-recovery.test.ts +++ b/src/agents/subagent-orphan-recovery.test.ts @@ -4,6 +4,11 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import * as sessions from "../config/sessions.js"; import * as gateway from "../gateway/call.js"; import * as sessionUtils from "../gateway/session-transcript-readers.js"; +import { + getActiveGatewayRootWorkCount, + resetGatewayWorkAdmission, + tryBeginGatewaySuspendAdmission, +} from "../process/gateway-work-admission.js"; import { resolveInternalSessionEffectsTranscriptPath } from "./internal-session-effects.js"; import * as announceDelivery from "./subagent-announce-delivery.js"; import { @@ -130,9 +135,11 @@ describe("subagent-orphan-recovery", () => { beforeEach(() => { vi.useFakeTimers(); vi.clearAllMocks(); + resetGatewayWorkAdmission(); }); afterEach(() => { + resetGatewayWorkAdmission(); vi.useRealTimers(); vi.restoreAllMocks(); }); @@ -630,7 +637,17 @@ describe("subagent-orphan-recovery", () => { abortedLastRun: true, }, }); - vi.mocked(gateway.callGateway).mockRejectedValue(new Error("service restart")); + const admittedRootCounts: number[] = []; + vi.mocked(gateway.callGateway).mockImplementation(async () => { + admittedRootCounts.push(getActiveGatewayRootWorkCount()); + throw new Error("service restart"); + }); + vi.mocked(subagentRegistrySteerRuntime.finalizeInterruptedSubagentRun).mockImplementation( + async () => { + admittedRootCounts.push(getActiveGatewayRootWorkCount()); + return 1; + }, + ); const activeRuns = createActiveRuns(createTestRunRecord()); @@ -642,10 +659,13 @@ describe("subagent-orphan-recovery", () => { await vi.advanceTimersByTimeAsync(1); await Promise.resolve(); + expect(getActiveGatewayRootWorkCount()).toBe(0); await vi.advanceTimersByTimeAsync(2); await Promise.resolve(); expect(subagentRegistrySteerRuntime.finalizeInterruptedSubagentRun).toHaveBeenCalledOnce(); + expect(admittedRootCounts).toEqual([1, 1, 1]); + expect(getActiveGatewayRootWorkCount()).toBe(0); const finalizeParams = requireRecord( firstCallParam( vi.mocked(subagentRegistrySteerRuntime.finalizeInterruptedSubagentRun).mock.calls, @@ -658,4 +678,29 @@ describe("subagent-orphan-recovery", () => { expect(finalizeParams.error).toContain("Automatic recovery failed after 2 attempts"); expect(finalizeParams.error).toContain("service restart"); }); + + it("waits for suspension to reopen before mutating an orphaned session", async () => { + mockSingleAbortedSession(); + vi.mocked(gateway.callGateway).mockResolvedValue({ runId: "resumed-run" }); + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + + scheduleOrphanRecovery({ + getActiveRuns: () => createActiveRuns(createTestRunRecord()), + delayMs: 1, + maxRetries: 0, + }); + + await vi.advanceTimersByTimeAsync(1); + expect(gateway.callGateway).not.toHaveBeenCalled(); + expect(sessions.updateSessionStore).not.toHaveBeenCalled(); + expect(getActiveGatewayRootWorkCount()).toBe(0); + + expect(suspension?.release()).toBe(true); + await vi.advanceTimersByTimeAsync(0); + + expect(gateway.callGateway).toHaveBeenCalledOnce(); + expect(sessions.updateSessionStore).toHaveBeenCalledOnce(); + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); }); diff --git a/src/agents/subagent-orphan-recovery.ts b/src/agents/subagent-orphan-recovery.ts index cb07103d40b3..92c6c947e822 100644 --- a/src/agents/subagent-orphan-recovery.ts +++ b/src/agents/subagent-orphan-recovery.ts @@ -23,6 +23,7 @@ import { callGateway } from "../gateway/call.js"; import { readSessionMessagesAsync } from "../gateway/session-transcript-readers.js"; import { formatErrorMessage } from "../infra/errors.js"; import { createSubsystemLogger } from "../logging/subsystem.js"; +import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js"; import { truncateUtf16Safe } from "../utils.js"; import { resolveInternalSessionEffectsTranscriptPath } from "./internal-session-effects.js"; import { @@ -437,49 +438,48 @@ export function scheduleOrphanRecovery(params: { const resumedSessionKeys = new Set(); const attemptRecovery = (attempt: number, delay: number) => { setTimeout(() => { - void recoverOrphanedSubagentSessions({ - ...params, - resumedSessionKeys, - }) - .then((result) => { - if (result.failed > 0 && attempt < maxRetries) { - const nextDelay = delay * RETRY_BACKOFF_MULTIPLIER; - log.info( - `orphan recovery had ${result.failed} failure(s); retrying in ${nextDelay}ms (attempt ${attempt + 1}/${maxRetries})`, - ); - attemptRecovery(attempt + 1, nextDelay); - return; - } - if (result.failedRuns.length === 0) { - return; - } - const attempts = attempt + 1; - void Promise.allSettled( - result.failedRuns.map((run) => - finalizeInterruptedSubagentRun({ - runId: run.runId, - childSessionKey: run.childSessionKey, - error: buildRecoveryFailureMessage({ - attempts, - error: run.error, - }), - }), - ), - ); - }) - .catch((err: unknown) => { - if (attempt < maxRetries) { - const nextDelay = delay * RETRY_BACKOFF_MULTIPLIER; - log.warn( - `scheduled orphan recovery failed: ${String(err)}; retrying in ${nextDelay}ms (attempt ${attempt + 1}/${maxRetries})`, - ); - attemptRecovery(attempt + 1, nextDelay); - } else { - log.warn( - `scheduled orphan recovery failed after ${maxRetries} retries: ${String(err)}`, - ); - } + // Every delayed/retry scan owns a fresh root lease. Keep terminal + // mutation in the same lease so suspension cannot become ready mid-attempt. + void runWithGatewayIndependentRootWorkAdmission(async () => { + const result = await recoverOrphanedSubagentSessions({ + ...params, + resumedSessionKeys, }); + if (result.failed > 0 && attempt < maxRetries) { + const nextDelay = delay * RETRY_BACKOFF_MULTIPLIER; + log.info( + `orphan recovery had ${result.failed} failure(s); retrying in ${nextDelay}ms (attempt ${attempt + 1}/${maxRetries})`, + ); + attemptRecovery(attempt + 1, nextDelay); + return; + } + if (result.failedRuns.length === 0) { + return; + } + const attempts = attempt + 1; + await Promise.allSettled( + result.failedRuns.map((run) => + finalizeInterruptedSubagentRun({ + runId: run.runId, + childSessionKey: run.childSessionKey, + error: buildRecoveryFailureMessage({ + attempts, + error: run.error, + }), + }), + ), + ); + }).catch((err: unknown) => { + if (attempt < maxRetries) { + const nextDelay = delay * RETRY_BACKOFF_MULTIPLIER; + log.warn( + `scheduled orphan recovery failed: ${String(err)}; retrying in ${nextDelay}ms (attempt ${attempt + 1}/${maxRetries})`, + ); + attemptRecovery(attempt + 1, nextDelay); + } else { + log.warn(`scheduled orphan recovery failed after ${maxRetries} retries: ${String(err)}`); + } + }); }, delay).unref?.(); }; diff --git a/src/agents/subagent-registry-lifecycle.test.ts b/src/agents/subagent-registry-lifecycle.test.ts index fdf391ff8381..a7077b0f2309 100644 --- a/src/agents/subagent-registry-lifecycle.test.ts +++ b/src/agents/subagent-registry-lifecycle.test.ts @@ -2,6 +2,11 @@ // detached task status, and resource retirement around child-run endings. import { beforeEach, describe, expect, it, vi } from "vitest"; import type { CallGatewayOptions } from "../gateway/call.js"; +import { + getActiveGatewayRootWorkCount, + markGatewayRestartDraining, + resetGatewayWorkAdmission, +} from "../process/gateway-work-admission.js"; import { SUBAGENT_KILL_TASK_ERROR } from "../tasks/detached-task-runtime-contract.js"; import { buildAnnounceIdFromChildRun, @@ -265,6 +270,7 @@ async function runNoReplyMirrorScenario(params: { describe("subagent registry lifecycle hardening", () => { beforeEach(() => { + resetGatewayWorkAdmission(); vi.clearAllMocks(); taskExecutorMocks.completeTaskRunByRunId.mockReset(); taskExecutorMocks.failTaskRunByRunId.mockReset(); @@ -276,6 +282,131 @@ describe("subagent registry lifecycle hardening", () => { bundleMcpRuntimeMocks.retireSessionMcpRuntimeForSessionKey.mockResolvedValue(true); }); + it("keeps task finalization, resource retirement, and announce cleanup root-admitted", async () => { + const entry = createRunEntry({ expectsCompletionMessage: true }); + let releaseBrowserCleanup: (() => void) | undefined; + let releaseAnnounce: ((didAnnounce: boolean) => void) | undefined; + browserLifecycleCleanupMocks.cleanupBrowserSessionsForLifecycleEnd.mockImplementationOnce( + () => + new Promise((resolve) => { + releaseBrowserCleanup = resolve; + }), + ); + const runSubagentAnnounceFlow = vi.fn( + () => + new Promise((resolve) => { + releaseAnnounce = resolve; + }), + ); + const controller = createLifecycleController({ entry, runSubagentAnnounceFlow }); + + const completion = controller.completeSubagentRun({ + runId: entry.runId, + endedAt: 4_000, + outcome: { status: "ok" }, + reason: SUBAGENT_ENDED_REASON_COMPLETE, + triggerCleanup: true, + }); + + await vi.waitFor(() => + expect( + browserLifecycleCleanupMocks.cleanupBrowserSessionsForLifecycleEnd, + ).toHaveBeenCalledOnce(), + ); + expect(taskExecutorMocks.completeTaskRunByRunId).toHaveBeenCalledOnce(); + expect(getActiveGatewayRootWorkCount()).toBe(1); + + releaseBrowserCleanup?.(); + await vi.waitFor(() => expect(runSubagentAnnounceFlow).toHaveBeenCalledOnce()); + await completion; + expect(getActiveGatewayRootWorkCount()).toBe(1); + + releaseAnnounce?.(true); + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + expect(entry.cleanupCompletedAt).toBeTypeOf("number"); + }); + + it("keeps direct delete cleanup root-admitted until the gateway call settles", async () => { + const entry = createRunEntry({ cleanup: "delete", expectsCompletionMessage: false }); + const runs = new Map([[entry.runId, entry]]); + let releaseDelete: (() => void) | undefined; + gatewayMocks.callGateway.mockImplementation((opts) => { + if (opts.method !== "sessions.delete") { + return Promise.resolve({}); + } + return new Promise>((resolve) => { + releaseDelete = () => resolve({}); + }); + }); + const controller = createLifecycleController({ entry, runs }); + + await controller.completeSubagentRun({ + runId: entry.runId, + endedAt: 4_000, + outcome: { status: "ok" }, + reason: SUBAGENT_ENDED_REASON_COMPLETE, + triggerCleanup: true, + }); + await vi.waitFor(() => expect(releaseDelete).toBeTypeOf("function")); + expect(getActiveGatewayRootWorkCount()).toBe(1); + + releaseDelete?.(); + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + expect(runs.has(entry.runId)).toBe(false); + }); + + it("retries a cleanup handoff rejected by restart drain", async () => { + vi.useFakeTimers(); + try { + const entry = createRunEntry({ expectsCompletionMessage: true }); + let releaseBrowserCleanup: (() => void) | undefined; + browserLifecycleCleanupMocks.cleanupBrowserSessionsForLifecycleEnd.mockImplementationOnce( + () => + new Promise((resolve) => { + releaseBrowserCleanup = resolve; + }), + ); + const runSubagentAnnounceFlow = vi.fn(async () => true); + const resumeSubagentRun = vi.fn((runId: string) => { + controller.startSubagentAnnounceCleanupFlow(runId, entry); + }); + const controller = createLifecycleController({ + entry, + resumeSubagentRun, + runSubagentAnnounceFlow, + }); + + const completion = controller.completeSubagentRun({ + runId: entry.runId, + endedAt: 4_000, + outcome: { status: "ok" }, + reason: SUBAGENT_ENDED_REASON_COMPLETE, + triggerCleanup: true, + }); + await vi.waitFor(() => expect(releaseBrowserCleanup).toBeTypeOf("function")); + markGatewayRestartDraining(); + releaseBrowserCleanup?.(); + await completion; + await vi.waitFor(() => + expect(runtimeMocks.log).toHaveBeenCalledWith( + expect.stringContaining("subagent cleanup admission failed"), + ), + ); + expect(runSubagentAnnounceFlow).not.toHaveBeenCalled(); + expect(entry.cleanupHandled).toBe(true); + + resetGatewayWorkAdmission(); + await vi.advanceTimersByTimeAsync(1_000); + await vi.waitFor(() => expect(runSubagentAnnounceFlow).toHaveBeenCalledOnce()); + await vi.waitFor(() => expect(entry.cleanupCompletedAt).toBeTypeOf("number")); + expect(resumeSubagentRun).toHaveBeenCalledWith(entry.runId); + expect(getActiveGatewayRootWorkCount()).toBe(0); + } finally { + resetGatewayWorkAdmission(); + vi.useRealTimers(); + } + }); + it("does not reject completion when task finalization throws", async () => { const persist = vi.fn(); const persistOrThrow = vi.fn(); @@ -1985,6 +2116,53 @@ describe("subagent registry lifecycle hardening", () => { }); }); + it("keeps a late superseded-delivery retirement root-admitted", async () => { + const entry = createRunEntry({ expectsCompletionMessage: true, generation: 1 }); + const runs = new Map([[entry.runId, entry]]); + let onDeliveryResult: ((delivery: SubagentAnnounceDeliveryResult) => void) | undefined; + const runSubagentAnnounceFlow: LifecycleControllerParams["runSubagentAnnounceFlow"] = vi.fn( + async (announceParams) => { + onDeliveryResult = announceParams.onDeliveryResult; + return true; + }, + ); + let releaseRetirement = () => {}; + const retirementPending = new Promise((resolve) => { + releaseRetirement = resolve; + }); + const retireSupersededRun = vi.fn(async () => { + await retirementPending; + }); + const controller = createLifecycleController({ + entry, + runs, + retireSupersededRun, + runSubagentAnnounceFlow, + }); + + await controller.completeSubagentRun({ + runId: entry.runId, + endedAt: 4_000, + outcome: { status: "ok" }, + reason: SUBAGENT_ENDED_REASON_COMPLETE, + triggerCleanup: true, + }); + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + const newer = createRunEntry({ + runId: "run-2", + childSessionKey: entry.childSessionKey, + generation: 2, + }); + runs.set(newer.runId, newer); + + onDeliveryResult?.({ delivered: false, path: "none" }); + + await vi.waitFor(() => expect(retireSupersededRun).toHaveBeenCalledWith(entry.runId, entry)); + expect(getActiveGatewayRootWorkCount()).toBe(1); + releaseRetirement(); + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + }); + it("finalizes terminal visible-send failures without scheduling completion retry", async () => { const persist = vi.fn(); const entry = createRunEntry({ diff --git a/src/agents/subagent-registry-lifecycle.ts b/src/agents/subagent-registry-lifecycle.ts index 299a434165cf..ee2db88606dc 100644 --- a/src/agents/subagent-registry-lifecycle.ts +++ b/src/agents/subagent-registry-lifecycle.ts @@ -8,6 +8,10 @@ import { isSilentReplyText, SILENT_REPLY_TOKEN } from "../auto-reply/tokens.js"; import type { cleanupBrowserSessionsForLifecycleEnd } from "../browser-lifecycle-cleanup.js"; import type { callGateway as defaultCallGateway } from "../gateway/call.js"; import { formatErrorMessage, readErrorName } from "../infra/errors.js"; +import { + isGatewayRestartDraining, + runWithGatewayIndependentRootWorkAdmission, +} from "../process/gateway-work-admission.js"; import { defaultRuntime } from "../runtime.js"; import { emitSessionLifecycleEvent } from "../sessions/session-lifecycle-events.js"; import { extractTextFromChatContent } from "../shared/chat-content.js"; @@ -215,13 +219,43 @@ export function createSubagentRegistryLifecycleController(params: { }; }; - const scheduleResumeSubagentRun = (runId: string, entry: SubagentRunRecord, delayMs: number) => { + const scheduleResumeSubagentRun = ( + runId: string, + entry: SubagentRunRecord, + delayMs: number, + cleanupGeneration?: number, + ) => { const timer = setTimeout(() => { scheduledResumeTimers.delete(timer); - if (params.runs.get(runId) !== entry) { - return; - } - params.resumeSubagentRun(runId); + void runWithGatewayIndependentRootWorkAdmission(async () => { + if (params.runs.get(runId) !== entry) { + return; + } + if (cleanupGeneration !== undefined) { + if (!isCleanupAttemptCurrent(runId, entry, cleanupGeneration)) { + return; + } + entry.cleanupHandled = false; + params.persist(); + } + params.resumedRuns.delete(runId); + params.resumeSubagentRun(runId); + }).catch((err: unknown) => { + defaultRuntime.log(`[warn] subagent cleanup resume failed (${runId}): ${String(err)}`); + const current = params.runs.get(runId); + if ( + isGatewayRestartDraining() && + current === entry && + typeof current.cleanupCompletedAt !== "number" + ) { + scheduleResumeSubagentRun( + runId, + entry, + Math.max(delayMs, MIN_ANNOUNCE_RETRY_DELAY_MS), + cleanupGeneration, + ); + } + }); }, delayMs); timer.unref?.(); scheduledResumeTimers.add(timer); @@ -234,6 +268,49 @@ export function createSubagentRegistryLifecycleController(params: { scheduledResumeTimers.clear(); }; + const runDetachedCleanupAttempt = (args: { + runId: string; + entry: SubagentRunRecord; + cleanupGeneration: number; + run: () => Promise; + }) => { + // Completion makes the task projection non-blocking before delivery and + // cleanup finish. This independent lease bridges that handoff and owns the + // full detached attempt, including its final durable registry write. + void runWithGatewayIndependentRootWorkAdmission(async () => { + try { + await args.run(); + } catch (err) { + defaultRuntime.log( + `[warn] subagent cleanup finalize failed (${args.runId}): ${String(err)}`, + ); + const current = params.runs.get(args.runId); + if ( + !current || + current.cleanupCompletedAt || + !isCleanupAttemptCurrent(args.runId, args.entry, args.cleanupGeneration) + ) { + return; + } + current.cleanupHandled = false; + params.resumedRuns.delete(args.runId); + params.persist(); + } + }).catch((err: unknown) => { + defaultRuntime.log( + `[warn] subagent cleanup admission failed (${args.runId}): ${String(err)}`, + ); + if (isGatewayRestartDraining()) { + scheduleResumeSubagentRun( + args.runId, + args.entry, + MIN_ANNOUNCE_RETRY_DELAY_MS, + args.cleanupGeneration, + ); + } + }); + }; + const maskRunId = (runId: string): string => { const trimmed = runId.trim(); if (!trimmed) { @@ -803,6 +880,22 @@ export function createSubagentRegistryLifecycleController(params: { return true; }; + const retireSupersededCleanupInBackground = ( + runId: string, + entry: SubagentRunRecord, + generation: number, + ) => { + // Delivery callbacks are synchronous and may arrive after their announce + // attempt returns. Give the async retirement tail its own snapshot blocker. + void runWithGatewayIndependentRootWorkAdmission(async () => { + await retireSupersededCleanupIfNeeded(runId, entry, generation); + }).catch((error: unknown) => { + defaultRuntime.log( + `[warn] subagent superseded cleanup retirement failed (${runId}): ${String(error)}`, + ); + }); + }; + const isTerminalCallbackCurrent = ( runId: string, entry: SubagentRunRecord, @@ -840,20 +933,17 @@ export function createSubagentRegistryLifecycleController(params: { if (!beginSubagentCleanup(runId)) { continue; } - void finalizeResumedAnnounceGiveUp({ + runDetachedCleanupAttempt({ runId, entry, - reason: "expiry", - }).catch((error: unknown) => { - defaultRuntime.log( - `[warn] Subagent expiry finalize failed during deferred retry for run ${runId}: ${String(error)}`, - ); - const current = params.runs.get(runId); - if (!current || current.cleanupCompletedAt) { - return; - } - current.cleanupHandled = false; - params.persist(); + cleanupGeneration: cleanupGenerations.get(entry)!, + run: async () => { + await finalizeResumedAnnounceGiveUp({ + runId, + entry, + reason: "expiry", + }); + }, }); continue; } @@ -870,31 +960,46 @@ export function createSubagentRegistryLifecycleController(params: { preserveTranscript?: boolean; provisionalKill?: boolean; }) => { + const runCleanupTail = (label: string, run: () => Promise) => { + // These best-effort tails can outlive the durable registry transition, + // but they still mutate session-owned resources and must block snapshots. + void runWithGatewayIndependentRootWorkAdmission(run).catch((error: unknown) => { + defaultRuntime.log( + `[warn] subagent ${label} failed (${cleanupParams.runId}): ${String(error)}`, + ); + }); + }; if (!cleanupParams.preserveTranscript) { - void removeInternalSessionEffectsTranscript(cleanupParams.entry.execution?.transcriptFile); + runCleanupTail("transcript cleanup", async () => { + await removeInternalSessionEffectsTranscript(cleanupParams.entry.execution?.transcriptFile); + }); } if (cleanupParams.entry.spawnMode !== "session") { - void retireSessionMcpRuntimeForSessionKey({ - sessionKey: cleanupParams.entry.childSessionKey, - reason: "subagent-run-cleanup", - onError: (error, sessionId) => { - params.warn("failed to retire subagent bundle MCP runtime", { - error: buildSafeLifecycleErrorMeta(error), - sessionId, - runId: maskRunId(cleanupParams.runId), - childSessionKey: maskSessionKey(cleanupParams.entry.childSessionKey), - }); - }, + runCleanupTail("bundle MCP cleanup", async () => { + await retireSessionMcpRuntimeForSessionKey({ + sessionKey: cleanupParams.entry.childSessionKey, + reason: "subagent-run-cleanup", + onError: (error, sessionId) => { + params.warn("failed to retire subagent bundle MCP runtime", { + error: buildSafeLifecycleErrorMeta(error), + sessionId, + runId: maskRunId(cleanupParams.runId), + childSessionKey: maskSessionKey(cleanupParams.entry.childSessionKey), + }); + }, + }); }); } if (cleanupParams.cleanup === "delete") { params.clearPendingLifecycleError(cleanupParams.runId); if (!cleanupParams.provisionalKill) { - void params.notifyContextEngineSubagentEnded({ - childSessionKey: cleanupParams.entry.childSessionKey, - reason: "deleted", - agentDir: cleanupParams.entry.agentDir, - workspaceDir: cleanupParams.entry.workspaceDir, + runCleanupTail("context-engine cleanup", async () => { + await params.notifyContextEngineSubagentEnded({ + childSessionKey: cleanupParams.entry.childSessionKey, + reason: "deleted", + agentDir: cleanupParams.entry.agentDir, + workspaceDir: cleanupParams.entry.workspaceDir, + }); }); } params.runs.delete(cleanupParams.runId); @@ -903,11 +1008,13 @@ export function createSubagentRegistryLifecycleController(params: { return; } if (!cleanupParams.provisionalKill) { - void params.notifyContextEngineSubagentEnded({ - childSessionKey: cleanupParams.entry.childSessionKey, - reason: "completed", - agentDir: cleanupParams.entry.agentDir, - workspaceDir: cleanupParams.entry.workspaceDir, + runCleanupTail("context-engine cleanup", async () => { + await params.notifyContextEngineSubagentEnded({ + childSessionKey: cleanupParams.entry.childSessionKey, + reason: "completed", + agentDir: cleanupParams.entry.agentDir, + workspaceDir: cleanupParams.entry.workspaceDir, + }); }); } if ( @@ -1161,20 +1268,15 @@ export function createSubagentRegistryLifecycleController(params: { return false; } const cleanupGeneration = cleanupGenerations.get(entry)!; - void finalizeSubagentCleanup(runId, cleanup, true, cleanupGeneration, { - skipAnnounce: true, - }).catch((err: unknown) => { - defaultRuntime.log(`[warn] subagent cleanup finalize failed (${runId}): ${String(err)}`); - const current = params.runs.get(runId); - if ( - !current || - current.cleanupCompletedAt || - !isCleanupAttemptCurrent(runId, entry, cleanupGeneration) - ) { - return; - } - current.cleanupHandled = false; - params.persist(); + runDetachedCleanupAttempt({ + runId, + entry, + cleanupGeneration, + run: async () => { + await finalizeSubagentCleanup(runId, cleanup, true, cleanupGeneration, { + skipAnnounce: true, + }); + }, }); return true; } @@ -1184,52 +1286,45 @@ export function createSubagentRegistryLifecycleController(params: { const cleanupGeneration = cleanupGenerations.get(entry)!; const skipRequesterDelivery = entry.suppressCompletionDelivery === true; if (entry.expectsCompletionMessage === false || skipRequesterDelivery) { - void (async () => { - // This driver is detached. Yield once so synchronous successor - // registration can invalidate it before sessions.delete is submitted. - await Promise.resolve(); - if (!isCleanupAttemptCurrent(runId, entry, cleanupGeneration)) { - await retireSupersededCleanupIfNeeded(runId, entry, cleanupGeneration); - return; - } - if (cleanup === "delete") { - // This durable boundary prevents a late yield from reviving a run - // after deletion may already have reached the gateway. - entry.deleteCleanupDispatchedAt ??= Date.now(); - params.persist(); - await deleteSubagentSessionForCleanup({ - callGateway: params.callGateway, - childSessionKey: entry.childSessionKey, - spawnMode: entry.spawnMode, - onError: (error) => - params.warn("sessions.delete failed during subagent cleanup", { - error: buildSafeLifecycleErrorMeta(error), - runId: maskRunId(runId), - childSessionKey: maskSessionKey(entry.childSessionKey), - }), + runDetachedCleanupAttempt({ + runId, + entry, + cleanupGeneration, + run: async () => { + // This driver is detached. Yield once so synchronous successor + // registration can invalidate it before sessions.delete is submitted. + await Promise.resolve(); + if (!isCleanupAttemptCurrent(runId, entry, cleanupGeneration)) { + await retireSupersededCleanupIfNeeded(runId, entry, cleanupGeneration); + return; + } + if (cleanup === "delete") { + // This durable boundary prevents a late yield from reviving a run + // after deletion may already have reached the gateway. + entry.deleteCleanupDispatchedAt ??= Date.now(); + params.persist(); + await deleteSubagentSessionForCleanup({ + callGateway: params.callGateway, + childSessionKey: entry.childSessionKey, + spawnMode: entry.spawnMode, + onError: (error) => + params.warn("sessions.delete failed during subagent cleanup", { + error: buildSafeLifecycleErrorMeta(error), + runId: maskRunId(runId), + childSessionKey: maskSessionKey(entry.childSessionKey), + }), + }); + } + if (!isCleanupAttemptCurrent(runId, entry, cleanupGeneration)) { + await retireSupersededCleanupIfNeeded(runId, entry, cleanupGeneration); + return; + } + await finalizeSubagentCleanup(runId, cleanup, true, cleanupGeneration, { + skipAnnounce: true, + skipDeliveryStatus: true, + skipRequesterDelivery, }); - } - if (!isCleanupAttemptCurrent(runId, entry, cleanupGeneration)) { - await retireSupersededCleanupIfNeeded(runId, entry, cleanupGeneration); - return; - } - await finalizeSubagentCleanup(runId, cleanup, true, cleanupGeneration, { - skipAnnounce: true, - skipDeliveryStatus: true, - skipRequesterDelivery, - }); - })().catch((err: unknown) => { - defaultRuntime.log(`[warn] subagent cleanup finalize failed (${runId}): ${String(err)}`); - const current = params.runs.get(runId); - if ( - !current || - current.cleanupCompletedAt || - !isCleanupAttemptCurrent(runId, entry, cleanupGeneration) - ) { - return; - } - current.cleanupHandled = false; - params.persist(); + }, }); return true; } @@ -1253,93 +1348,87 @@ export function createSubagentRegistryLifecycleController(params: { if (!didAnnounce && latestDeliveryError) { ensureDeliveryState(entry).lastError = latestDeliveryError; } - void finalizeSubagentCleanup( + await finalizeSubagentCleanup( runId, cleanup, didAnnounce || shouldCreditPriorDelivery, cleanupGeneration, - ).catch((err: unknown) => { - defaultRuntime.log(`[warn] subagent cleanup finalize failed (${runId}): ${String(err)}`); - const current = params.runs.get(runId); - if ( - !current || - current.cleanupCompletedAt || - !isCleanupAttemptCurrent(runId, entry, cleanupGeneration) - ) { - return; - } - current.cleanupHandled = false; - params.persist(); - }); + ); }; - void params - .runSubagentAnnounceFlow({ - childSessionKey: pendingPayload.childSessionKey, - childRunId: pendingPayload.childRunId, - requesterSessionKey: pendingPayload.requesterSessionKey, - requesterOrigin, - requesterDisplayKey: pendingPayload.requesterDisplayKey, - task: pendingPayload.task, - timeoutMs: params.subagentAnnounceTimeoutMs, - cleanup, - roundOneReply: pendingPayload.frozenResultText ?? undefined, - fallbackReply: pendingPayload.fallbackFrozenResultText ?? undefined, - waitForCompletion: false, - startedAt: pendingPayload.startedAt, - endedAt: pendingPayload.endedAt, - label: pendingPayload.label, - outcome: pendingPayload.outcome, - spawnMode: pendingPayload.spawnMode, - expectsCompletionMessage: pendingPayload.expectsCompletionMessage, - wakeOnDescendantSettle: pendingPayload.wakeOnDescendantSettle === true, - onBeforeDeleteChildSession: - cleanup === "delete" - ? () => { - if (!isCleanupAttemptCurrent(runId, entry, cleanupGeneration)) { - return false; - } - // Announce owns delete submission; fence late yields at the - // exact handoff instead of when cleanup merely starts. - entry.deleteCleanupDispatchedAt ??= Date.now(); - params.persist(); - return true; + const announceParams: Parameters[0] = { + childSessionKey: pendingPayload.childSessionKey, + childRunId: pendingPayload.childRunId, + requesterSessionKey: pendingPayload.requesterSessionKey, + requesterOrigin, + requesterDisplayKey: pendingPayload.requesterDisplayKey, + task: pendingPayload.task, + timeoutMs: params.subagentAnnounceTimeoutMs, + cleanup, + roundOneReply: pendingPayload.frozenResultText ?? undefined, + fallbackReply: pendingPayload.fallbackFrozenResultText ?? undefined, + waitForCompletion: false, + startedAt: pendingPayload.startedAt, + endedAt: pendingPayload.endedAt, + label: pendingPayload.label, + outcome: pendingPayload.outcome, + spawnMode: pendingPayload.spawnMode, + expectsCompletionMessage: pendingPayload.expectsCompletionMessage, + wakeOnDescendantSettle: pendingPayload.wakeOnDescendantSettle === true, + onBeforeDeleteChildSession: + cleanup === "delete" + ? () => { + if (!isCleanupAttemptCurrent(runId, entry, cleanupGeneration)) { + return false; } - : undefined, - onDeliveryResult: (delivery) => { - if (!isCleanupAttemptCurrent(runId, entry, cleanupGeneration)) { - void retireSupersededCleanupIfNeeded(runId, entry, cleanupGeneration); - return; - } - recordAnnounceDeliveryResult(entry, delivery); - if (delivery.delivered) { - const deliveryState = ensureDeliveryState(entry); - if (deliveryState.lastError !== undefined) { - deliveryState.lastError = undefined; + // Announce owns delete submission; fence late yields at the + // exact handoff instead of when cleanup merely starts. + entry.deleteCleanupDispatchedAt ??= Date.now(); params.persist(); + return true; } - latestDeliveryError = undefined; - return; - } - if (delivery.path === "none") { - ensureDeliveryState(entry).lastDropReason = "sink_unavailable"; - } - latestDeliveryError = formatAnnounceDeliveryError(delivery); - if (ensureDeliveryState(entry).lastError !== latestDeliveryError) { - ensureDeliveryState(entry).lastError = latestDeliveryError; + : undefined, + onDeliveryResult: (delivery) => { + if (!isCleanupAttemptCurrent(runId, entry, cleanupGeneration)) { + retireSupersededCleanupInBackground(runId, entry, cleanupGeneration); + return; + } + recordAnnounceDeliveryResult(entry, delivery); + if (delivery.delivered) { + const deliveryState = ensureDeliveryState(entry); + if (deliveryState.lastError !== undefined) { + deliveryState.lastError = undefined; params.persist(); } - }, - }) - .then((didAnnounce) => { - void finalizeAnnounceCleanup(didAnnounce); - }) - .catch((error: unknown) => { - defaultRuntime.log( - `[warn] Subagent announce flow failed during cleanup for run ${runId}: ${String(error)}`, - ); - void finalizeAnnounceCleanup(false); - }); + latestDeliveryError = undefined; + return; + } + if (delivery.path === "none") { + ensureDeliveryState(entry).lastDropReason = "sink_unavailable"; + } + latestDeliveryError = formatAnnounceDeliveryError(delivery); + if (ensureDeliveryState(entry).lastError !== latestDeliveryError) { + ensureDeliveryState(entry).lastError = latestDeliveryError; + params.persist(); + } + }, + }; + runDetachedCleanupAttempt({ + runId, + entry, + cleanupGeneration, + run: async () => { + let didAnnounce = false; + try { + didAnnounce = await params.runSubagentAnnounceFlow(announceParams); + } catch (error) { + defaultRuntime.log( + `[warn] Subagent announce flow failed during cleanup for run ${runId}: ${String(error)}`, + ); + } + await finalizeAnnounceCleanup(didAnnounce); + }, + }); return true; }; @@ -1356,7 +1445,7 @@ export function createSubagentRegistryLifecycleController(params: { completionSnapshot?: { resultText: string | null; capturedAt: number }; }; - const completeSubagentRun = async (completeParams: CompleteSubagentRunParams) => { + const completeSubagentRunAttempt = async (completeParams: CompleteSubagentRunParams) => { const releaseCompletionLock = await acquireTerminalCompletionLock(completeParams.runId); let entry: SubagentRunRecord | undefined; let terminalGeneration = 0; @@ -1843,6 +1932,16 @@ export function createSubagentRegistryLifecycleController(params: { startSubagentAnnounceCleanupFlow(completeParams.runId, entry); }; + const completeSubagentRun = async (completeParams: CompleteSubagentRunParams) => { + // Task finalization can make the run disappear from suspension blockers + // before browser/MCP retirement and cleanup delivery hand off. Own this + // entire transition as an independent root so that boundary stays atomic. + // Callers can detach while retaining parent ALS, so nesting is intentional. + await runWithGatewayIndependentRootWorkAdmission(async () => { + await completeSubagentRunAttempt(completeParams); + }); + }; + return { clearScheduledResumeTimers, completeCleanupBookkeeping, diff --git a/src/agents/subagent-registry-run-manager.ts b/src/agents/subagent-registry-run-manager.ts index b3fab6444fb4..8481495abb9b 100644 --- a/src/agents/subagent-registry-run-manager.ts +++ b/src/agents/subagent-registry-run-manager.ts @@ -7,6 +7,7 @@ import { getRuntimeConfig } from "../config/config.js"; import type { OpenClawConfig } from "../config/types.openclaw.js"; import { callGateway } from "../gateway/call.js"; import { createSubsystemLogger } from "../logging/subsystem.js"; +import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js"; import { SUBAGENT_KILL_TASK_ERROR, type DetachedTaskFindResult, @@ -990,18 +991,28 @@ export function createSubagentRunManager(params: { finalizeKilledTask(pending.entry, pending.endedAt); } for (const entry of entriesByChildSessionKey.values()) { - void persistSubagentSessionTiming(entry, { - isCurrentGeneration: () => currentRunOwnsSession(entry), + // Task finalization removes the suspension blocker before these session-owned + // writes finish. Join them under one independent root so snapshots stay atomic. + void runWithGatewayIndependentRootWorkAdmission(async () => { + await Promise.all([ + persistSubagentSessionTiming(entry, { + isCurrentGeneration: () => currentRunOwnsSession(entry), + }).catch((err: unknown) => { + log.warn("failed to persist killed subagent session timing", { + err, + runId: entry.runId, + childSessionKey: entry.childSessionKey, + }); + }), + shouldDeleteAttachments(entry) ? safeRemoveAttachmentsDir(entry) : Promise.resolve(), + ]); }).catch((err: unknown) => { - log.warn("failed to persist killed subagent session timing", { + log.warn("failed to run killed subagent cleanup tail", { err, runId: entry.runId, childSessionKey: entry.childSessionKey, }); }); - if (shouldDeleteAttachments(entry)) { - void safeRemoveAttachmentsDir(entry); - } params.completeCleanupBookkeeping({ runId: entry.runId, entry, diff --git a/src/agents/subagent-registry.test.ts b/src/agents/subagent-registry.test.ts index 57eb73326fb8..ff1d84ef11b1 100644 --- a/src/agents/subagent-registry.test.ts +++ b/src/agents/subagent-registry.test.ts @@ -11,6 +11,11 @@ import type { SessionEntryPatchOptions, } from "../config/sessions/session-accessor.js"; import type { AgentEventPayload } from "../infra/agent-events.js"; +import { + getActiveGatewayRootWorkCount, + markGatewayRestartDraining, + resetGatewayWorkAdmission, +} from "../process/gateway-work-admission.js"; import { SUBAGENT_KILL_TASK_ERROR } from "../tasks/detached-task-runtime-contract.js"; import { createRunningTaskRun, @@ -255,6 +260,7 @@ describe("subagent registry seam flow", () => { }); beforeEach(() => { + resetGatewayWorkAdmission(); resetDetachedTaskLifecycleRuntimeForTests(); vi.clearAllMocks(); vi.useFakeTimers(); @@ -311,12 +317,172 @@ describe("subagent registry seam flow", () => { }); afterEach(() => { + resetGatewayWorkAdmission(); resetDetachedTaskLifecycleRuntimeForTests(); mod.testing.setDepsForTest(); mod.resetSubagentRegistryForTests({ persist: false }); vi.useRealTimers(); }); + it("keeps a sweeper archive mutation root-admitted until deletion settles", async () => { + const now = Date.now(); + let releaseDelete: (() => void) | undefined; + mocks.callGateway.mockImplementation((request: { method?: string }) => { + if (request.method !== "sessions.delete") { + return Promise.resolve({}); + } + return new Promise>((resolve) => { + releaseDelete = () => resolve({}); + }); + }); + mod.addSubagentRunForTests({ + runId: "run-sweep-admission", + childSessionKey: "agent:main:subagent:sweep-admission", + requesterSessionKey: "agent:main:main", + requesterDisplayKey: "main", + task: "archive while suspension is possible", + cleanup: "delete", + createdAt: now - 10_000, + endedAt: now - 5_000, + cleanupCompletedAt: now - 4_000, + archiveAtMs: now - 1, + }); + + const sweep = mod.testing.runSweeperTickForTests(); + await waitForFast(() => expect(releaseDelete).toBeTypeOf("function")); + expect(getActiveGatewayRootWorkCount()).toBe(1); + + releaseDelete?.(); + await sweep; + await waitForFast(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + }); + + it("tracks missing-entry lifecycle result refresh until capture and persistence settle", async () => { + const childSessionKey = "agent:main:subagent:refresh-admission"; + mocks.callGateway.mockImplementation(async (request: { method?: string }) => + request.method === "agent.wait" ? { status: "pending" } : {}, + ); + mod.registerSubagentRun({ + runId: "run-refresh-admission-old", + childSessionKey, + requesterSessionKey: "agent:main:main", + requesterDisplayKey: "main", + task: "capture replacement completion", + cleanup: "keep", + expectsCompletionMessage: true, + }); + await waitForFast(() => expect(mocks.callGateway).toHaveBeenCalled()); + const entry = mod.getSubagentRunByChildSessionKey(childSessionKey); + expect(entry).not.toBeNull(); + Object.assign(entry ?? {}, { + endedAt: Date.now(), + outcome: { status: "ok" }, + }); + await waitForFast(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + + let finishCapture: ((value: string) => void) | undefined; + mocks.captureSubagentCompletionReply.mockImplementationOnce( + async () => + await new Promise((resolve) => { + finishCapture = resolve; + }), + ); + mocks.persistSubagentRunsToDisk.mockClear(); + const lifecycleHandler = mocks.onAgentEvent.mock.calls.at(-1)?.[0]; + expect(lifecycleHandler).toBeTypeOf("function"); + + lifecycleHandler?.({ + runId: "run-refresh-admission-new", + seq: 1, + stream: "lifecycle", + ts: Date.now(), + sessionKey: childSessionKey, + data: { phase: "end" }, + }); + + await waitForFast(() => expect(finishCapture).toBeTypeOf("function")); + expect(getActiveGatewayRootWorkCount()).toBe(1); + expect(entry?.completion?.resultText).toBeUndefined(); + + finishCapture?.("replacement final reply"); + await waitForFast(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + expect(entry?.completion?.resultText).toBe("replacement final reply"); + expect(mocks.persistSubagentRunsToDisk).toHaveBeenCalledOnce(); + }); + + it("retries a terminal completion deferred by restart drain", async () => { + const now = Date.now(); + const runId = "run-terminal-restart-retry"; + mod.addSubagentRunForTests({ + runId, + childSessionKey: "agent:main:subagent:terminal-restart-retry", + requesterSessionKey: "agent:main:main", + requesterDisplayKey: "main", + task: "deliver terminal completion after restart", + cleanup: "keep", + expectsCompletionMessage: true, + createdAt: now - 10_000, + startedAt: now - 9_000, + endedAt: now - 1_000, + endedReason: SUBAGENT_ENDED_REASON_ERROR, + outcome: { status: "error", error: "provider interrupted" }, + }); + + markGatewayRestartDraining(); + await expect( + mod.finalizeInterruptedSubagentRun({ + runId, + error: "provider interrupted", + endedAt: now - 1_000, + }), + ).resolves.toBe(1); + expect(mocks.runSubagentAnnounceFlow).not.toHaveBeenCalled(); + + resetGatewayWorkAdmission(); + await vi.advanceTimersByTimeAsync(1_000); + await waitForFast(() => expect(mocks.runSubagentAnnounceFlow).toHaveBeenCalledOnce()); + await waitForFast(() => { + const entry = mod + .listSubagentRunsForRequester("agent:main:main") + .find((candidate) => candidate.runId === runId); + expect(entry?.cleanupCompletedAt).toBeTypeOf("number"); + }); + }); + + it("keeps killed session timing root-admitted after task finalization", async () => { + let finishTiming: (() => void) | undefined; + mocks.patchSessionEntry.mockImplementationOnce(async () => { + await new Promise((resolve) => { + finishTiming = resolve; + }); + return null; + }); + mocks.callGateway.mockImplementation(async (request: { method?: string }) => + request.method === "agent.wait" ? { status: "pending" } : {}, + ); + const runId = "run-kill-tail-admission"; + mod.registerSubagentRun({ + runId, + childSessionKey: "agent:main:subagent:child", + requesterSessionKey: "agent:main:main", + requesterDisplayKey: "main", + task: "persist killed session state", + cleanup: "keep", + spawnMode: "session", + }); + await waitForFast(() => expect(mocks.callGateway).toHaveBeenCalled()); + + expect(mod.markSubagentRunTerminated({ runId, reason: "manual kill" })).toBe(1); + await waitForFast(() => expect(finishTiming).toBeTypeOf("function")); + expect(findTaskByRunIdForStatus(runId)).toMatchObject({ status: "cancelled" }); + // Bookkeeping can launch additional tracked cleanup tails; the held timing + // write must keep at least one independent root visible until it settles. + expect(getActiveGatewayRootWorkCount()).toBeGreaterThan(0); + + finishTiming?.(); + await waitForFast(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + }); + it("lists active and pending-delivery child sessions for maintenance preservation", () => { const now = Date.now(); mod.addSubagentRunForTests({ diff --git a/src/agents/subagent-registry.ts b/src/agents/subagent-registry.ts index 57c15d6da8c2..2974919e7529 100644 --- a/src/agents/subagent-registry.ts +++ b/src/agents/subagent-registry.ts @@ -11,6 +11,10 @@ import type { ContextEngine, SubagentEndReason } from "../context-engine/types.j import { callGateway } from "../gateway/call.js"; import { getAgentRunContext, onAgentEvent } from "../infra/agent-events.js"; import { createSubsystemLogger } from "../logging/subsystem.js"; +import { + isGatewayRestartDraining, + runWithGatewayIndependentRootWorkAdmission, +} from "../process/gateway-work-admission.js"; import { formatAbandonedLivenessError, formatBlockedLivenessError, @@ -225,6 +229,7 @@ let restoreAttempted = false; const ORPHAN_RECOVERY_DEBOUNCE_MS = 1_000; let lastOrphanRecoveryScheduleAt = 0; const SUBAGENT_ANNOUNCE_TIMEOUT_MS = 120_000; +const GATEWAY_ADMISSION_RETRY_DELAY_MS = 1_000; /** * Embedded runs can emit transient lifecycle `error` events while provider/model * retry is still in progress. Defer terminal error cleanup briefly so a @@ -379,6 +384,8 @@ export function scheduleSubagentOrphanRecovery(params?: { delayMs?: number; maxR lastOrphanRecoveryScheduleAt = now; void import("./subagent-orphan-recovery.js").then( ({ scheduleOrphanRecovery }) => { + // This import only installs timers. Each delayed or retrying recovery + // attempt owns independent root admission inside the recovery module. scheduleOrphanRecovery({ getActiveRuns: () => subagentRuns, delayMs: params?.delayMs, @@ -455,7 +462,10 @@ type CompleteSubagentRunParams = { suppressSessionEffects?: boolean; }; -async function completeSubagentRunWithRecovery(params: CompleteSubagentRunParams, source: string) { +async function completeSubagentRunWithRecoveryAttempt( + params: CompleteSubagentRunParams, + source: string, +) { try { await completeSubagentRun(params); return; @@ -506,6 +516,52 @@ async function completeSubagentRunWithRecovery(params: CompleteSubagentRunParams resumeSubagentRun(params.runId); } +function scheduleSubagentCompletionRetryAfterRestart( + params: CompleteSubagentRunParams, + source: string, + expectedEntry: SubagentRunRecord, +) { + const expectedGeneration = expectedEntry.generation; + const timer = setTimeout(() => { + resumeRetryTimers.delete(timer); + const current = subagentRuns.get(params.runId); + if (current !== expectedEntry || current.generation !== expectedGeneration) { + return; + } + void completeSubagentRunWithRecovery(params, source).catch((error: unknown) => { + log.warn("failed to retry subagent completion after gateway restart", { + source, + runId: params.runId, + error, + }); + }); + }, GATEWAY_ADMISSION_RETRY_DELAY_MS); + timer.unref?.(); + resumeRetryTimers.add(timer); +} + +async function completeSubagentRunWithRecovery(params: CompleteSubagentRunParams, source: string) { + // Each controller attempt owns its terminal transition, while this outer + // lease closes the gap between failed attempts and fallback cleanup. + try { + await runWithGatewayIndependentRootWorkAdmission(async () => { + await completeSubagentRunWithRecoveryAttempt(params, source); + }); + } catch (error) { + if (!isGatewayRestartDraining()) { + throw error; + } + log.warn("subagent completion deferred during gateway restart", { + source, + runId: params.runId, + }); + const current = subagentRuns.get(params.runId); + if (current) { + scheduleSubagentCompletionRetryAfterRestart(params, source, current); + } + } +} + function completeSubagentRunInBackground(params: CompleteSubagentRunParams, source: string) { void completeSubagentRunWithRecovery(params, source); } @@ -717,6 +773,61 @@ const { startSubagentAnnounceCleanupFlow, } = subagentLifecycleController; +function scheduleSubagentDeliveryResumeRetry( + runId: string, + scheduledEntry: SubagentRunRecord, + waitMs: number, +) { + const timer = setTimeout(() => { + resumeRetryTimers.delete(timer); + void runWithGatewayIndependentRootWorkAdmission(async () => { + if (subagentRuns.get(runId) !== scheduledEntry) { + resumedRuns.delete(runId); + return; + } + resumedRuns.delete(runId); + resumeSubagentRun(runId); + }).catch((error: unknown) => { + log.warn("failed to resume subagent delivery retry", { runId, error }); + if ( + isGatewayRestartDraining() && + subagentRuns.get(runId) === scheduledEntry && + typeof scheduledEntry.cleanupCompletedAt !== "number" + ) { + scheduleSubagentDeliveryResumeRetry( + runId, + scheduledEntry, + Math.max(waitMs, GATEWAY_ADMISSION_RETRY_DELAY_MS), + ); + return; + } + resumedRuns.delete(runId); + }); + }, waitMs); + timer.unref?.(); + resumeRetryTimers.add(timer); +} + +function finalizeResumedAnnounceGiveUpInBackground( + runId: string, + entry: SubagentRunRecord, + reason: "retry-limit" | "expiry", +) { + void runWithGatewayIndependentRootWorkAdmission(async () => { + await finalizeResumedAnnounceGiveUp({ runId, entry, reason }); + }).catch((error: unknown) => { + log.warn("failed to finalize exhausted subagent delivery", { runId, reason, error }); + if ( + isGatewayRestartDraining() && + subagentRuns.get(runId) === entry && + typeof entry.cleanupCompletedAt !== "number" + ) { + scheduleSubagentDeliveryResumeRetry(runId, entry, GATEWAY_ADMISSION_RETRY_DELAY_MS); + resumedRuns.add(runId); + } + }); +} + function resumeSubagentRun(runId: string) { if (!runId || resumedRuns.has(runId)) { return; @@ -738,11 +849,7 @@ function resumeSubagentRun(runId: string) { } // Skip entries that have exhausted their retry budget or expired (#18264). if (getDeliveryAttemptCount(entry) >= MAX_ANNOUNCE_RETRY_COUNT) { - void finalizeResumedAnnounceGiveUp({ - runId, - entry, - reason: "retry-limit", - }); + finalizeResumedAnnounceGiveUpInBackground(runId, entry, "retry-limit"); return; } if ( @@ -750,11 +857,7 @@ function resumeSubagentRun(runId: string) { typeof entry.endedAt === "number" && Date.now() - entry.endedAt > ANNOUNCE_EXPIRY_MS ) { - void finalizeResumedAnnounceGiveUp({ - runId, - entry, - reason: "expiry", - }); + finalizeResumedAnnounceGiveUpInBackground(runId, entry, "expiry"); return; } @@ -764,17 +867,7 @@ function resumeSubagentRun(runId: string) { const earliestRetryAt = (lastAttemptAt ?? 0) + delayMs; if (entry.expectsCompletionMessage === true && lastAttemptAt && now < earliestRetryAt) { const waitMs = Math.max(1, earliestRetryAt - now); - const scheduledEntry = entry; - const timer = setTimeout(() => { - resumeRetryTimers.delete(timer); - if (subagentRuns.get(runId) !== scheduledEntry) { - return; - } - resumedRuns.delete(runId); - resumeSubagentRun(runId); - }, waitMs); - timer.unref?.(); - resumeRetryTimers.add(timer); + scheduleSubagentDeliveryResumeRetry(runId, entry, waitMs); resumedRuns.add(runId); return; } @@ -884,12 +977,20 @@ function startSweeper() { async function runSubagentSweep() { try { - await sweepSubagentRuns(); + await runWithGatewayIndependentRootWorkAdmission(async () => { + await sweepSubagentRuns(); + }); } catch (err) { log.warn(`subagent run sweep failed: ${err instanceof Error ? err.message : String(err)}`); } } +function runSubagentSweepCleanupTail(runId: string, label: string, run: () => Promise) { + void runWithGatewayIndependentRootWorkAdmission(run).catch((error: unknown) => { + log.warn(`subagent sweep ${label} failed`, { runId, error }); + }); +} + function stopSweeper() { if (!sweeper) { return; @@ -1342,11 +1443,13 @@ async function sweepSubagentRuns() { now - entry.cleanupCompletedAt > SESSION_RUN_TTL_MS ) { clearPendingLifecycleError(runId); - void notifyContextEngineSubagentEnded({ - childSessionKey: entry.childSessionKey, - reason: "swept", - agentDir: entry.agentDir, - workspaceDir: entry.workspaceDir, + runSubagentSweepCleanupTail(runId, "context-engine cleanup", async () => { + await notifyContextEngineSubagentEnded({ + childSessionKey: entry.childSessionKey, + reason: "swept", + agentDir: entry.agentDir, + workspaceDir: entry.workspaceDir, + }); }); subagentRuns.delete(runId); mutated = true; @@ -1382,11 +1485,13 @@ async function sweepSubagentRuns() { mutated = true; // Archive/purge is terminal for the run record; remove any retained attachments too. await safeRemoveAttachmentsDir(entry); - void notifyContextEngineSubagentEnded({ - childSessionKey: entry.childSessionKey, - reason: "swept", - agentDir: entry.agentDir, - workspaceDir: entry.workspaceDir, + runSubagentSweepCleanupTail(runId, "context-engine cleanup", async () => { + await notifyContextEngineSubagentEnded({ + childSessionKey: entry.childSessionKey, + reason: "swept", + agentDir: entry.agentDir, + workspaceDir: entry.workspaceDir, + }); }); } // Sweep orphaned pendingLifecycleError entries (absolute TTL). @@ -1426,7 +1531,12 @@ function ensureListener() { const entry = subagentRuns.get(evt.runId); if (!entry) { if (phase === "end" && typeof evt.sessionKey === "string") { - await refreshFrozenResultFromSession(evt.sessionKey); + const sessionKey = evt.sessionKey; + // A replacement generation can finish after its predecessor row is + // terminal. Keep capture + persistence inside the suspension fence. + await runWithGatewayIndependentRootWorkAdmission(async () => { + await refreshFrozenResultFromSession(sessionKey); + }); } return; } @@ -1573,7 +1683,9 @@ const subagentRunManager = createSubagentRunManager({ resolveSubagentSessionStartedAt, notifyContextEngineSubagentEnded, completeCleanupBookkeeping, - completeSubagentRun, + completeSubagentRun: async (params) => { + await completeSubagentRunWithRecovery(params, "subagent-wait"); + }, resolveSubagentTask: findSubagentTaskForRun, }); diff --git a/src/auto-reply/reply/queue.drain-restart.test.ts b/src/auto-reply/reply/queue.drain-restart.test.ts index 9c62527aefb0..fb91e30a36a6 100644 --- a/src/auto-reply/reply/queue.drain-restart.test.ts +++ b/src/auto-reply/reply/queue.drain-restart.test.ts @@ -1,9 +1,18 @@ // Tests queue drain restart behavior when follow-up runs chain together. import { importFreshModule } from "openclaw/plugin-sdk/test-fixtures"; import { describe, expect, it, vi } from "vitest"; +import { + getActiveGatewayRootWorkCount, + isGatewaySubordinateWorkAdmissionClosed, + resetGatewayWorkAdmission, + tryBeginGatewayRootWorkAdmission, + tryBeginGatewaySuspendAdmission, +} from "../../process/gateway-work-admission.js"; +import { captureEnv, setTestEnvValue } from "../../test-utils/env.js"; import type { FollowupRun, QueueSettings } from "./queue.js"; import { clearFollowupQueue, + clearSessionQueues, enqueueFollowupRun, FollowupRunDeferredError, scheduleFollowupDrain, @@ -18,6 +27,81 @@ import { getExistingFollowupQueue } from "./queue/state.js"; installQueueRuntimeErrorSilencer(); describe("followup queue drain restart after idle window", () => { + it("keeps a detached drain on a live root after its enqueue request returns", async () => { + resetGatewayWorkAdmission(); + const key = `test-detached-drain-root-${Date.now()}`; + const settings: QueueSettings = { mode: "followup", debounceMs: 0, cap: 50 }; + const parentReleased = createDeferred(); + const drained = createDeferred(); + const parent = tryBeginGatewayRootWorkAdmission(); + if (!parent) { + throw new Error("expected parent Gateway work admission"); + } + let suspensionStarted = false; + let subordinateAdmissionClosed: boolean | undefined; + let activeRootCountDuringDrain: number | undefined; + + try { + await parent.run(async () => { + enqueueFollowupRun(key, createRun({ prompt: "detached" }), settings); + scheduleFollowupDrain(key, async () => { + await parentReleased.promise; + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + suspensionStarted = suspension !== null; + try { + subordinateAdmissionClosed = isGatewaySubordinateWorkAdmissionClosed(); + activeRootCountDuringDrain = getActiveGatewayRootWorkCount(); + } finally { + suspension?.rollback(); + drained.resolve(); + } + }); + }); + + parent.release(); + parentReleased.resolve(); + await drained.promise; + + expect(suspensionStarted).toBe(true); + expect(subordinateAdmissionClosed).toBe(false); + expect(activeRootCountDuringDrain).toBe(1); + await vi.waitFor(() => { + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); + } finally { + parent.release(); + parentReleased.resolve(); + clearSessionQueues([key]); + resetGatewayWorkAdmission(); + } + }); + + it("releases a detached drain root when its queue is cleared during debounce", async () => { + resetGatewayWorkAdmission(); + const env = captureEnv(["OPENCLAW_TEST_FAST"]); + setTestEnvValue("OPENCLAW_TEST_FAST", "0"); + const key = `test-cleared-debounce-root-${Date.now()}`; + const settings: QueueSettings = { mode: "followup", debounceMs: 60_000, cap: 50 }; + + try { + enqueueFollowupRun(key, createRun({ prompt: "clear during debounce" }), settings); + scheduleFollowupDrain(key, async () => {}); + await vi.waitFor(() => { + expect(getActiveGatewayRootWorkCount()).toBe(1); + }); + + clearSessionQueues([key]); + + await vi.waitFor(() => { + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); + } finally { + clearSessionQueues([key]); + env.restore(); + resetGatewayWorkAdmission(); + } + }); + it("does not retain stale callbacks when scheduleFollowupDrain runs with an empty queue", async () => { const key = `test-no-stale-callback-${Date.now()}`; const settings: QueueSettings = { mode: "followup", debounceMs: 0, cap: 50 }; diff --git a/src/auto-reply/reply/queue/drain.ts b/src/auto-reply/reply/queue/drain.ts index fc56b9399f05..95ce8be5a60f 100644 --- a/src/auto-reply/reply/queue/drain.ts +++ b/src/auto-reply/reply/queue/drain.ts @@ -9,6 +9,7 @@ import { channelRouteCompactKey, channelRouteDedupeKey, } from "../../../plugin-sdk/channel-route.js"; +import { runWithGatewayIndependentRootWorkContinuation } from "../../../process/gateway-work-admission.js"; import { defaultRuntime } from "../../../runtime.js"; import { buildPersistedUserTurnMediaInputsFromFields, @@ -1129,7 +1130,9 @@ export function scheduleFollowupDrain( // Cache callback only when a drain actually starts. Avoid keeping stale // callbacks around from finalize calls where no queue work is pending. rememberFollowupDrainCallback(key, effectiveRunFollowup); - void (async () => { + // Queue drains outlive their enqueue request across debounce and retries. + // Give the detached chain its own root so inherited request admission cannot go stale. + void runWithGatewayIndependentRootWorkContinuation(async () => { let retryDeferred = false; try { const collectState = { forceIndividualCollect: false }; @@ -1138,7 +1141,7 @@ export function scheduleFollowupDrain( if (queue.items.length === 0 && queue.droppedCount === 0) { break; } - await waitForQueueDebounce(queue); + await waitForQueueDebounce(queue, queue.abortController.signal); await dropAbortedFollowups(queue.items, effectiveRunFollowup); if (queue.items.length === 0 && queue.droppedCount === 0) { break; @@ -1364,5 +1367,8 @@ export function scheduleFollowupDrain( scheduleFollowupDrain(key, effectiveRunFollowup); } } - })(); + }).catch((err: unknown) => { + queue.draining = false; + defaultRuntime.error?.(`followup queue drain admission failed for ${key}: ${String(err)}`); + }); } diff --git a/src/auto-reply/reply/session-hooks-context.test.ts b/src/auto-reply/reply/session-hooks-context.test.ts index 245599f9dcd9..78c234c2af8a 100644 --- a/src/auto-reply/reply/session-hooks-context.test.ts +++ b/src/auto-reply/reply/session-hooks-context.test.ts @@ -5,6 +5,12 @@ import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } import type { OpenClawConfig } from "../../config/config.js"; import type { SessionEntry } from "../../config/sessions.js"; import type { HookRunner } from "../../plugins/hooks.js"; +import { + getActiveGatewayRootWorkCount, + markGatewayRestartDraining, + resetGatewayWorkAdmission, + tryBeginGatewayRootWorkAdmission, +} from "../../process/gateway-work-admission.js"; import { createSuiteTempRootTracker } from "../../test-helpers/temp-dir.js"; import { initSessionState } from "./session.js"; @@ -176,10 +182,12 @@ describe("session hook context wiring", () => { }); beforeEach(() => { + resetGatewayWorkAdmission(); hookRunnerMocks.hasHooks.mockReset(); hookRunnerMocks.runSessionStart.mockReset(); hookRunnerMocks.runSessionEnd.mockReset(); sessionCleanupMocks.closeTrackedBrowserTabsForSessions.mockClear(); + sessionCleanupMocks.closeTrackedBrowserTabsForSessions.mockResolvedValue(0); sessionCleanupMocks.resetRegisteredAgentHarnessSessions.mockClear(); sessionCleanupMocks.retireSessionMcpRuntime.mockClear(); hookRunnerMocks.runSessionStart.mockResolvedValue(undefined); @@ -190,6 +198,7 @@ describe("session hook context wiring", () => { }); afterEach(() => { + resetGatewayWorkAdmission(); vi.restoreAllMocks(); }); @@ -246,6 +255,85 @@ describe("session hook context wiring", () => { expectFields(startContext, { sessionId: startEvent?.sessionId }); }); + it("keeps rollover hooks and browser cleanup root-admitted until they settle", async () => { + const releases: Array<() => void> = []; + const held = () => + new Promise((resolve) => { + releases.push(resolve); + }); + hookRunnerMocks.runSessionEnd.mockImplementationOnce(held); + hookRunnerMocks.runSessionStart.mockImplementationOnce(held); + sessionCleanupMocks.closeTrackedBrowserTabsForSessions.mockImplementationOnce( + () => + new Promise((resolve) => { + releases.push(() => resolve(0)); + }), + ); + const sessionKey = "agent:main:telegram:direct:held-rollover"; + const { storePath } = await createStoredSession({ + prefix: "openclaw-session-hook-held-rollover", + sessionKey, + sessionId: "old-held-session", + }); + + await initSessionState({ + ctx: { Body: "/new", SessionKey: sessionKey }, + cfg: { session: { store: storePath } } as OpenClawConfig, + commandAuthorized: true, + }); + + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(3)); + await vi.waitFor(() => expect(releases).toHaveLength(3)); + for (const release of releases) { + release(); + } + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + }); + + it("hands rollover hooks off after restart drain closes admission", async () => { + const releases: Array<() => void> = []; + const held = () => + new Promise((resolve) => { + releases.push(resolve); + }); + hookRunnerMocks.runSessionEnd.mockImplementationOnce(held); + hookRunnerMocks.runSessionStart.mockImplementationOnce(held); + sessionCleanupMocks.closeTrackedBrowserTabsForSessions.mockImplementationOnce( + () => + new Promise((resolve) => { + releases.push(() => resolve(0)); + }), + ); + const sessionKey = "agent:main:telegram:direct:restart-handoff"; + const { storePath } = await createStoredSession({ + prefix: "openclaw-session-hook-restart-handoff", + sessionKey, + sessionId: "old-restart-session", + }); + const admission = tryBeginGatewayRootWorkAdmission(); + expect(admission).not.toBeNull(); + + await admission?.run(async () => { + markGatewayRestartDraining(); + await initSessionState({ + ctx: { Body: "/new", SessionKey: sessionKey }, + cfg: { session: { store: storePath } } as OpenClawConfig, + commandAuthorized: true, + }); + await vi.waitFor(() => expect(releases).toHaveLength(3)); + expect(getActiveGatewayRootWorkCount()).toBe(4); + }); + + admission?.release(); + expect(getActiveGatewayRootWorkCount()).toBe(3); + for (const release of releases) { + release(); + } + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + expect(hookRunnerMocks.runSessionEnd).toHaveBeenCalledTimes(1); + expect(hookRunnerMocks.runSessionStart).toHaveBeenCalledTimes(1); + }); + it("marks explicit /reset rollovers with reason reset", async () => { const sessionKey = "agent:main:telegram:direct:456"; const { storePath } = await createStoredSession({ diff --git a/src/auto-reply/reply/session-updates.lifecycle.test.ts b/src/auto-reply/reply/session-updates.lifecycle.test.ts index 1911375230ce..6856f4e2e281 100644 --- a/src/auto-reply/reply/session-updates.lifecycle.test.ts +++ b/src/auto-reply/reply/session-updates.lifecycle.test.ts @@ -6,6 +6,12 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { OpenClawConfig } from "../../config/config.js"; import { loadSessionStore, type SessionEntry } from "../../config/sessions.js"; import type { HookRunner } from "../../plugins/hooks.js"; +import { + getActiveGatewayRootWorkCount, + markGatewayRestartDraining, + resetGatewayWorkAdmission, + tryBeginGatewayRootWorkAdmission, +} from "../../process/gateway-work-admission.js"; const hookRunnerMocks = vi.hoisted(() => ({ hasHooks: vi.fn(), @@ -46,6 +52,7 @@ function firstSessionStartCall() { describe("session-updates lifecycle hooks", () => { beforeEach(async () => { + resetGatewayWorkAdmission(); vi.resetModules(); vi.doMock("../../plugins/hook-runner-global.js", () => ({ getGlobalHookRunner: () => @@ -67,6 +74,7 @@ describe("session-updates lifecycle hooks", () => { }); afterEach(async () => { + resetGatewayWorkAdmission(); vi.restoreAllMocks(); await Promise.all( tempDirs.splice(0).map((dir) => fs.rm(dir, { recursive: true, force: true })), @@ -109,6 +117,69 @@ describe("session-updates lifecycle hooks", () => { expect(startContext?.agentId).toBe("main"); }); + it("keeps compaction lifecycle hooks root-admitted until both settle", async () => { + const { storePath, sessionKey, sessionStore, entry } = await createFixture(); + const releases: Array<() => void> = []; + const heldHook = () => + new Promise((resolve) => { + releases.push(resolve); + }); + hookRunnerMocks.runSessionEnd.mockImplementationOnce(heldHook); + hookRunnerMocks.runSessionStart.mockImplementationOnce(heldHook); + + await incrementCompactionCount({ + cfg: { session: { store: storePath } } as OpenClawConfig, + sessionEntry: entry, + sessionStore, + sessionKey, + storePath, + newSessionId: "s2", + }); + + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(2)); + await vi.waitFor(() => expect(releases).toHaveLength(2)); + for (const release of releases) { + release(); + } + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + }); + + it("hands compaction lifecycle hooks off after restart drain closes admission", async () => { + const { storePath, sessionKey, sessionStore, entry } = await createFixture(); + const releases: Array<() => void> = []; + const heldHook = () => + new Promise((resolve) => { + releases.push(resolve); + }); + hookRunnerMocks.runSessionEnd.mockImplementationOnce(heldHook); + hookRunnerMocks.runSessionStart.mockImplementationOnce(heldHook); + const admission = tryBeginGatewayRootWorkAdmission(); + expect(admission).not.toBeNull(); + + await admission?.run(async () => { + markGatewayRestartDraining(); + await incrementCompactionCount({ + cfg: { session: { store: storePath } } as OpenClawConfig, + sessionEntry: entry, + sessionStore, + sessionKey, + storePath, + newSessionId: "s2", + }); + await vi.waitFor(() => expect(releases).toHaveLength(2)); + expect(getActiveGatewayRootWorkCount()).toBe(3); + }); + + admission?.release(); + expect(getActiveGatewayRootWorkCount()).toBe(2); + for (const release of releases) { + release(); + } + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + expect(hookRunnerMocks.runSessionEnd).toHaveBeenCalledTimes(1); + expect(hookRunnerMocks.runSessionStart).toHaveBeenCalledTimes(1); + }); + it("recreates a complete persisted row when compaction updates a missing store row", async () => { const { storePath, sessionKey, sessionStore, entry } = await createFixture(); await fs.writeFile(storePath, JSON.stringify({}, null, 2), "utf-8"); diff --git a/src/auto-reply/reply/session-updates.ts b/src/auto-reply/reply/session-updates.ts index 45b59869dbb2..900713365171 100644 --- a/src/auto-reply/reply/session-updates.ts +++ b/src/auto-reply/reply/session-updates.ts @@ -13,6 +13,7 @@ import { import { resolveStableSessionEndTranscript } from "../../gateway/session-transcript-files.fs.js"; import { logVerbose } from "../../globals.js"; import { getGlobalHookRunner } from "../../plugins/hook-runner-global.js"; +import { runWithGatewayIndependentRootWorkContinuation } from "../../process/gateway-work-admission.js"; import { resolveAgentIdFromSessionKey } from "../../routing/session-key.js"; import { getRemoteSkillEligibility } from "../../skills/runtime/remote.js"; import { resolveReusableWorkspaceSkillSnapshot } from "../../skills/runtime/session-snapshot.js"; @@ -96,7 +97,9 @@ function emitCompactionSessionLifecycleHooks(params: { transcriptArchived: transcript.transcriptArchived, nextSessionId: params.nextEntry.sessionId, }); - void hookRunner.runSessionEnd(payload.event, payload.context).catch((err: unknown) => { + void runWithGatewayIndependentRootWorkContinuation(async () => { + await hookRunner.runSessionEnd(payload.event, payload.context); + }).catch((err: unknown) => { logVerbose(`session_end hook failed: ${String(err)}`); }); } @@ -108,7 +111,9 @@ function emitCompactionSessionLifecycleHooks(params: { cfg: params.cfg, resumedFrom: params.previousEntry.sessionId, }); - void hookRunner.runSessionStart(payload.event, payload.context).catch((err: unknown) => { + void runWithGatewayIndependentRootWorkContinuation(async () => { + await hookRunner.runSessionStart(payload.event, payload.context); + }).catch((err: unknown) => { logVerbose(`session_start hook failed: ${String(err)}`); }); } diff --git a/src/auto-reply/reply/session.ts b/src/auto-reply/reply/session.ts index 706026420737..e25b7773331d 100644 --- a/src/auto-reply/reply/session.ts +++ b/src/auto-reply/reply/session.ts @@ -55,6 +55,7 @@ import { deliverSessionMaintenanceWarning } from "../../infra/session-maintenanc import { createSubsystemLogger } from "../../logging/subsystem.js"; import { getGlobalHookRunner } from "../../plugins/hook-runner-global.js"; import type { PluginHookSessionEndReason } from "../../plugins/hook-types.js"; +import { runWithGatewayIndependentRootWorkContinuation } from "../../process/gateway-work-admission.js"; import { buildAgentMainSessionKey, isAcpSessionKey, @@ -1062,11 +1063,15 @@ async function initSessionStateAttemptLocked( // their transcript aliases main; cleanup must carry both exact keys. const runtimePolicySessionKey = resolveRuntimePolicySessionKey({ cfg, ctx: sessionCtxForState, sessionKey }) ?? sessionKey; - void cleanupBrowserSessionsForLifecycleEnd({ - cfg, - sessionKeys: [previousSessionEntry.sessionId, sessionKey, runtimePolicySessionKey], - onWarn: (message) => log.warn(message), - onError: (error) => log.warn(`browser tab cleanup failed: ${String(error)}`), + void runWithGatewayIndependentRootWorkContinuation(async () => { + await cleanupBrowserSessionsForLifecycleEnd({ + cfg, + sessionKeys: [previousSessionEntry.sessionId, sessionKey, runtimePolicySessionKey], + onWarn: (message) => log.warn(message), + onError: (error) => log.warn(`browser tab cleanup failed: ${String(error)}`), + }); + }).catch((error: unknown) => { + log.warn(`browser tab cleanup admission failed: ${String(error)}`); }); } @@ -1108,7 +1113,9 @@ async function initSessionStateAttemptLocked( transcriptArchived: previousSessionTranscript.transcriptArchived, nextSessionId: effectiveSessionId, }); - void hookRunner.runSessionEnd(payload.event, payload.context).catch(() => {}); + void runWithGatewayIndependentRootWorkContinuation(async () => { + await hookRunner.runSessionEnd(payload.event, payload.context); + }).catch(() => {}); } } @@ -1133,7 +1140,9 @@ async function initSessionStateAttemptLocked( cfg, resumedFrom: previousSessionEntry?.sessionId, }); - void hookRunner.runSessionStart(payload.event, payload.context).catch(() => {}); + void runWithGatewayIndependentRootWorkContinuation(async () => { + await hookRunner.runSessionStart(payload.event, payload.context); + }).catch(() => {}); } } diff --git a/src/cli/gateway-cli/lifecycle.runtime.ts b/src/cli/gateway-cli/lifecycle.runtime.ts index 1f0f6c577c32..0411f481abd5 100644 --- a/src/cli/gateway-cli/lifecycle.runtime.ts +++ b/src/cli/gateway-cli/lifecycle.runtime.ts @@ -22,6 +22,7 @@ export { markGatewaySigusr1RestartHandled, peekGatewaySigusr1RestartReason, resetGatewayRestartStateForInProcessRestart, + rollbackGatewayRestartSignalAdmission, scheduleGatewaySigusr1Restart, } from "../../infra/restart.js"; export { writeGatewayRestartHandoffSync } from "../../infra/restart-handoff.js"; diff --git a/src/cli/gateway-cli/run-loop.test.ts b/src/cli/gateway-cli/run-loop.test.ts index 775ea870f4b5..86cf30c6719e 100644 --- a/src/cli/gateway-cli/run-loop.test.ts +++ b/src/cli/gateway-cli/run-loop.test.ts @@ -19,6 +19,7 @@ const isGatewaySigusr1RestartExternallyAllowed = vi.fn(() => false); const markGatewaySigusr1RestartHandled = vi.fn(); const peekGatewaySigusr1RestartReason = vi.fn<() => string | undefined>(() => undefined); const resetGatewayRestartStateForInProcessRestart = vi.fn(); +const rollbackGatewayRestartSignalAdmission = vi.fn(); const writeGatewayRestartHandoffSync = vi.fn((_opts: unknown) => ({ kind: "gateway-supervisor-restart-handoff" as const, version: 1 as const, @@ -129,6 +130,7 @@ vi.mock("../../infra/restart.js", () => ({ markGatewaySigusr1RestartHandled: () => markGatewaySigusr1RestartHandled(), peekGatewaySigusr1RestartReason: () => peekGatewaySigusr1RestartReason(), resetGatewayRestartStateForInProcessRestart: () => resetGatewayRestartStateForInProcessRestart(), + rollbackGatewayRestartSignalAdmission: () => rollbackGatewayRestartSignalAdmission(), resolveGatewayRestartDeferralTimeoutMs: (timeoutMs: unknown) => { if (typeof timeoutMs !== "number" || !Number.isFinite(timeoutMs)) { return DEFAULT_RESTART_DEFERRAL_TIMEOUT_MS; @@ -492,6 +494,9 @@ describe("runGatewayLoop", () => { expect(consumeGatewayRestartIntentPayloadSync).toHaveBeenCalledOnce(); expect(markGatewayDraining).toHaveBeenCalledOnce(); + expect(markGatewayDraining.mock.invocationCallOrder[0]).toBeLessThan( + loadConfig.mock.invocationCallOrder[0] ?? Number.POSITIVE_INFINITY, + ); expect(waitForActiveTasks).toHaveBeenCalledWith(90_000); expectRestartCloseCall(closeFirst, 90_000); await startedSecond; @@ -898,6 +903,9 @@ describe("runGatewayLoop", () => { await startedSecond; expect(start).toHaveBeenCalledTimes(2); + expect(markGatewayDraining.mock.invocationCallOrder[0]).toBeLessThan( + markGatewaySigusr1RestartHandled.mock.invocationCallOrder[0] ?? 0, + ); await new Promise((resolve) => { setImmediate(resolve); }); @@ -1477,6 +1485,9 @@ describe("runGatewayLoop", () => { expect(consumeGatewaySigusr1RestartAuthorization).toHaveBeenCalledOnce(); expect(markGatewaySigusr1RestartHandled).toHaveBeenCalledOnce(); + expect(markGatewayDraining.mock.invocationCallOrder[0]).toBeLessThan( + markGatewaySigusr1RestartHandled.mock.invocationCallOrder[0] ?? 0, + ); expect(markRestartAbortedMainSessions).toHaveBeenCalledWith({ cfg: { gateway: { @@ -2028,6 +2039,7 @@ describe("runGatewayLoop", () => { "SIGUSR1 handler failed: lifecycle module corrupted", ); expect(markGatewaySigusr1RestartHandled).toHaveBeenCalled(); + expect(rollbackGatewayRestartSignalAdmission).toHaveBeenCalledOnce(); expect(close).not.toHaveBeenCalled(); expect(start).toHaveBeenCalledTimes(1); @@ -2063,6 +2075,7 @@ describe("runGatewayLoop", () => { // Restart token must be cleared so future SIGUSR1 restarts are not // permanently coalesced as "already in-flight". expect(markGatewaySigusr1RestartHandled).toHaveBeenCalled(); + expect(rollbackGatewayRestartSignalAdmission).toHaveBeenCalledOnce(); expect(close).not.toHaveBeenCalled(); expect(start).toHaveBeenCalledTimes(1); @@ -2070,6 +2083,37 @@ describe("runGatewayLoop", () => { await expect(exited).resolves.toBe(0); }); }); + + it("recloses restart admission after a failed SIGUSR1 handler rolls it back", async () => { + vi.clearAllMocks(); + markGatewaySigusr1RestartHandled.mockImplementationOnce(() => { + throw new Error("restart token cleanup failed"); + }); + + await withIsolatedSignals(async ({ captureSignal }) => { + const { close, start, exited } = await createSignaledLoopHarness(); + const sigusr1 = captureSignal("SIGUSR1"); + const sigterm = captureSignal("SIGTERM"); + + sigusr1(); + await waitForLoopCondition( + () => rollbackGatewayRestartSignalAdmission.mock.calls.length === 1, + "failed SIGUSR1 handler did not roll back restart admission", + ); + + sigusr1(); + await waitForLoopCondition( + () => start.mock.calls.length === 2, + "second SIGUSR1 did not start a restart", + ); + + expect(close).toHaveBeenCalledTimes(1); + expect(markGatewayDraining).toHaveBeenCalledTimes(2); + + sigterm(); + await expect(exited).resolves.toBe(0); + }); + }); }); describe("gateway discover routing helpers", () => { diff --git a/src/cli/gateway-cli/run-loop.ts b/src/cli/gateway-cli/run-loop.ts index fed1298f088b..2ecb6073836c 100644 --- a/src/cli/gateway-cli/run-loop.ts +++ b/src/cli/gateway-cli/run-loop.ts @@ -146,7 +146,7 @@ export async function runGatewayLoop(params: { let activeRestartRequest: GatewayRunSignalRequest | null = null; let forceActiveRestartExit: (() => void) | null = null; let pendingStartupForceExitTimer: ReturnType | null = null; - let restartDrainingMarkPromise: Promise | null = null; + let restartDrainingMarked = false; let startupFailedWithoutServerHandle = false; const processInstanceId = randomUUID(); const waitForHealthyChild = params.waitForHealthyChild ?? waitForHealthyGatewayChild; @@ -408,17 +408,15 @@ export async function runGatewayLoop(params: { return DEFAULT_RESTART_DRAIN_TIMEOUT_MS; } }; - const markRestartDraining = async () => { - if (!restartDrainingMarkPromise) { - restartDrainingMarkPromise = (async () => { - const { markGatewayDraining } = await loadGatewayLifecycleRuntimeModule(); - markGatewayDraining(); - })().catch((err: unknown) => { - restartDrainingMarkPromise = null; - throw err; - }); + const markRestartDraining = () => { + if (restartDrainingMarked) { + return; } - await restartDrainingMarkPromise; + // The lifecycle module is primed before listeners are installed. Keep this + // transition synchronous so an accepted signal cannot yield between token + // handling and closing process-wide root admission. + eagerLifecycleRuntime.markGatewayDraining(); + restartDrainingMarked = true; }; const runAcceptedRequest = (acceptedRequest: GatewayRunSignalRequest) => { @@ -572,7 +570,7 @@ export async function runGatewayLoop(params: { // Reject new enqueues immediately during the drain window so // sessions get an explicit restart error instead of silent task loss. - await markRestartDraining(); + markRestartDraining(); const activeTasks = getActiveTaskCount(); const activeRuns = getActiveEmbeddedRunCount(); activeTasksAtDrainStart = activeTasks; @@ -736,8 +734,11 @@ export async function runGatewayLoop(params: { gatewayLog.info(`received ${signal} during shutdown; ignoring`); return; } - shuttingDown = true; const isRestart = action === "restart"; + if (isRestart) { + markRestartDraining(); + } + shuttingDown = true; gatewayLog.info(`received ${signal}; ${isRestart ? "restarting" : "shutting down"}`); if (isRestart) { startGatewayRestartTrace("restart.signal.received", [ @@ -758,9 +759,6 @@ export async function runGatewayLoop(params: { } if (!server || !restartResolver) { pendingStartupRequest = acceptedRequest; - void markRestartDraining().catch((err: unknown) => { - gatewayLog.warn(`failed to mark gateway draining for startup restart: ${String(err)}`); - }); armPendingStartupForceExitTimer(); return; } @@ -803,7 +801,9 @@ export async function runGatewayLoop(params: { const restartIntent = consumeGatewayRestartIntentPayloadSync(); if (restartIntent) { abortPendingChannelReloads(); - if (consumeGatewaySigusr1RestartAuthorization()) { + const authorized = consumeGatewaySigusr1RestartAuthorization(); + markRestartDraining(); + if (authorized) { markGatewaySigusr1RestartHandled(); } request("restart", "SIGUSR1", restartIntent.reason ?? "gateway.restart", restartIntent); @@ -836,6 +836,7 @@ export async function runGatewayLoop(params: { abortPendingChannelReloads(); const sigusr1RestartIntent = consumeGatewaySigusr1RestartIntent(); const restartReason = peekGatewaySigusr1RestartReason(); + markRestartDraining(); markGatewaySigusr1RestartHandled(); request( "restart", @@ -856,6 +857,14 @@ export async function runGatewayLoop(params: { } catch { // Best-effort: the eager reference itself is the recovery path. } + try { + eagerLifecycleRuntime.rollbackGatewayRestartSignalAdmission(); + // A later signal must repeat the synchronous close transition even if + // this handler failed after marking the one-way drain. + restartDrainingMarked = false; + } catch { + // Keep admission recovery independent from restart-token recovery. + } }); }; @@ -908,7 +917,7 @@ export async function runGatewayLoop(params: { let isFirstStart = true; for (;;) { await onIteration(); - restartDrainingMarkPromise = null; + restartDrainingMarked = false; startupStartedAt = Date.now(); let startupFailedBeforeServerHandle = false; try { diff --git a/src/cron/service.every-jobs-fire.test.ts b/src/cron/service.every-jobs-fire.test.ts index c80860c96caf..384b6547a677 100644 --- a/src/cron/service.every-jobs-fire.test.ts +++ b/src/cron/service.every-jobs-fire.test.ts @@ -1,5 +1,15 @@ // Every-job firing tests cover repeated schedule execution semantics. import { describe, expect, it, vi } from "vitest"; +import { + getGatewaySuspendStatus, + prepareGatewaySuspend, + resetGatewaySuspendCoordinatorForTest, +} from "../infra/gateway-suspend-coordinator.js"; +import { + beginGatewayRestartSignalAdmission, + isGatewayWorkAdmissionClosed, + resetGatewayWorkAdmission, +} from "../process/gateway-work-admission.js"; import { CronService } from "./service.js"; import { createStartedCronServiceWithFinishedBarrier, @@ -101,6 +111,156 @@ describe("CronService interval/cron jobs fire on time", () => { await store.cleanup(); }); + it("keeps a due timer frozen while scheduling is paused and fires it after resume", async () => { + const store = await makeStorePath(); + const { cron, enqueueSystemEvent, finished } = createStartedCronServiceWithFinishedBarrier({ + storePath: store.storePath, + logger: noopLogger, + }); + + await cron.start(); + const job = await cron.add({ + name: "suspension pause check", + enabled: true, + schedule: { kind: "every", everyMs: 10_000 }, + sessionTarget: "main", + wakeMode: "next-heartbeat", + payload: { kind: "systemEvent", text: "resumed-tick" }, + }); + + cron.pauseScheduling(); + await vi.advanceTimersByTimeAsync(10_005); + expect(enqueueSystemEvent).not.toHaveBeenCalled(); + + const finishedRun = finished.waitForOk(job.id); + cron.resumeScheduling(); + await vi.runOnlyPendingTimersAsync(); + await finishedRun; + expectMainSystemEvent(enqueueSystemEvent, "resumed-tick", job.id); + + cron.stop(); + await store.cleanup(); + }); + + it("rolls a failed scheduler resume back so a retry can rearm cron", async () => { + const store = await makeStorePath(); + const logger = createNoopLogger(); + const cron = new CronService({ + storePath: store.storePath, + cronEnabled: true, + log: logger, + enqueueSystemEvent: vi.fn(), + requestHeartbeat: vi.fn(), + runIsolatedAgentJob: vi.fn(async () => ({ status: "ok" as const })), + }); + + await cron.start(); + await cron.add({ + name: "resume retry check", + enabled: true, + schedule: { kind: "every", everyMs: 10_000 }, + sessionTarget: "main", + wakeMode: "next-heartbeat", + payload: { kind: "systemEvent", text: "resume" }, + }); + cron.pauseScheduling(); + logger.debug.mockImplementationOnce(() => { + throw new Error("arm failed"); + }); + + expect(() => cron.resumeScheduling()).toThrow("arm failed"); + expect(vi.getTimerCount()).toBe(0); + expect(() => cron.resumeScheduling()).not.toThrow(); + expect(vi.getTimerCount()).toBe(1); + + cron.stop(); + await store.cleanup(); + }); + + it("keeps admission closed until a real cron scheduler resume retry succeeds", async () => { + const store = await makeStorePath(); + const logger = createNoopLogger(); + const { cron, enqueueSystemEvent, finished } = createStartedCronServiceWithFinishedBarrier({ + storePath: store.storePath, + logger, + }); + resetGatewaySuspendCoordinatorForTest(); + resetGatewayWorkAdmission(); + + try { + await cron.start(); + const job = await cron.add({ + name: "coordinator resume retry check", + enabled: true, + schedule: { kind: "every", everyMs: 10_000 }, + sessionTarget: "main", + wakeMode: "next-heartbeat", + payload: { kind: "systemEvent", text: "recovered-tick" }, + }); + logger.debug.mockImplementationOnce(() => { + throw new Error("arm failed"); + }); + + expect( + prepareGatewaySuspend({ + requestId: "cron-resume-retry", + pauseScheduling: () => cron.pauseScheduling(), + resumeScheduling: () => cron.resumeScheduling(), + inspect: { getQueueSize: () => 1 }, + }), + ).toMatchObject({ status: "recovering" }); + expect(isGatewayWorkAdmissionClosed()).toBe(true); + + await vi.advanceTimersByTimeAsync(1_000); + expect(getGatewaySuspendStatus("stale-id")).toEqual({ status: "running" }); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + + const finishedRun = finished.waitForOk(job.id); + await vi.advanceTimersByTimeAsync(9_005); + await finishedRun; + expectMainSystemEvent(enqueueSystemEvent, "recovered-tick", job.id); + } finally { + cron.stop(); + resetGatewaySuspendCoordinatorForTest(); + resetGatewayWorkAdmission(); + await store.cleanup(); + } + }); + + it("keeps a due timer pending when restart signal admission rolls back", async () => { + const store = await makeStorePath(); + const { cron, enqueueSystemEvent, finished } = createStartedCronServiceWithFinishedBarrier({ + storePath: store.storePath, + logger: noopLogger, + }); + resetGatewayWorkAdmission(); + + try { + await cron.start(); + const job = await cron.add({ + name: "restart signal rollback check", + enabled: true, + schedule: { kind: "every", everyMs: 10_000 }, + sessionTarget: "main", + wakeMode: "next-heartbeat", + payload: { kind: "systemEvent", text: "rollback-tick" }, + }); + + const pendingSignal = beginGatewayRestartSignalAdmission(); + const finishedRun = finished.waitForOk(job.id); + await vi.advanceTimersByTimeAsync(10_005); + expect(enqueueSystemEvent).not.toHaveBeenCalled(); + + expect(pendingSignal.rollback()).toBe(true); + await finishedRun; + expectMainSystemEvent(enqueueSystemEvent, "rollback-tick", job.id); + } finally { + cron.stop(); + resetGatewayWorkAdmission(); + await store.cleanup(); + } + }); + it("fires a cron-expression job when the timer fires a few ms late", async () => { const store = await makeStorePath(); const { cron, enqueueSystemEvent, finished } = createStartedCronServiceWithFinishedBarrier({ diff --git a/src/cron/service.test-harness.ts b/src/cron/service.test-harness.ts index 316e608e5ec9..224458f66532 100644 --- a/src/cron/service.test-harness.ts +++ b/src/cron/service.test-harness.ts @@ -235,6 +235,8 @@ export function createMockCronStateForJobs(params: { durableNextRunAtMsByJobId: new Map(), running: false, stopped: false, + schedulingPaused: false, + schedulerStarted: false, restartRecoveryPending: false, pendingCatchupDeferralJobIds: new Set(), activeManualRunJobIds: new Set(), diff --git a/src/cron/service.ts b/src/cron/service.ts index 8691b787e821..9204bb535516 100644 --- a/src/cron/service.ts +++ b/src/cron/service.ts @@ -20,19 +20,74 @@ export type { CronEvent, CronServiceDeps } from "./service/state.js"; /** Public cron service facade that owns mutable scheduler state and delegates to locked ops. */ export class CronService implements CronServiceContract { private readonly state; + private startInProgress = 0; + private startState: { generation: number; promise: Promise } | null = null; + private lifecycleGeneration = 0; constructor(deps: CronServiceDeps) { this.state = createCronServiceState(deps); } async start() { - await ops.start(this.state); + const generation = this.lifecycleGeneration; + const pending = this.startState; + if (pending) { + try { + await pending.promise; + } catch (err) { + if (pending.generation === generation) { + throw err; + } + } + if (pending.generation === generation) { + return; + } + await this.start(); + return; + } + const promise = this.startOnce(generation); + this.startState = { generation, promise }; + try { + await promise; + } finally { + if (this.startState?.promise === promise) { + this.startState = null; + } + } + } + + private async startOnce(generation: number) { + this.startInProgress += 1; + this.state.schedulerStarted = false; + try { + await ops.start(this.state); + if (generation !== this.lifecycleGeneration) { + ops.stop(this.state); + return; + } + this.state.schedulerStarted = !this.state.stopped; + } finally { + this.startInProgress -= 1; + } } stop() { + this.lifecycleGeneration += 1; ops.stop(this.state); } + pauseScheduling() { + ops.pauseScheduling(this.state); + } + + resumeScheduling() { + ops.resumeScheduling(this.state); + } + + getSuspensionBlockerCount() { + return this.startInProgress; + } + async status() { return await ops.status(this.state); } diff --git a/src/cron/service/ops.ts b/src/cron/service/ops.ts index 418079385226..211bd5291dc2 100644 --- a/src/cron/service/ops.ts +++ b/src/cron/service/ops.ts @@ -316,9 +316,35 @@ export async function start(state: CronServiceState) { /** Stops the cron service timer without mutating persisted job state. */ export function stop(state: CronServiceState) { state.stopped = true; + state.schedulerStarted = false; stopTimer(state); } +/** Temporarily stops automatic ticks without running startup recovery on resume. */ +export function pauseScheduling(state: CronServiceState) { + state.schedulingPaused = true; + stopTimer(state); +} + +export function resumeScheduling(state: CronServiceState) { + if (!state.schedulingPaused) { + return; + } + state.schedulingPaused = false; + if (!state.schedulerStarted) { + return; + } + try { + armTimer(state); + } catch (err) { + // armTimer can install a timer before a later dependency throws. Roll the + // whole transition back so a suspension retry cannot reopen without cron. + state.schedulingPaused = true; + stopTimer(state); + throw err; + } +} + /** Returns cron service status after a read-only maintenance pass. */ export async function status(state: CronServiceState) { return await locked(state, async () => { diff --git a/src/cron/service/state.ts b/src/cron/service/state.ts index 75afdf88933c..1c7b78bcda65 100644 --- a/src/cron/service/state.ts +++ b/src/cron/service/state.ts @@ -206,6 +206,8 @@ export type CronServiceState = { timer: NodeJS.Timeout | null; running: boolean; stopped: boolean; + schedulingPaused: boolean; + schedulerStarted: boolean; restartRecoveryPending: boolean; /** Prevents maintenance reads from advancing deferred startup catch-up slots. * Entries are removed when the deferred job runs or becomes irrelevant. */ @@ -234,6 +236,8 @@ export function createCronServiceState(deps: CronServiceDeps): CronServiceState timer: null, running: false, stopped: false, + schedulingPaused: false, + schedulerStarted: false, restartRecoveryPending: false, pendingCatchupDeferralJobIds: new Set(), activeManualRunJobIds: new Set(), diff --git a/src/cron/service/timer.ts b/src/cron/service/timer.ts index d78991295ad6..d4c9ef52bb33 100644 --- a/src/cron/service/timer.ts +++ b/src/cron/service/timer.ts @@ -10,6 +10,10 @@ import { HEARTBEAT_SKIP_CRON_IN_PROGRESS, isRetryableHeartbeatBusySkipReason, } from "../../infra/heartbeat-wake.js"; +import { + beginGatewayRootWorkAdmissionWhenOpen, + GatewayDrainingError, +} from "../../process/gateway-work-admission.js"; import { DEFAULT_AGENT_ID, normalizeAgentId, @@ -1214,7 +1218,7 @@ export function armTimer(state: CronServiceState) { clearTimeout(state.timer); } state.timer = null; - if (state.stopped) { + if (state.stopped || state.schedulingPaused) { state.deps.log.debug({}, "cron: armTimer skipped - scheduler stopped"); return; } @@ -1276,7 +1280,7 @@ export function armTimer(state: CronServiceState) { } function armRunningRecheckTimer(state: CronServiceState) { - if (state.stopped) { + if (state.stopped || state.schedulingPaused) { return; } if (state.timer) { @@ -1289,9 +1293,29 @@ function armRunningRecheckTimer(state: CronServiceState) { }, MAX_TIMER_DELAY_MS); } -/** Handles one cron timer tick: load due jobs, reserve them, execute, persist, and re-arm. */ +/** Handles one cron timer tick under the process-wide root work admission. */ export async function onTimer(state: CronServiceState) { - if (state.stopped) { + let admission; + try { + // A restart signal can be rejected after temporarily closing admission. + // Wait for that decision so the consumed timer is not silently lost. + admission = await beginGatewayRootWorkAdmissionWhenOpen(); + } catch (err) { + if (err instanceof GatewayDrainingError) { + return; + } + throw err; + } + try { + await admission.run(async () => await onAdmittedTimer(state)); + } finally { + admission.release(); + } +} + +/** Loads due jobs, reserves them, executes, persists, and re-arms. */ +async function onAdmittedTimer(state: CronServiceState) { + if (state.stopped || state.schedulingPaused) { return; } if (state.restartRecoveryPending) { diff --git a/src/gateway/chat-abort.test.ts b/src/gateway/chat-abort.test.ts index 61848bff4878..3d72503d882a 100644 --- a/src/gateway/chat-abort.test.ts +++ b/src/gateway/chat-abort.test.ts @@ -2,6 +2,7 @@ // abort fanout, history snapshots, and cleanup of buffered streaming state. import { afterEach, describe, expect, it, vi } from "vitest"; import { isAgentRunRestartAbortReason } from "../agents/run-termination.js"; +import { onAgentEvent } from "../infra/agent-events.js"; import { abortChatRunById, abortChatRunsForProvider, @@ -307,6 +308,39 @@ describe("registerChatAbortController", () => { }); describe("abortChatRunById", () => { + it("retains terminal persistence ownership observed during abort", () => { + const { runId, sessionKey, entry, ops } = createAbortRunFixture({}); + let terminalEvents = 0; + const unsubscribe = onAgentEvent((event) => { + if (event.runId === runId && event.stream === "lifecycle" && event.data.phase === "end") { + terminalEvents += 1; + entry.projectSessionTerminalPending = true; + entry.projectSessionTerminalObservedAt = event.ts; + } + }); + + try { + const result = abortChatRunById(ops, { runId, sessionKey, stopReason: "user" }); + + expect(result).toEqual({ aborted: true }); + expect(entry.controller.signal.aborted).toBe(true); + expect(entry.projectSessionActive).toBe(false); + expect(entry.registrationCleanupRequested).toBe(true); + expect(entry.projectSessionTerminalPending).toBe(true); + expect(entry.projectSessionTerminalObservedAt).toEqual(expect.any(Number)); + expect(ops.chatAbortControllers.get(runId)).toBe(entry); + + expect(abortChatRunById(ops, { runId, sessionKey, stopReason: "user" })).toEqual({ + aborted: false, + }); + expect(terminalEvents).toBe(1); + expect(ops.broadcast).toHaveBeenCalledOnce(); + expect(ops.removeChatRun).toHaveBeenCalledOnce(); + } finally { + unsubscribe(); + } + }); + it("broadcasts aborted payload with partial message when buffered text exists", () => { const now = new Date("2026-01-02T03:04:05.000Z"); const { runId, sessionKey, entry, ops } = createAbortRunFixture({ diff --git a/src/gateway/chat-abort.ts b/src/gateway/chat-abort.ts index f69ce218f339..f7a7ccf32b98 100644 --- a/src/gateway/chat-abort.ts +++ b/src/gateway/chat-abort.ts @@ -470,6 +470,9 @@ function resolveDefaultGlobalAgentId(ops: ChatAbortOps): string | undefined { } export function isChatAbortControllerEntryAbortable(entry: ChatAbortControllerEntry): boolean { + if (entry.controller.signal.aborted) { + return false; + } try { return entry.isAbortable?.(entry) !== false; } catch { @@ -521,8 +524,13 @@ export function abortChatRunById( if (stopReason) { active.abortStopReason = stopReason; } + active.projectSessionActive = false; + // Reserve terminal ownership before abort listeners run; synchronous caller + // cleanup must not erase the entry before Gateway observes the event below. + active.projectSessionTerminalPending = true; + active.projectSessionTerminalObservedAt = undefined; + active.registrationCleanupRequested = true; active.controller.abort(createChatAbortSignalReason(stopReason)); - removeChatAbortControllerEntry(ops.chatAbortControllers, runId, active); ops.clearChatRunState(runId); const removed = ops.removeChatRun(runId, runId, sessionKey); if (active.controlUiVisible !== false) { @@ -551,6 +559,15 @@ export function abortChatRunById( endedAt: Date.now(), }, }); + // Gateway listeners synchronously stamp the terminal observation. Keep the + // entry as suspension-visible ownership until its persistence write settles. + if ( + ops.chatAbortControllers.get(runId) === active && + active.projectSessionTerminalObservedAt === undefined && + !active.projectSessionTerminalPersistence + ) { + removeChatAbortControllerEntry(ops.chatAbortControllers, runId, active); + } ops.agentRunSeq.delete(runId); if (removed?.clientRunId) { ops.agentRunSeq.delete(removed.clientRunId); diff --git a/src/gateway/config-reload.test.ts b/src/gateway/config-reload.test.ts index 46bc7c0a3558..0b75b2a8753e 100644 --- a/src/gateway/config-reload.test.ts +++ b/src/gateway/config-reload.test.ts @@ -16,6 +16,11 @@ import { resetPluginRuntimeStateForTest, setActivePluginRegistry, } from "../plugins/runtime.js"; +import { + getActiveGatewayRootWorkCount, + resetGatewayWorkAdmission, + runWithGatewayIndependentRootWorkAdmission, +} from "../process/gateway-work-admission.js"; import { getSkillsSnapshotVersion, resetSkillsRefreshStateForTest, @@ -723,6 +728,8 @@ function createReloaderHarness( promoteSnapshot?: (snapshot: ConfigFileSnapshot, reason: string) => Promise; initialPluginInstallRecords?: Record; readPluginInstallRecords?: () => Promise>; + runTransaction?: (run: () => Promise) => Promise; + onRestart?: (plan: GatewayReloadPlan, nextConfig: OpenClawConfig) => void | Promise; } = {}, ) { const watcher = createWatcherMock(); @@ -735,7 +742,9 @@ function createReloaderHarness( async (_plan: GatewayReloadPlan, _nextConfig: OpenClawConfig) => {}, ); const onHotReload = vi.fn(async (_plan: GatewayReloadPlan, _nextConfig: OpenClawConfig) => {}); - const onRestart = vi.fn((_plan: GatewayReloadPlan, _nextConfig: OpenClawConfig) => {}); + const onRestart = vi.fn( + options.onRestart ?? ((_plan: GatewayReloadPlan, _nextConfig: OpenClawConfig) => {}), + ); let writeListener: ((event: ConfigWriteNotification) => void) | null = null; const subscribeToWrites = vi.fn((listener: (event: ConfigWriteNotification) => void) => { writeListener = listener; @@ -765,6 +774,7 @@ function createReloaderHarness( onNoopConfigCommit, onHotReload, onRestart, + ...(options.runTransaction ? { runTransaction: options.runTransaction } : {}), log, watchPath: "/tmp/openclaw.json", }); @@ -816,10 +826,12 @@ function getOnlyPromoteSnapshotCall(promoteSnapshot: { describe("startGatewayConfigReloader", () => { beforeEach(() => { + resetGatewayWorkAdmission(); vi.useFakeTimers(); }); afterEach(() => { + resetGatewayWorkAdmission(); vi.useRealTimers(); vi.restoreAllMocks(); }); @@ -932,6 +944,43 @@ describe("startGatewayConfigReloader", () => { await harness.reloader.stop(); }); + it("keeps restart preparation inside the accepted config root", async () => { + let releaseRestart = () => {}; + let noteRestartStarted = () => {}; + const restartStarted = new Promise((resolve) => { + noteRestartStarted = resolve; + }); + const restartPending = new Promise((resolve) => { + releaseRestart = resolve; + }); + const initialConfig: OpenClawConfig = { + gateway: { reload: { debounceMs: 0 }, terminal: { enabled: true } }, + }; + const nextConfig: OpenClawConfig = { + gateway: { reload: { debounceMs: 0 }, terminal: { enabled: false } }, + }; + const harness = createReloaderHarness( + async () => makeSnapshot({ config: nextConfig, hash: "restart-root" }), + { + initialConfig, + runTransaction: runWithGatewayIndependentRootWorkAdmission, + onRestart: async () => { + noteRestartStarted(); + await restartPending; + }, + }, + ); + + harness.watcher.emit("change"); + await vi.runOnlyPendingTimersAsync(); + await restartStarted; + + expect(getActiveGatewayRootWorkCount()).toBe(1); + releaseRestart(); + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + await harness.reloader.stop(); + }); + it("does not notify lifecycle owners when reload mode ignores the change", async () => { const initialConfig: OpenClawConfig = { gateway: { reload: { mode: "off", debounceMs: 0 }, terminal: { enabled: true } }, diff --git a/src/gateway/config-reload.ts b/src/gateway/config-reload.ts index c12a6a9212ac..929063fb5029 100644 --- a/src/gateway/config-reload.ts +++ b/src/gateway/config-reload.ts @@ -116,6 +116,8 @@ export function startGatewayConfigReloader(opts: { onNoopConfigCommit: (plan: GatewayReloadPlan, nextConfig: OpenClawConfig) => Promise; onHotReload: (plan: GatewayReloadPlan, nextConfig: OpenClawConfig) => Promise; onRestart: (plan: GatewayReloadPlan, nextConfig: OpenClawConfig) => void | Promise; + /** Keeps one accepted config transaction inside the Gateway work fence. */ + runTransaction?: (run: () => Promise) => Promise; promoteSnapshot?: (snapshot: ConfigFileSnapshot, reason: string) => Promise; initialPluginInstallRecords?: PluginInstallRecords; readPluginInstallRecords?: () => Promise; @@ -164,21 +166,21 @@ export function startGatewayConfigReloader(opts: { const schedule = () => { scheduleAfter(settings.debounceMs); }; - const queueRestart = (plan: GatewayReloadPlan, nextConfig: OpenClawConfig) => { + const queueRestart = async (plan: GatewayReloadPlan, nextConfig: OpenClawConfig) => { if (restartQueued) { return; } restartQueued = true; - void (async () => { - try { - await opts.onRestart(plan, nextConfig); - } catch (err) { - // Restart checks can fail (for example unresolved SecretRefs). Keep the - // reloader alive and allow a future change to retry restart scheduling. - restartQueued = false; - opts.log.error(`config restart failed: ${String(err)}`); - } - })(); + try { + // Restart preparation reads secrets and can mutate auth/runtime state. + // Keep it inside the accepted config transaction instead of detaching it. + await opts.onRestart(plan, nextConfig); + } catch (err) { + // Restart checks can fail (for example unresolved SecretRefs). Keep the + // reloader alive and allow a future change to retry restart scheduling. + restartQueued = false; + opts.log.error(`config restart failed: ${String(err)}`); + } }; const handleMissingSnapshot = (snapshot: ConfigFileSnapshot): boolean => { @@ -297,12 +299,12 @@ export function startGatewayConfigReloader(opts: { restartReasons: [...plan.restartReasons, followUp.reason], }; await opts.onConfigChange?.(restartPlan, nextConfig); - queueRestart(restartPlan, nextConfig); + await queueRestart(restartPlan, nextConfig); return; } if (settings.mode === "restart") { await opts.onConfigChange?.({ ...plan, restartGateway: true }, nextConfig); - queueRestart(plan, nextConfig); + await queueRestart(plan, nextConfig); return; } if (plan.restartGateway) { @@ -315,7 +317,7 @@ export function startGatewayConfigReloader(opts: { return; } await opts.onConfigChange?.(plan, nextConfig); - queueRestart(plan, nextConfig); + await queueRestart(plan, nextConfig); return; } @@ -335,6 +337,14 @@ export function startGatewayConfigReloader(opts: { } }; + const runAcceptedTransaction = async (run: () => Promise) => { + if (opts.runTransaction) { + await opts.runTransaction(run); + return; + } + await run(); + }; + const promoteAcceptedInProcessWrite = async (persistedHash: string) => { if (!opts.promoteSnapshot) { return; @@ -368,12 +378,14 @@ export function startGatewayConfigReloader(opts: { const pendingWrite = pendingInProcessConfig; pendingInProcessConfig = null; missingConfigRetries = 0; - await applySnapshot( - pendingWrite.config, - pendingWrite.compareConfig, - pendingWrite.afterWrite, - ); - await promoteAcceptedInProcessWrite(pendingWrite.persistedHash); + await runAcceptedTransaction(async () => { + await applySnapshot( + pendingWrite.config, + pendingWrite.compareConfig, + pendingWrite.afterWrite, + ); + await promoteAcceptedInProcessWrite(pendingWrite.persistedHash); + }); return; } const snapshot = await opts.readSnapshot(); @@ -390,8 +402,10 @@ export function startGatewayConfigReloader(opts: { handleInvalidSnapshot(snapshot); return; } - await applySnapshot(snapshot.config, snapshot.sourceConfig); - await promoteAcceptedSnapshot(snapshot, "valid-config"); + await runAcceptedTransaction(async () => { + await applySnapshot(snapshot.config, snapshot.sourceConfig); + await promoteAcceptedSnapshot(snapshot, "valid-config"); + }); } catch (err) { opts.log.error(`config reload failed: ${String(err)}`); } finally { diff --git a/src/gateway/cron-exit-watchers.test.ts b/src/gateway/cron-exit-watchers.test.ts index 5a4b1b629e2b..35feedbb7d33 100644 --- a/src/gateway/cron-exit-watchers.test.ts +++ b/src/gateway/cron-exit-watchers.test.ts @@ -289,7 +289,9 @@ describe("createCronExitWatchers", () => { // The orphaned child is killed and the job never fires. expect(fake.runCancels.length).toBe(1); expect(fireOnExit).not.toHaveBeenCalled(); - expect(w.activeJobIds()).toEqual([]); + expect(w.activeJobIds()).toEqual(["job-a"]); + fake.runs[0].deferred.resolve({ exitCode: null, reason: "manual-cancel" }); + await vi.waitFor(() => expect(w.activeJobIds()).toEqual([])); }); it("does not arm a watcher for time-based or disabled jobs", async () => { @@ -325,8 +327,8 @@ describe("createCronExitWatchers", () => { expect(supervisor.spawn).toHaveBeenCalledTimes(1); }); - it("cancels the watcher when the job is removed from the set", async () => { - const { supervisor, cancelled } = makeFakeSupervisor(); + it("keeps a cancelled watcher blocking until the supervised child settles", async () => { + const { supervisor, cancelled, runs } = makeFakeSupervisor(); const w = createCronExitWatchers({ getProcessSupervisor: () => supervisor as never, persistCompletion: vi.fn(async () => {}), @@ -337,7 +339,10 @@ describe("createCronExitWatchers", () => { await flush(); w.reconcile([]); expect(cancelled).toContain("cron-exit:job-a"); - expect(w.activeJobIds()).toEqual([]); + expect(w.activeJobIds()).toEqual(["job-a"]); + + runs[0].deferred.resolve({ exitCode: null, reason: "manual-cancel" }); + await vi.waitFor(() => expect(w.activeJobIds()).toEqual([])); }); it("does not fire a job whose watcher was cancelled before exit", async () => { @@ -357,6 +362,35 @@ describe("createCronExitWatchers", () => { expect(fireOnExit).not.toHaveBeenCalled(); }); + it("retains a blocker and suppresses stale fire when removed during terminal persistence", async () => { + const { supervisor, runs } = makeFakeSupervisor(); + let releasePersist = () => {}; + const persistCompletion = vi.fn( + () => + new Promise((resolve) => { + releasePersist = resolve; + }), + ); + const fireOnExit = vi.fn(async () => {}); + const w = createCronExitWatchers({ + getProcessSupervisor: () => supervisor as never, + persistCompletion, + fireOnExit, + logger: noopLogger, + }); + w.reconcile([onExitJob("job-a")]); + await flush(); + + runs[0].deferred.resolve({ exitCode: 0, reason: "exit" }); + await vi.waitFor(() => expect(persistCompletion).toHaveBeenCalledOnce()); + w.reconcile([]); + expect(w.activeJobIds()).toEqual(["job-a"]); + + releasePersist(); + await vi.waitFor(() => expect(w.activeJobIds()).toEqual([])); + expect(fireOnExit).not.toHaveBeenCalled(); + }); + it("is one-shot: a fired job is not re-armed on a later reconcile", async () => { const { supervisor, runs } = makeFakeSupervisor(); const w = createCronExitWatchers({ diff --git a/src/gateway/cron-exit-watchers.ts b/src/gateway/cron-exit-watchers.ts index 1bd594a91f4f..3baebc16195b 100644 --- a/src/gateway/cron-exit-watchers.ts +++ b/src/gateway/cron-exit-watchers.ts @@ -79,17 +79,32 @@ export function createCronExitWatchers(params: { job: OnExitCronJob; run: ManagedRun | undefined; fired: boolean; + terminalPersisting: boolean; + cancelled: boolean; + lifecycleSettled: boolean; command: string; cwd: string | undefined; }; const active = new Map(); + // A cancelled child can keep running until the supervisor observes exit. + // Retain those slots separately so replacement arms can own the job id while + // suspension still sees every predecessor that is settling. + const settlingCancelledSlots = new Set(); const cancel = (jobId: string) => { const slot = active.get(jobId); if (!slot) { return; } - active.delete(jobId); + slot.cancelled = true; + if (!slot.lifecycleSettled) { + settlingCancelledSlots.add(slot); + } + // Terminal persistence is user-visible state. Keep the slot as a suspend + // blocker until that write settles even when hot reload cancels the watcher. + if (!slot.terminalPersisting) { + active.delete(jobId); + } // Cancel an already-spawned child; an in-flight spawn (run undefined) is // killed by the arm() ownership check once it resolves. slot.run?.cancel("manual-cancel"); @@ -106,7 +121,17 @@ export function createCronExitWatchers(params: { const armToken: object = {}; // Reserve the slot synchronously so a concurrent cancel/replace can observe // and act on this arm before the child is spawned. - const slot: WatcherSlot = { armToken, job, run: undefined, fired: false, command, cwd }; + const slot: WatcherSlot = { + armToken, + job, + run: undefined, + fired: false, + terminalPersisting: false, + cancelled: false, + lifecycleSettled: false, + command, + cwd, + }; active.set(job.id, slot); const owns = () => active.get(job.id) === slot && slot.armToken === armToken; void (async () => { @@ -136,8 +161,14 @@ export function createCronExitWatchers(params: { } if (!owns()) { // Cancelled or re-armed (changed command/cwd) while the spawn was in - // flight — kill this now-orphaned child instead of leaking it. + // flight — kill this now-orphaned child instead of leaking it. Wait for + // supervisor settlement so suspension cannot snapshot a live child. run.cancel("manual-cancel"); + try { + await run.wait(); + } catch { + // The watcher was already cancelled; settlement, not outcome, matters. + } return; } slot.run = run; @@ -165,6 +196,7 @@ export function createCronExitWatchers(params: { { jobId: job.id, exitCode: exit.exitCode, reason: exit.reason }, "cron-exit: watched command exited; firing job", ); + slot.terminalPersisting = true; // Persist the terminal one-shot state BEFORE firing. FAIL CLOSED: if the // store write fails we do NOT wake — waking without a persisted terminal // state would let a gateway restart re-arm and re-run the command. @@ -180,6 +212,13 @@ export function createCronExitWatchers(params: { ); return; } + slot.terminalPersisting = false; + if (!owns() || slot.cancelled) { + if (active.get(job.id) === slot) { + active.delete(job.id); + } + return; + } slot.fired = true; try { await params.fireOnExit(slot.job, { @@ -196,7 +235,13 @@ export function createCronExitWatchers(params: { "cron-exit: fireOnExit after exit failed", ); } - })(); + })().finally(() => { + slot.lifecycleSettled = true; + settlingCancelledSlots.delete(slot); + if (slot.cancelled && active.get(job.id) === slot) { + active.delete(job.id); + } + }); }; const reconcile = (jobs: CronJob[]) => { @@ -236,6 +281,14 @@ export function createCronExitWatchers(params: { reconcile, cancel, cancelAll, - activeJobIds: () => Array.from(active.keys()), + activeJobIds: () => + Array.from( + new Set([ + ...Array.from(active.entries()) + .filter(([, slot]) => !slot.fired) + .map(([jobId]) => jobId), + ...Array.from(settlingCancelledSlots, (slot) => slot.job.id), + ]), + ), }; } diff --git a/src/gateway/local-request-context.ts b/src/gateway/local-request-context.ts index fdbcbb12293b..2d2bb6b98c60 100644 --- a/src/gateway/local-request-context.ts +++ b/src/gateway/local-request-context.ts @@ -3,7 +3,6 @@ import { loadManifestModelCatalog } from "../agents/model-catalog.js"; import type { CliDeps } from "../cli/deps.types.js"; import type { OpenClawConfig } from "../config/types.openclaw.js"; -import type { CronServiceContract } from "../cron/service-contract.js"; import { createSubsystemLogger } from "../logging/subsystem.js"; import { getPluginRuntimeGatewayRequestScope, @@ -12,6 +11,7 @@ import { import { NodeRegistry } from "./node-registry.js"; import type { ChannelRuntimeSnapshot } from "./server-channel-runtime.types.js"; import { createChatRunEntry, type ChatRunEntry } from "./server-chat-state.js"; +import type { GatewayCronServiceContract } from "./server-cron.js"; import type { GatewayRequestContext } from "./server-methods/types.js"; // Embedded/local agent calls need enough GatewayRequestContext to reuse server @@ -26,11 +26,13 @@ function cronUnavailable(): never { throw new Error("Cron is unavailable in local embedded agent gateway context."); } -const unavailableCron: CronServiceContract = { +const unavailableCron: GatewayCronServiceContract = { start: async () => { cronUnavailable(); }, stop: () => {}, + pauseScheduling: () => {}, + resumeScheduling: () => {}, status: async () => cronUnavailable(), list: async () => cronUnavailable(), listPage: async () => cronUnavailable(), diff --git a/src/gateway/methods/core-descriptors.ts b/src/gateway/methods/core-descriptors.ts index c13cb743fdd6..cf027147c63f 100644 --- a/src/gateway/methods/core-descriptors.ts +++ b/src/gateway/methods/core-descriptors.ts @@ -288,6 +288,15 @@ export const CORE_GATEWAY_METHOD_SPECS: readonly CoreGatewayMethodSpec[] = [ // Session PR chips read the session's own checkout metadata, matching the // sessions.files.* trusted-operator read domain. { name: "controlUi.sessionPullRequests", scope: "operator.read" }, + { + name: "gateway.suspend.prepare", + scope: "operator.admin", + startup: true, + controlPlaneWrite: true, + }, + { name: "gateway.suspend.status", scope: "operator.read" }, + // Resume is the safety escape hatch and must not sit behind write-rate limiting. + { name: "gateway.suspend.resume", scope: "operator.admin" }, ] as const; const CORE_GATEWAY_METHOD_SPEC_BY_NAME: ReadonlyMap = new Map( diff --git a/src/gateway/server-active-work.test.ts b/src/gateway/server-active-work.test.ts new file mode 100644 index 000000000000..98703bf9b21f --- /dev/null +++ b/src/gateway/server-active-work.test.ts @@ -0,0 +1,57 @@ +// Covers server-local chat, cron watcher, queued-turn, and terminal blockers. +import { describe, expect, it, vi } from "vitest"; +import { createGatewayServerActiveWorkInspectors } from "./server-active-work.js"; +import type { GatewayRequestContext } from "./server-methods/shared-types.js"; + +vi.mock("../cron/active-jobs.js", () => ({ + getActiveCronJobCount: vi.fn(() => 2), +})); + +vi.mock("../tasks/cron-task-cancel.js", () => ({ + getSuspensionVisibleCronTaskRunCount: vi.fn(() => 4), +})); + +function controller(aborted = false): AbortController { + const value = new AbortController(); + if (aborted) { + value.abort(); + } + return value; +} + +describe("gateway server active work inspectors", () => { + it("filters completed chat entries while retaining persistence and watcher blockers", () => { + const context = { + cron: { getSuspensionBlockerCount: () => 1 }, + chatAbortControllers: new Map([ + ["active", { controller: controller() }], + ["aborted", { controller: controller(true) }], + [ + "persisting", + { + controller: controller(), + registrationCleanupRequested: true, + controlUiVisible: true, + projectSessionTerminalPending: true, + }, + ], + ]), + chatQueuedTurns: new Map([ + ["queued", { controller: controller() }], + ["cancelled", { controller: controller(true) }], + ]), + terminalSessions: { size: 2 }, + } as unknown as Pick< + GatewayRequestContext, + "chatAbortControllers" | "chatQueuedTurns" | "cron" | "terminalSessions" + >; + + const inspectors = createGatewayServerActiveWorkInspectors(context); + + expect(inspectors.getCronRuns?.()).toBe(5); + expect(inspectors.getChatRuns?.()).toBe(1); + expect(inspectors.getQueuedTurns?.()).toBe(1); + expect(inspectors.getTerminalPersistence?.()).toBe(1); + expect(inspectors.getTerminalSessions?.()).toBe(2); + }); +}); diff --git a/src/gateway/server-active-work.ts b/src/gateway/server-active-work.ts new file mode 100644 index 000000000000..b9ab2c441a11 --- /dev/null +++ b/src/gateway/server-active-work.ts @@ -0,0 +1,35 @@ +// Adapts server-local chat and terminal state to the shared activity inspector. +import { getActiveCronJobCount } from "../cron/active-jobs.js"; +import type { GatewayActiveWorkInspectors } from "../infra/gateway-active-work.js"; +import { getSuspensionVisibleCronTaskRunCount } from "../tasks/cron-task-cancel.js"; +import type { GatewayRequestContext } from "./server-methods/shared-types.js"; + +export function createGatewayServerActiveWorkInspectors( + context: Pick< + GatewayRequestContext, + "chatAbortControllers" | "chatQueuedTurns" | "cron" | "terminalSessions" + >, +): Partial { + return { + getCronRuns: () => + Math.max(getActiveCronJobCount(), getSuspensionVisibleCronTaskRunCount()) + + (context.cron.getSuspensionBlockerCount?.() ?? 0), + getChatRuns: () => + Array.from(context.chatAbortControllers.values()).filter( + (entry) => !entry.controller.signal.aborted && entry.registrationCleanupRequested !== true, + ).length, + getQueuedTurns: () => + Array.from(context.chatQueuedTurns.values()).filter( + (entry) => !entry.controller.signal.aborted, + ).length, + getTerminalPersistence: () => + Array.from(context.chatAbortControllers.values()).filter( + (entry) => + entry.controlUiVisible !== false && + entry.projectSessionTerminalPersisted !== true && + (entry.projectSessionTerminalPending === true || + entry.projectSessionTerminalPersistence !== undefined), + ).length, + getTerminalSessions: () => context.terminalSessions?.size ?? 0, + }; +} diff --git a/src/gateway/server-cron-lazy.test.ts b/src/gateway/server-cron-lazy.test.ts index cd7993a3b531..2e71019fe366 100644 --- a/src/gateway/server-cron-lazy.test.ts +++ b/src/gateway/server-cron-lazy.test.ts @@ -4,8 +4,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import type { CliDeps } from "../cli/deps.types.js"; import type { OpenClawConfig } from "../config/types.openclaw.js"; -import type { CronServiceContract } from "../cron/service-contract.js"; -import type { GatewayCronState } from "./server-cron.js"; +import type { GatewayCronServiceContract, GatewayCronState } from "./server-cron.js"; const hoisted = vi.hoisted(() => { let state: unknown; @@ -23,6 +22,14 @@ vi.mock("./server-cron.js", () => ({ const { createLazyGatewayCronState } = await import("./server-cron-lazy.js"); +function deferred() { + let resolve = () => {}; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + describe("createLazyGatewayCronState", () => { beforeEach(() => { vi.unstubAllEnvs(); @@ -111,6 +118,26 @@ describe("createLazyGatewayCronState", () => { expect(cron["start"]).toHaveBeenCalledTimes(2); }); + it("restarts after stop interrupts an in-flight startup", async () => { + const finishFirstStart = deferred(); + const cron = createCronService(); + cron.start = vi + .fn() + .mockImplementationOnce(async () => await finishFirstStart.promise) + .mockResolvedValueOnce(undefined); + hoisted.setState(createCronState(cron)); + const lazy = createLazyGatewayCronState(createParams()); + + const firstStart = lazy.cron.start(); + await vi.waitFor(() => expect(cron["start"]).toHaveBeenCalledOnce()); + lazy.cron.stop(); + const restarted = lazy.cron.start(); + finishFirstStart.resolve(); + await Promise.all([firstStart, restarted]); + + expect(cron["start"]).toHaveBeenCalledTimes(2); + }); + it("keeps synchronous wake non-blocking before the cron service is loaded", async () => { const cron = createCronService(); hoisted.setState(createCronState(cron)); @@ -134,6 +161,104 @@ describe("createLazyGatewayCronState", () => { expect(hoisted.buildGatewayCronService).not.toHaveBeenCalled(); }); + it("does not arm a read-loaded scheduler when suspension ends", async () => { + const cron = createCronService(); + hoisted.setState(createCronState(cron)); + const lazy = createLazyGatewayCronState(createParams()); + + lazy.cron.pauseScheduling(); + expect(hoisted.buildGatewayCronService).not.toHaveBeenCalled(); + + await lazy.cron.status(); + expect(cron["pauseScheduling"]).toHaveBeenCalledOnce(); + lazy.cron.resumeScheduling(); + expect(cron["resumeScheduling"]).not.toHaveBeenCalled(); + + await lazy.cron.start(); + expect(cron["resumeScheduling"]).toHaveBeenCalledOnce(); + lazy.cron.pauseScheduling(); + lazy.cron.resumeScheduling(); + expect(cron["resumeScheduling"]).toHaveBeenCalledTimes(2); + }); + + it("waits to start while scheduling is paused", async () => { + const cron = createCronService(); + hoisted.setState(createCronState(cron)); + const lazy = createLazyGatewayCronState(createParams()); + + lazy.cron.pauseScheduling(); + const startPromise = lazy.cron.start(); + await vi.waitFor(() => expect(lazy.cron.getSuspensionBlockerCount?.()).toBe(1)); + expect(cron["start"]).not.toHaveBeenCalled(); + + lazy.cron.resumeScheduling(); + await startPromise; + expect(cron["start"]).toHaveBeenCalledOnce(); + expect(lazy.cron.getSuspensionBlockerCount?.()).toBe(0); + }); + + it("keeps in-flight startup as a blocker until startup settles", async () => { + const finishStart = deferred(); + const cron = createCronService(); + cron.start = vi.fn(async () => await finishStart.promise); + hoisted.setState(createCronState(cron)); + const lazy = createLazyGatewayCronState(createParams()); + + const startPromise = lazy.cron.start(); + await vi.waitFor(() => expect(cron["start"]).toHaveBeenCalledOnce()); + expect(lazy.cron.getSuspensionBlockerCount?.()).toBe(1); + + lazy.cron.pauseScheduling(); + lazy.cron.resumeScheduling(); + expect(cron["resumeScheduling"]).toHaveBeenCalledOnce(); + expect(lazy.cron.getSuspensionBlockerCount?.()).toBe(1); + + finishStart.resolve(); + await startPromise; + expect(lazy.cron.getSuspensionBlockerCount?.()).toBe(0); + }); + + it("resumes a paused scheduler while exit watchers are reconciling", async () => { + const reconcileStarted = deferred(); + const finishReconcile = deferred(); + const cron = createCronService(); + hoisted.setState({ + ...createCronState(cron), + reconcileExitWatchers: vi.fn(async () => { + reconcileStarted.resolve(); + await finishReconcile.promise; + }), + }); + const lazy = createLazyGatewayCronState(createParams()); + + const startPromise = lazy.cron.start(); + await reconcileStarted.promise; + expect(lazy.cron.getSuspensionBlockerCount?.()).toBe(1); + + lazy.cron.pauseScheduling(); + lazy.cron.resumeScheduling(); + expect(cron["resumeScheduling"]).toHaveBeenCalledOnce(); + + finishReconcile.resolve(); + await startPromise; + expect(lazy.cron.getSuspensionBlockerCount?.()).toBe(0); + }); + + it("allows startup to retry after the underlying service rejects", async () => { + const cron = createCronService(); + cron.start = vi + .fn() + .mockRejectedValueOnce(new Error("startup failed")) + .mockResolvedValueOnce(undefined); + hoisted.setState(createCronState(cron)); + const lazy = createLazyGatewayCronState(createParams()); + + await expect(lazy.cron.start()).rejects.toThrow("startup failed"); + await lazy.cron.start(); + + expect(cron["start"]).toHaveBeenCalledTimes(2); + }); + it("does not reconcile exit watchers when cron is disabled", async () => { const cron = createCronService(); const reconcileExitWatchers = vi.fn(async () => {}); @@ -161,7 +286,7 @@ function createParams(overrides: Partial = {}) { }; } -function createCronState(cron: CronServiceContract): GatewayCronState { +function createCronState(cron: GatewayCronServiceContract): GatewayCronState { return { cron, storePath: "/tmp/openclaw-cron.json", @@ -169,10 +294,12 @@ function createCronState(cron: CronServiceContract): GatewayCronState { } as GatewayCronState; } -function createCronService(): CronServiceContract { +function createCronService(): GatewayCronServiceContract { return { start: vi.fn(async () => undefined), stop: vi.fn(), + pauseScheduling: vi.fn(), + resumeScheduling: vi.fn(), status: vi.fn(async () => ({ enabled: true }) as never), list: vi.fn(async () => [] as never), listPage: vi.fn(async () => ({ items: [], total: 0 }) as never), diff --git a/src/gateway/server-cron-lazy.ts b/src/gateway/server-cron-lazy.ts index b11c19c7fab3..64d66ae47f78 100644 --- a/src/gateway/server-cron-lazy.ts +++ b/src/gateway/server-cron-lazy.ts @@ -2,10 +2,9 @@ // Defers scheduler startup until cron is touched by runtime or API handlers. import type { CliDeps } from "../cli/deps.types.js"; import type { OpenClawConfig } from "../config/types.openclaw.js"; -import type { CronServiceContract } from "../cron/service-contract.js"; import { resolveCronJobsStorePath } from "../cron/store.js"; import { createLazyPromiseLoader } from "../shared/lazy-runtime.js"; -import type { GatewayCronState } from "./server-cron.js"; +import type { GatewayCronServiceContract, GatewayCronState } from "./server-cron.js"; type LazyGatewayCronParams = { cfg: OpenClawConfig; @@ -15,7 +14,12 @@ type LazyGatewayCronParams = { type LoadedGatewayCronState = { state: GatewayCronState; - started: boolean; + phase: "idle" | "starting" | "started" | "stopped"; + startPromise: Promise | null; + startGeneration: number | null; + schedulingPaused: boolean; + underlyingStartInFlight: boolean; + underlyingStarted: boolean; }; /** Creates a cron state proxy that imports the real cron service on first use. */ @@ -24,13 +28,40 @@ export function createLazyGatewayCronState(params: LazyGatewayCronParams): Gatew const cronEnabled = process.env.OPENCLAW_SKIP_CRON !== "1" && params.cfg.cron?.enabled !== false; let loaded: LoadedGatewayCronState | null = null; let stopped = false; + let lifecycleGeneration = 0; + let schedulingPaused = false; + const schedulingResumeWaiters = new Set<() => void>(); + const releaseSchedulingResumeWaiters = () => { + const waiters = Array.from(schedulingResumeWaiters); + schedulingResumeWaiters.clear(); + for (const resolve of waiters) { + resolve(); + } + }; + const waitForSchedulingResume = async () => { + if (!schedulingPaused) { + return; + } + await new Promise((resolve) => { + schedulingResumeWaiters.add(resolve); + }); + }; const cronStateLoader = createLazyPromiseLoader( () => import("./server-cron.js").then(({ buildGatewayCronService }) => { loaded = { state: buildGatewayCronService(params), - started: false, + phase: "idle", + startPromise: null, + startGeneration: null, + schedulingPaused: false, + underlyingStartInFlight: false, + underlyingStarted: false, }; + if (schedulingPaused) { + loaded.state.cron.pauseScheduling(); + loaded.schedulingPaused = true; + } return loaded; }), { cacheRejections: true }, @@ -45,35 +76,106 @@ export function createLazyGatewayCronState(params: LazyGatewayCronParams): Gatew return await cronStateLoader.load(); }; - const cron: CronServiceContract = { + const cron: GatewayCronServiceContract = { async start() { stopped = false; + const generation = lifecycleGeneration; + const startCancelled = () => stopped || generation !== lifecycleGeneration; const resolved = await load(); - if (stopped) { + const hasStarted = () => resolved.phase === "started"; + if (startCancelled()) { return; } - if (resolved.started) { + if (hasStarted()) { return; } - resolved.started = true; - await resolved.state.cron.start(); - // Arm on-exit watchers for jobs loaded from the store at startup (no - // change event fires for already-persisted jobs). - if (resolved.state.cronEnabled) { - await resolved.state.reconcileExitWatchers?.(); + if (resolved.startPromise) { + const pendingGeneration = resolved.startGeneration; + try { + await resolved.startPromise; + } catch (err) { + if (pendingGeneration === generation) { + throw err; + } + } + if (startCancelled() || hasStarted()) { + return; + } + if (pendingGeneration !== generation) { + await cron.start(); + return; + } } - // If stop raced the lazy import/start path, immediately stop the loaded - // scheduler so shutdown does not leave a background loop alive. - if (stopped && resolved.started) { - resolved.started = false; - resolved.state.cron.stop(); - resolved.state.stopExitWatchers?.(); + resolved.phase = "starting"; + resolved.startGeneration = generation; + const startPromise = (async () => { + await waitForSchedulingResume(); + if (startCancelled()) { + resolved.phase = "stopped"; + return; + } + if (resolved.schedulingPaused) { + resolved.state.cron.resumeScheduling(); + resolved.schedulingPaused = false; + } + resolved.underlyingStartInFlight = true; + try { + await resolved.state.cron.start(); + resolved.underlyingStarted = true; + } catch (err) { + resolved.underlyingStarted = false; + resolved.phase = startCancelled() ? "stopped" : "idle"; + throw err; + } finally { + resolved.underlyingStartInFlight = false; + } + if (startCancelled()) { + resolved.phase = "stopped"; + resolved.underlyingStarted = false; + resolved.state.cron.stop(); + resolved.state.stopExitWatchers?.(); + return; + } + if (schedulingPaused) { + resolved.state.cron.pauseScheduling(); + resolved.schedulingPaused = true; + } + // Arm on-exit watchers for jobs loaded from the store at startup (no + // change event fires for already-persisted jobs). + try { + if (resolved.state.cronEnabled) { + await resolved.state.reconcileExitWatchers?.(); + } + } catch (err) { + resolved.phase = startCancelled() ? "stopped" : "started"; + throw err; + } + if (startCancelled()) { + resolved.phase = "stopped"; + resolved.underlyingStarted = false; + resolved.state.cron.stop(); + resolved.state.stopExitWatchers?.(); + return; + } + resolved.phase = "started"; + })(); + resolved.startPromise = startPromise; + try { + await startPromise; + } finally { + if (resolved.startPromise === startPromise) { + resolved.startPromise = null; + resolved.startGeneration = null; + } } }, stop() { stopped = true; + lifecycleGeneration += 1; + releaseSchedulingResumeWaiters(); if (loaded) { - loaded.started = false; + loaded.phase = "stopped"; + loaded.underlyingStarted = false; loaded.state.cron.stop(); loaded.state.stopExitWatchers?.(); return; @@ -87,13 +189,37 @@ export function createLazyGatewayCronState(params: LazyGatewayCronParams): Gatew if (!stopped) { return; } - resolved.started = false; + resolved.phase = "stopped"; + resolved.underlyingStarted = false; resolved.state.cron.stop(); resolved.state.stopExitWatchers?.(); }) .catch(() => {}); } }, + pauseScheduling() { + schedulingPaused = true; + if (loaded) { + loaded.state.cron.pauseScheduling(); + loaded.schedulingPaused = true; + } + }, + resumeScheduling() { + schedulingPaused = false; + releaseSchedulingResumeWaiters(); + if ( + loaded && + loaded.schedulingPaused && + (loaded.underlyingStarted || loaded.underlyingStartInFlight) + ) { + loaded.state.cron.resumeScheduling(); + loaded.schedulingPaused = false; + } + }, + getSuspensionBlockerCount() { + const loadedBlockers = loaded?.state.cron.getSuspensionBlockerCount?.() ?? 0; + return loaded?.phase === "starting" ? Math.max(1, loadedBlockers) : loadedBlockers; + }, async status() { return await (await load()).state.cron.status(); }, diff --git a/src/gateway/server-cron-notifications.test.ts b/src/gateway/server-cron-notifications.test.ts index 1e45b62a4d6b..bfc2b83f76ec 100644 --- a/src/gateway/server-cron-notifications.test.ts +++ b/src/gateway/server-cron-notifications.test.ts @@ -1,13 +1,20 @@ // Cron notification tests protect completion-delivery warning behavior, // including URL redaction for invalid webhook destinations. -import { beforeEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { CliDeps } from "../cli/deps.types.js"; import { makeCronJob } from "../cron/delivery.test-helpers.js"; import type { CronJob } from "../cron/types.js"; +import { + getActiveGatewayRootWorkCount, + resetGatewayWorkAdmission, + tryBeginGatewayRootWorkAdmission, + tryBeginGatewaySuspendAdmission, +} from "../process/gateway-work-admission.js"; const mocks = vi.hoisted(() => ({ fetchWithSsrFGuard: vi.fn(async (_request: unknown) => ({ release: vi.fn() })), sendFailureNotificationAnnounce: vi.fn(), + sendCronAnnouncePayloadStrict: vi.fn(), })); vi.mock("../infra/net/fetch-guard.js", () => ({ @@ -19,10 +26,14 @@ vi.mock("../cron/delivery.js", async (importOriginal) => { return { ...actual, sendFailureNotificationAnnounce: mocks.sendFailureNotificationAnnounce, + sendCronAnnouncePayloadStrict: mocks.sendCronAnnouncePayloadStrict, }; }); -import { dispatchGatewayCronFinishedNotifications } from "./server-cron-notifications.js"; +import { + dispatchGatewayCronFinishedNotifications, + sendGatewayCronFailureAlert, +} from "./server-cron-notifications.js"; function requireRecord(value: unknown, label: string): Record { if (!value || typeof value !== "object") { @@ -44,9 +55,184 @@ function webhookRequestBody() { return JSON.parse(init.body); } +function createVoidDeferred(): { promise: Promise; resolve: () => void } { + let resolve = () => {}; + const promise = new Promise((resolvePromise) => { + resolve = resolvePromise; + }); + return { promise, resolve }; +} + +function createWebhookJob(delivery: NonNullable): CronJob { + return { + id: "cron-notification-admission", + name: "notification admission", + enabled: true, + createdAtMs: 1, + updatedAtMs: 1, + schedule: { kind: "every", everyMs: 60_000 }, + sessionTarget: "isolated", + wakeMode: "next-heartbeat", + payload: { kind: "agentTurn", message: "hello" }, + delivery, + state: {}, + }; +} + describe("dispatchGatewayCronFinishedNotifications", () => { beforeEach(() => { + resetGatewayWorkAdmission(); vi.clearAllMocks(); + mocks.fetchWithSsrFGuard.mockImplementation(async () => ({ release: vi.fn() })); + mocks.sendFailureNotificationAnnounce.mockResolvedValue(undefined); + mocks.sendCronAnnouncePayloadStrict.mockResolvedValue(undefined); + }); + + afterEach(() => { + resetGatewayWorkAdmission(); + }); + + it("independently admits detached completion webhook delivery", async () => { + const deferred = createVoidDeferred(); + mocks.fetchWithSsrFGuard.mockImplementationOnce(async () => { + await deferred.promise; + return { release: vi.fn() }; + }); + const job = createWebhookJob({ + mode: "webhook", + to: "https://example.invalid/cron", + }); + const parentAdmission = tryBeginGatewayRootWorkAdmission(); + expect(parentAdmission).not.toBeNull(); + if (!parentAdmission) { + throw new Error("expected parent Gateway work admission"); + } + + try { + await parentAdmission.run(async () => { + dispatchGatewayCronFinishedNotifications({ + evt: { jobId: job.id, action: "finished", status: "ok", summary: "done" }, + job, + deps: {} as CliDeps, + logger: { warn: vi.fn() }, + resolveCronAgent: () => ({ agentId: "main", cfg: {} }), + }); + + await vi.waitFor(() => expect(mocks.fetchWithSsrFGuard).toHaveBeenCalledTimes(1)); + expect(getActiveGatewayRootWorkCount()).toBe(2); + }); + } finally { + parentAdmission.release(); + } + + expect(getActiveGatewayRootWorkCount()).toBe(1); + deferred.resolve(); + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + }); + + it("independently admits immediate failure alerts", async () => { + const deferred = createVoidDeferred(); + mocks.sendCronAnnouncePayloadStrict.mockImplementationOnce(async () => { + await deferred.promise; + }); + const job = createWebhookJob({ mode: "announce", channel: "discord", to: "channel:ops" }); + + const delivery = sendGatewayCronFailureAlert({ + deps: {} as CliDeps, + logger: { warn: vi.fn() }, + resolveCronAgent: () => ({ agentId: "main", cfg: {} }), + job, + text: "cron failed", + channel: "discord", + to: "channel:ops", + mode: "announce", + }); + + await vi.waitFor(() => expect(mocks.sendCronAnnouncePayloadStrict).toHaveBeenCalledOnce()); + expect(getActiveGatewayRootWorkCount()).toBe(1); + deferred.resolve(); + await delivery; + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); + + it("defers detached completion delivery while suspension is prepared", async () => { + const job = createWebhookJob({ + mode: "webhook", + to: "https://example.invalid/cron", + }); + const suspensionAdmission = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspensionAdmission?.commit()).toBe(true); + + dispatchGatewayCronFinishedNotifications({ + evt: { jobId: job.id, action: "finished", status: "ok", summary: "done" }, + job, + deps: {} as CliDeps, + logger: { warn: vi.fn() }, + resolveCronAgent: () => ({ agentId: "main", cfg: {} }), + }); + + await Promise.resolve(); + expect(mocks.fetchWithSsrFGuard).not.toHaveBeenCalled(); + expect(getActiveGatewayRootWorkCount()).toBe(0); + + expect(suspensionAdmission?.release()).toBe(true); + await vi.waitFor(() => expect(mocks.fetchWithSsrFGuard).toHaveBeenCalledTimes(1)); + }); + + it("independently admits failure destination webhook delivery", async () => { + const deferred = createVoidDeferred(); + mocks.fetchWithSsrFGuard.mockImplementationOnce(async () => { + await deferred.promise; + return { release: vi.fn() }; + }); + const job = createWebhookJob({ + mode: "announce", + channel: "last", + failureDestination: { + mode: "webhook", + to: "https://example.invalid/failure", + }, + }); + + dispatchGatewayCronFinishedNotifications({ + evt: { jobId: job.id, action: "finished", status: "error", error: "boom" }, + job, + deps: {} as CliDeps, + logger: { warn: vi.fn() }, + resolveCronAgent: () => ({ agentId: "main", cfg: {} }), + }); + + await vi.waitFor(() => expect(mocks.fetchWithSsrFGuard).toHaveBeenCalledTimes(1)); + expect(getActiveGatewayRootWorkCount()).toBe(1); + deferred.resolve(); + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + }); + + it("independently admits failure destination announce delivery", async () => { + const deferred = createVoidDeferred(); + mocks.sendFailureNotificationAnnounce.mockImplementationOnce(() => deferred.promise); + const job = createWebhookJob({ + mode: "announce", + channel: "last", + failureDestination: { + mode: "announce", + channel: "telegram", + to: "-1001234567890", + }, + }); + + dispatchGatewayCronFinishedNotifications({ + evt: { jobId: job.id, action: "finished", status: "error", error: "boom" }, + job, + deps: {} as CliDeps, + logger: { warn: vi.fn() }, + resolveCronAgent: () => ({ agentId: "main", cfg: {} }), + }); + + await vi.waitFor(() => expect(mocks.sendFailureNotificationAnnounce).toHaveBeenCalledTimes(1)); + expect(getActiveGatewayRootWorkCount()).toBe(1); + deferred.resolve(); + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); }); it("redacts invalid completion webhook targets in warnings", () => { diff --git a/src/gateway/server-cron-notifications.ts b/src/gateway/server-cron-notifications.ts index 12ec5ed531f9..cb849a6ffe76 100644 --- a/src/gateway/server-cron-notifications.ts +++ b/src/gateway/server-cron-notifications.ts @@ -21,6 +21,7 @@ import { normalizeHttpWebhookUrl } from "../cron/webhook-url.js"; import { formatErrorMessage } from "../infra/errors.js"; import { fetchWithSsrFGuard } from "../infra/net/fetch-guard.js"; import { SsrFBlockedError } from "../infra/net/ssrf.js"; +import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js"; const CRON_WEBHOOK_TIMEOUT_MS = 10_000; @@ -38,6 +39,19 @@ type CronWebhookTarget = { source: "delivery" | "completionDestination"; }; +type CronFailureAlertParams = { + deps: CliDeps; + logger: CronLogger; + resolveCronAgent: CronAgentResolver; + webhookToken?: unknown; + job: CronJob; + text: string; + channel: CronMessageChannel; + to?: string; + mode?: "announce" | "webhook"; + accountId?: string; +}; + function redactWebhookUrl(url: string): string { try { const parsed = new URL(url); @@ -232,19 +246,30 @@ async function postCronWebhook(params: { } } -/** Sends the immediate failure alert for cron jobs that failed before normal completion delivery. */ -export async function sendGatewayCronFailureAlert(params: { - deps: CliDeps; +/** Detached sends outlive cron ticks; own roots block mid-delivery suspension snapshots. */ +function dispatchDetachedCronNotification(params: { + jobId: string; logger: CronLogger; - resolveCronAgent: CronAgentResolver; - webhookToken?: unknown; - job: CronJob; - text: string; - channel: CronMessageChannel; - to?: string; - mode?: "announce" | "webhook"; - accountId?: string; -}): Promise { + deliver: () => Promise; +}): void { + void runWithGatewayIndependentRootWorkAdmission(params.deliver).catch((err: unknown) => { + params.logger.warn( + { jobId: params.jobId, err: formatErrorMessage(err) }, + "cron: detached notification delivery failed", + ); + }); +} + +/** Sends the immediate failure alert for cron jobs that failed before normal completion delivery. */ +export async function sendGatewayCronFailureAlert(params: CronFailureAlertParams): Promise { + await runWithGatewayIndependentRootWorkAdmission(async () => { + await sendGatewayCronFailureAlertUnderAdmission(params); + }); +} + +async function sendGatewayCronFailureAlertUnderAdmission( + params: CronFailureAlertParams, +): Promise { const { agentId, cfg: runtimeConfig } = params.resolveCronAgent(params.job.agentId); const webhookToken = normalizeOptionalString(params.webhookToken); @@ -355,17 +380,20 @@ export function dispatchGatewayCronFinishedNotifications(params: { const payload = buildCronFinishedWebhookPayload(redactedWebhookEvent); // Completion notification fanout is best-effort; the cron service has // already recorded the run result and must not wait on slow webhooks. - void (async () => { - await postCronWebhook({ - webhookUrl: webhookTarget.url, - webhookToken, - payload, - logContext: { jobId: params.evt.jobId, source: webhookTarget.source }, - blockedLog: "cron: webhook delivery blocked by SSRF guard", - failedLog: "cron: webhook delivery failed", - logger: params.logger, - }); - })(); + dispatchDetachedCronNotification({ + jobId: params.evt.jobId, + logger: params.logger, + deliver: () => + postCronWebhook({ + webhookUrl: webhookTarget.url, + webhookToken, + payload, + logContext: { jobId: params.evt.jobId, source: webhookTarget.source }, + blockedLog: "cron: webhook delivery blocked by SSRF guard", + failedLog: "cron: webhook delivery failed", + logger: params.logger, + }), + }); } } @@ -393,9 +421,10 @@ function dispatchCronFailureDestinationNotifications(params: { return; } - const failureDest = resolveFailureDestination(params.job, params.globalFailureDestination); - const deliverySessionKey = resolveCronDeliverySessionKey(params.job); - const failurePayload = buildCronFailureWebhookPayload({ evt: params.evt, job: params.job }); + const job = params.job; + const failureDest = resolveFailureDestination(job, params.globalFailureDestination); + const deliverySessionKey = resolveCronDeliverySessionKey(job); + const failurePayload = buildCronFailureWebhookPayload({ evt: params.evt, job }); if (failureDest) { if (failureDest.mode === "webhook" && failureDest.to) { @@ -403,17 +432,20 @@ function dispatchCronFailureDestinationNotifications(params: { if (webhookUrl) { // Failure destinations mirror completion webhooks: notify in the // background and log failures without rewriting the cron event result. - void (async () => { - await postCronWebhook({ - webhookUrl, - webhookToken: params.webhookToken, - payload: failurePayload, - logContext: { jobId: params.evt.jobId }, - blockedLog: "cron: failure destination webhook blocked by SSRF guard", - failedLog: "cron: failure destination webhook failed", - logger: params.logger, - }); - })(); + dispatchDetachedCronNotification({ + jobId: params.evt.jobId, + logger: params.logger, + deliver: () => + postCronWebhook({ + webhookUrl, + webhookToken: params.webhookToken, + payload: failurePayload, + logContext: { jobId: params.evt.jobId }, + blockedLog: "cron: failure destination webhook blocked by SSRF guard", + failedLog: "cron: failure destination webhook failed", + logger: params.logger, + }), + }); } else { params.logger.warn( { @@ -427,44 +459,54 @@ function dispatchCronFailureDestinationNotifications(params: { } if (failureDest.mode === "announce") { - const { agentId, cfg: runtimeConfig } = params.resolveCronAgent(params.job.agentId); - void sendFailureNotificationAnnounce( - params.deps, - runtimeConfig, - agentId, - params.job.id, - { - channel: failureDest.channel, - to: failureDest.to, - accountId: failureDest.accountId, - sessionKey: deliverySessionKey, - // A configured failure route is already explicit; keep the cron run - // session only for context, not for reattaching the primary topic. - inheritSessionThread: false, - }, - `⚠️ ${failurePayload.message}`, - ); + const { agentId, cfg: runtimeConfig } = params.resolveCronAgent(job.agentId); + dispatchDetachedCronNotification({ + jobId: job.id, + logger: params.logger, + deliver: () => + sendFailureNotificationAnnounce( + params.deps, + runtimeConfig, + agentId, + job.id, + { + channel: failureDest.channel, + to: failureDest.to, + accountId: failureDest.accountId, + sessionKey: deliverySessionKey, + // A configured failure route is already explicit; keep the cron run + // session only for context, not for reattaching the primary topic. + inheritSessionThread: false, + }, + `⚠️ ${failurePayload.message}`, + ), + }); } return; } - const primaryPlan = resolveCronDeliveryPlan(params.job); + const primaryPlan = resolveCronDeliveryPlan(job); if (primaryPlan.mode !== "announce" || !primaryPlan.requested) { return; } - const { agentId, cfg: runtimeConfig } = params.resolveCronAgent(params.job.agentId); - void sendFailureNotificationAnnounce( - params.deps, - runtimeConfig, - agentId, - params.job.id, - { - channel: primaryPlan.channel, - to: primaryPlan.to, - accountId: primaryPlan.accountId, - sessionKey: deliverySessionKey, - }, - `⚠️ ${failurePayload.message}`, - ); + const { agentId, cfg: runtimeConfig } = params.resolveCronAgent(job.agentId); + dispatchDetachedCronNotification({ + jobId: job.id, + logger: params.logger, + deliver: () => + sendFailureNotificationAnnounce( + params.deps, + runtimeConfig, + agentId, + job.id, + { + channel: primaryPlan.channel, + to: primaryPlan.to, + accountId: primaryPlan.accountId, + sessionKey: deliverySessionKey, + }, + `⚠️ ${failurePayload.message}`, + ), + }); } diff --git a/src/gateway/server-cron.test.ts b/src/gateway/server-cron.test.ts index 1bfeacd001a8..664f3da05f19 100644 --- a/src/gateway/server-cron.test.ts +++ b/src/gateway/server-cron.test.ts @@ -6,6 +6,10 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import type { CliDeps } from "../cli/deps.js"; import type { OpenClawConfig } from "../config/config.js"; import { SsrFBlockedError } from "../infra/net/ssrf.js"; +import { + getActiveGatewayRootWorkCount, + resetGatewayWorkAdmission, +} from "../process/gateway-work-admission.js"; import { createDeferred } from "../test-utils/deferred.js"; type RunCronIsolatedAgentTurnMock = (params: { @@ -435,7 +439,6 @@ describe("buildGatewayCronService", () => { vi.setSystemTime(new Date("2026-07-10T12:00:00.000Z")); const cfg = createCronConfig("server-cron-hook-scheduled"); loadConfigMock.mockReturnValue(cfg); - const state = buildGatewayCronService({ cfg, deps: {} as CliDeps, @@ -486,6 +489,39 @@ describe("buildGatewayCronService", () => { } }); + it("keeps detached cron_changed hooks root-admitted until they settle", async () => { + resetGatewayWorkAdmission(); + const deferred = createDeferred(); + runCronChangedMock.mockImplementationOnce(async () => await deferred.promise); + const cfg = createCronConfig("server-cron-hook-admission"); + loadConfigMock.mockReturnValue(cfg); + const state = buildGatewayCronService({ + cfg, + deps: {} as CliDeps, + broadcast: () => {}, + }); + + try { + await state.cron.add({ + name: "held hook", + enabled: true, + schedule: { kind: "every", everyMs: 60_000 }, + sessionTarget: "main", + wakeMode: "next-heartbeat", + payload: { kind: "systemEvent", text: "hello" }, + }); + await vi.waitFor(() => expect(runCronChangedMock).toHaveBeenCalledTimes(1)); + expect(getActiveGatewayRootWorkCount()).toBe(1); + + deferred.resolve(); + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + } finally { + deferred.resolve(); + state.cron.stop(); + resetGatewayWorkAdmission(); + } + }); + it("cron_changed removed events include the deleted job snapshot", async () => { const cfg = createCronConfig("server-cron-hook-removed"); loadConfigMock.mockReturnValue(cfg); diff --git a/src/gateway/server-cron.ts b/src/gateway/server-cron.ts index c7b7d5a1509e..eb749c8951f8 100644 --- a/src/gateway/server-cron.ts +++ b/src/gateway/server-cron.ts @@ -49,6 +49,7 @@ import type { PluginHookGatewayCronService, PluginHookGatewayContext, } from "../plugins/hook-types.js"; +import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js"; import { getProcessSupervisor } from "../process/supervisor/index.js"; import { normalizeAgentId, @@ -63,8 +64,16 @@ import { sendGatewayCronFailureAlert, } from "./server-cron-notifications.js"; +export type GatewayCronServiceContract = CronServiceContract & { + /** Temporarily disarm ticks without running startup recovery on resume. */ + pauseScheduling(): void; + resumeScheduling(): void; + /** Scheduler-owned work not represented by active cron run markers. */ + getSuspensionBlockerCount?(): number; +}; + export type GatewayCronState = { - cron: CronServiceContract; + cron: GatewayCronServiceContract; storePath: string; cronEnabled: boolean; reconcileExitWatchers?: () => Promise; @@ -371,7 +380,11 @@ export function buildGatewayCronService(params: { config: getRuntimeConfig(), getCron: () => cron as PluginHookGatewayCronService, }; - void hookRunner.runCronChanged(evt, hookCtx).catch((err: unknown) => { + // Hook execution is detached from the cron mutation/tick that emitted it. + // Keep the whole plugin callback visible until its user-state effects settle. + void runWithGatewayIndependentRootWorkAdmission(async () => { + await hookRunner.runCronChanged(evt, hookCtx); + }).catch((err: unknown) => { cronLogger.warn( { err: formatErrorMessage(err), jobId: evt.jobId }, "cron_changed hook failed", @@ -383,12 +396,19 @@ export function buildGatewayCronService(params: { const exitWatchersRef: { current: ReturnType | undefined } = { current: undefined, }; + let exitWatcherReconciliations = 0; + let exitWatcherGeneration = 0; const reconcileExitWatchers = async () => { - if (!exitWatchersRef.current) { - return; - } + const generation = exitWatcherGeneration; + exitWatcherReconciliations += 1; try { + if (!exitWatchersRef.current) { + return; + } const result = await cron.list({ includeDisabled: true }); + if (generation !== exitWatcherGeneration) { + return; + } const jobs: CronJob[] = Array.isArray(result) ? result : (result as { jobs: CronJob[] }).jobs; reconcileCronExitWatchers({ cronEnabled, @@ -397,6 +417,8 @@ export function buildGatewayCronService(params: { }); } catch (err) { cronLogger.warn({ err: String(err) }, "cron-exit: reconcile failed"); + } finally { + exitWatcherReconciliations -= 1; } }; @@ -728,33 +750,35 @@ export function buildGatewayCronService(params: { globalFailureDestination: params.cfg.cron?.failureDestination, }); - void appendCronRunLog({ - storePath, - entry: { - ts: Date.now(), - jobId: evt.jobId, - action: "finished", - status: evt.status, - error: evt.error, - summary: evt.summary, - diagnostics: evt.diagnostics, - delivered: evt.delivered, - deliveryStatus: evt.deliveryStatus, - deliveryError: evt.deliveryError, - failureNotificationDelivery: evt.failureNotificationDelivery, - delivery: evt.delivery, - sessionId: evt.sessionId, - sessionKey: evt.sessionKey, - runId: evt.runId, - runAtMs: evt.runAtMs, - durationMs: evt.durationMs, - nextRunAtMs: evt.nextRunAtMs, - triggerFired: evt.triggerFired, - model: evt.model, - provider: evt.provider, - usage: evt.usage, - }, - opts: { keepLines: runLogPrune.keepLines }, + void runWithGatewayIndependentRootWorkAdmission(async () => { + await appendCronRunLog({ + storePath, + entry: { + ts: Date.now(), + jobId: evt.jobId, + action: "finished", + status: evt.status, + error: evt.error, + summary: evt.summary, + diagnostics: evt.diagnostics, + delivered: evt.delivered, + deliveryStatus: evt.deliveryStatus, + deliveryError: evt.deliveryError, + failureNotificationDelivery: evt.failureNotificationDelivery, + delivery: evt.delivery, + sessionId: evt.sessionId, + sessionKey: evt.sessionKey, + runId: evt.runId, + runAtMs: evt.runAtMs, + durationMs: evt.durationMs, + nextRunAtMs: evt.nextRunAtMs, + triggerFired: evt.triggerFired, + model: evt.model, + provider: evt.provider, + usage: evt.usage, + }, + opts: { keepLines: runLogPrune.keepLines }, + }); }).catch((err: unknown) => { cronLogger.warn( { err: String(err), storePath, jobId: evt.jobId }, @@ -767,19 +791,31 @@ export function buildGatewayCronService(params: { exitWatchersRef.current = createCronExitWatchers({ getProcessSupervisor, - persistCompletion: async (jobId) => { - await cron.update(jobId, { enabled: false }); - }, - fireOnExit: (job, exit) => - fireOnExitJob(job, exit, { - run: (jobId, payload) => cron.run(jobId, "force", payload ? { payload } : undefined), + persistCompletion: async (jobId) => + await runWithGatewayIndependentRootWorkAdmission(async () => { + await cron.update(jobId, { enabled: false }); }), + fireOnExit: (job, exit) => + runWithGatewayIndependentRootWorkAdmission(async () => + fireOnExitJob(job, exit, { + run: (jobId, payload) => cron.run(jobId, "force", payload ? { payload } : undefined), + }), + ), logger: cronLogger, }); + const getCronSuspensionBlockerCount = cron.getSuspensionBlockerCount.bind(cron); + cron.getSuspensionBlockerCount = () => + getCronSuspensionBlockerCount() + + exitWatcherReconciliations + + (exitWatchersRef.current?.activeJobIds().length ?? 0); + const stopExitWatchers = () => { + exitWatcherGeneration += 1; + exitWatchersRef.current?.cancelAll(); + }; const stopCron = cron.stop.bind(cron); cron.stop = () => { stopCron(); - exitWatchersRef.current?.cancelAll(); + stopExitWatchers(); }; return { @@ -787,6 +823,6 @@ export function buildGatewayCronService(params: { storePath, cronEnabled, reconcileExitWatchers, - stopExitWatchers: () => exitWatchersRef.current?.cancelAll(), + stopExitWatchers, }; } diff --git a/src/gateway/server-http.probe.test.ts b/src/gateway/server-http.probe.test.ts index 15740834fd26..8a91cd8b96f8 100644 --- a/src/gateway/server-http.probe.test.ts +++ b/src/gateway/server-http.probe.test.ts @@ -1,6 +1,18 @@ // Server HTTP probe tests cover readiness, health, disabled compat routes, and // auth handling through the in-memory HTTP harness. +import type { IncomingMessage, ServerResponse } from "node:http"; import { describe, expect, it, vi } from "vitest"; +import { + prepareGatewaySuspend, + resetGatewaySuspendCoordinatorForTest, + resumeGatewaySuspend, +} from "../infra/gateway-suspend-coordinator.js"; +import { isGatewayDraining } from "../process/command-queue.js"; +import { + getActiveGatewayRootWorkCount, + resetGatewayWorkAdmission, +} from "../process/gateway-work-admission.js"; +import type { ChannelManager } from "./server-channels.js"; import { AUTH_TOKEN, AUTH_NONE, @@ -9,7 +21,7 @@ import { dispatchRequest, withGatewayServer, } from "./server-http.test-harness.js"; -import type { ReadinessChecker } from "./server/readiness.js"; +import { createReadinessChecker, type ReadinessChecker } from "./server/readiness.js"; import { withTempConfig } from "./test-temp-config.js"; type GatewayServerHarness = Parameters[0]; @@ -44,6 +56,160 @@ describe("gateway OpenAI-compatible disabled HTTP routes", () => { }); describe("gateway probe endpoints", () => { + it("keeps liveness green while a prepared suspension lease makes readiness red", async () => { + resetGatewaySuspendCoordinatorForTest(); + resetGatewayWorkAdmission(); + const channelManager = { + getRuntimeSnapshot: () => ({ channels: {}, channelAccounts: {} }), + getAutostartSuppression: () => null, + } as unknown as ChannelManager; + const getReadiness = createReadinessChecker({ + channelManager, + startedAt: Date.now(), + getGatewayDraining: isGatewayDraining, + cacheTtlMs: 0, + }); + + try { + await withGatewayServer({ + prefix: "probe-suspension-lease", + resolvedAuth: AUTH_NONE, + overrides: { getReadiness, openAiChatCompletionsEnabled: true }, + run: async (server) => { + const prepared = prepareGatewaySuspend({ + requestId: "request-readiness-probe", + pauseScheduling: vi.fn(), + resumeScheduling: vi.fn(), + createSuspensionId: () => "suspension-readiness-probe", + inspect: { + getQueueSize: () => 0, + getPendingReplies: () => 0, + getEmbeddedRuns: () => 0, + getCronRuns: () => 0, + getActiveTasks: () => 0, + getTaskBlockers: () => [], + getRootRequests: () => 0, + getSessionAdmissions: () => 0, + getSessionMutations: () => 0, + getChatRuns: () => 0, + getQueuedTurns: () => 0, + getTerminalPersistence: () => 0, + getTerminalSessions: () => 0, + }, + }); + if (prepared.status !== "ready") { + throw new Error(`expected prepared suspension, received ${prepared.status}`); + } + + const health = await sendGatewayRequest(server, { path: "/healthz" }); + expect(health.res.statusCode).toBe(200); + expect(JSON.parse(health.getBody())).toEqual({ ok: true, status: "live" }); + + const suspendedReadiness = await sendGatewayRequest(server, { path: "/readyz" }); + expect(suspendedReadiness.res.statusCode).toBe(503); + expect(JSON.parse(suspendedReadiness.getBody())).toMatchObject({ + ready: false, + failing: ["gateway-draining"], + }); + + const blockedChat = await sendGatewayRequest(server, { + path: "/v1/chat/completions", + method: "POST", + }); + expect(blockedChat.res.statusCode).toBe(503); + expect(JSON.parse(blockedChat.getBody())).toMatchObject({ + error: { code: "gateway_unavailable" }, + }); + + expect(resumeGatewaySuspend(prepared.suspensionId)).toEqual({ + ok: true, + status: "running", + resumed: true, + }); + + const resumedReadiness = await sendGatewayRequest(server, { path: "/readyz" }); + expect(resumedReadiness.res.statusCode).toBe(200); + expect(JSON.parse(resumedReadiness.getBody())).toMatchObject({ + ready: true, + failing: [], + }); + }, + }); + } finally { + resetGatewaySuspendCoordinatorForTest(); + resetGatewayWorkAdmission(); + } + }); + + it("keeps in-flight core HTTP work visible to suspension preparation", async () => { + resetGatewaySuspendCoordinatorForTest(); + resetGatewayWorkAdmission(); + let releaseWatch = () => {}; + let markWatchStarted = () => {}; + const watchStarted = new Promise((resolve) => { + markWatchStarted = resolve; + }); + const heldWatch = new Promise((resolve) => { + releaseWatch = resolve; + }); + const handleWatchNodeRequest = vi.fn(async (_req: IncomingMessage, res: ServerResponse) => { + markWatchStarted(); + await heldWatch; + res.statusCode = 200; + res.end("ok"); + return true; + }); + + try { + await withGatewayServer({ + prefix: "probe-http-work-admission", + resolvedAuth: AUTH_NONE, + overrides: { handleWatchNodeRequest }, + run: async (server) => { + const request = createRequest({ path: "/api/nodes/watch/node-1" }); + const response = createResponse(); + const pendingRequest = dispatchRequest(server, request, response.res); + await watchStarted; + expect(getActiveGatewayRootWorkCount()).toBe(1); + + const prepared = prepareGatewaySuspend({ + requestId: "request-http-work", + pauseScheduling: vi.fn(), + resumeScheduling: vi.fn(), + inspect: { + getQueueSize: () => 0, + getPendingReplies: () => 0, + getEmbeddedRuns: () => 0, + getCronRuns: () => 0, + getActiveTasks: () => 0, + getTaskBlockers: () => [], + getSessionAdmissions: () => 0, + getSessionMutations: () => 0, + getChatRuns: () => 0, + getQueuedTurns: () => 0, + getTerminalPersistence: () => 0, + getTerminalSessions: () => 0, + }, + }); + expect(prepared).toMatchObject({ + status: "busy", + reason: "active-work", + activeCount: 1, + }); + + releaseWatch(); + await pendingRequest; + expect(response.res.statusCode).toBe(200); + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + }, + }); + } finally { + releaseWatch(); + resetGatewaySuspendCoordinatorForTest(); + resetGatewayWorkAdmission(); + } + }); + it("returns detailed readiness payload for local /ready requests", async () => { const getReadiness: ReadinessChecker = () => ({ ready: true, diff --git a/src/gateway/server-http.ts b/src/gateway/server-http.ts index 371bfbfdd3d5..8e9faac17478 100644 --- a/src/gateway/server-http.ts +++ b/src/gateway/server-http.ts @@ -16,6 +16,10 @@ import { createDiagnosticTraceContext, runWithDiagnosticTraceContext, } from "../infra/diagnostic-trace-context.js"; +import { + isGatewayWorkAdmissionClosed, + tryBeginGatewayRootWorkAdmission, +} 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"; @@ -52,6 +56,7 @@ type PluginHttpRequestHandler = ( gatewayAuthSatisfied?: boolean; gatewayRequestAuth?: AuthorizedGatewayHttpRequest; gatewayRequestOperatorScopes?: readonly string[]; + gatewayRequestClientIp?: string; }, ) => Promise; @@ -66,6 +71,7 @@ type PluginHttpUpgradeHandler = ( gatewayAuthSatisfied?: boolean; gatewayRequestAuth?: AuthorizedGatewayHttpRequest; gatewayRequestOperatorScopes?: readonly string[]; + gatewayRequestClientIp?: string; }, ) => Promise; @@ -348,6 +354,35 @@ export async function runGatewayHttpRequestStages( return false; } +/** Runs one core HTTP user-work route under the same root fence as Gateway RPCs. */ +export async function runWithGatewayHttpWorkAdmission( + res: ServerResponse, + run: () => Promise | boolean, +): Promise { + const admission = tryBeginGatewayRootWorkAdmission(); + if (!admission) { + res.statusCode = 503; + res.setHeader("Content-Type", "application/json; charset=utf-8"); + res.setHeader("Cache-Control", "no-store"); + res.setHeader("Retry-After", "1"); + res.end( + JSON.stringify({ + error: { + message: "Gateway is temporarily unavailable while suspending or restarting", + type: "service_unavailable", + code: "gateway_unavailable", + }, + }), + ); + return true; + } + try { + return await admission.run(async () => await run()); + } finally { + admission.release(); + } +} + function buildPluginRequestStages(params: { req: IncomingMessage; res: ServerResponse; @@ -364,6 +399,11 @@ function buildPluginRequestStages(params: { if (!params.handlePluginRequest) { return []; } + const requestClientIp = resolveRequestClientIp( + params.req, + params.trustedProxies, + params.allowRealIpFallback, + ); let pluginGatewayAuthSatisfied = false; let pluginGatewayRequestAuth: AuthorizedGatewayHttpRequest | undefined; let pluginRequestOperatorScopes: string[] | undefined; @@ -422,6 +462,7 @@ function buildPluginRequestStages(params: { gatewayAuthSatisfied: pluginGatewayAuthSatisfied, gatewayRequestAuth: pluginGatewayRequestAuth, gatewayRequestOperatorScopes: pluginRequestOperatorScopes, + gatewayRequestClientIp: requestClientIp, }) ?? false ); }, @@ -559,94 +600,112 @@ export function createGatewayHttpServer(opts: { if (opts.handleWatchNodeRequest && scopedRequestPath.startsWith("/api/nodes/watch/")) { requestStages.push({ name: "watch-node", - run: () => opts.handleWatchNodeRequest?.(req, res) ?? false, + run: () => + runWithGatewayHttpWorkAdmission( + res, + () => opts.handleWatchNodeRequest?.(req, res) ?? false, + ), }); } if (openAiCompatEnabled && isOpenAiModelsPath(scopedRequestPath)) { requestStages.push({ name: "models", run: async () => - (await getModelsHttpModule()).handleOpenAiModelsHttpRequest(req, res, { - auth: resolvedAuthValue, - trustedProxies, - allowRealIpFallback, - rateLimiter, - }), + await runWithGatewayHttpWorkAdmission(res, async () => + (await getModelsHttpModule()).handleOpenAiModelsHttpRequest(req, res, { + auth: resolvedAuthValue, + trustedProxies, + allowRealIpFallback, + rateLimiter, + }), + ), }); } if (openAiCompatEnabled && isEmbeddingsPath(scopedRequestPath)) { requestStages.push({ name: "embeddings", run: async () => - (await getEmbeddingsHttpModule()).handleOpenAiEmbeddingsHttpRequest(req, res, { - auth: resolvedAuthValue, - trustedProxies, - allowRealIpFallback, - rateLimiter, - }), + await runWithGatewayHttpWorkAdmission(res, async () => + (await getEmbeddingsHttpModule()).handleOpenAiEmbeddingsHttpRequest(req, res, { + auth: resolvedAuthValue, + trustedProxies, + allowRealIpFallback, + rateLimiter, + }), + ), }); } if (isToolsInvokePath(scopedRequestPath)) { requestStages.push({ name: "tools-invoke", run: async () => - (await getToolsInvokeHttpModule()).handleToolsInvokeHttpRequest(req, res, { - auth: resolvedAuthValue, - trustedProxies, - allowRealIpFallback, - rateLimiter, - }), + await runWithGatewayHttpWorkAdmission(res, async () => + (await getToolsInvokeHttpModule()).handleToolsInvokeHttpRequest(req, res, { + auth: resolvedAuthValue, + trustedProxies, + allowRealIpFallback, + rateLimiter, + }), + ), }); } if (isSessionKillPath(scopedRequestPath)) { requestStages.push({ name: "sessions-kill", run: async () => - (await getSessionKillHttpModule()).handleSessionKillHttpRequest(req, res, { - auth: resolvedAuthValue, - trustedProxies, - allowRealIpFallback, - rateLimiter, - }), + await runWithGatewayHttpWorkAdmission(res, async () => + (await getSessionKillHttpModule()).handleSessionKillHttpRequest(req, res, { + auth: resolvedAuthValue, + trustedProxies, + allowRealIpFallback, + rateLimiter, + }), + ), }); } if (isSessionHistoryPath(scopedRequestPath)) { requestStages.push({ name: "sessions-history", run: async () => - (await getSessionHistoryHttpModule()).handleSessionHistoryHttpRequest(req, res, { - auth: resolvedAuthValue, - getResolvedAuth, - trustedProxies, - allowRealIpFallback, - rateLimiter, - }), + await runWithGatewayHttpWorkAdmission(res, async () => + (await getSessionHistoryHttpModule()).handleSessionHistoryHttpRequest(req, res, { + auth: resolvedAuthValue, + getResolvedAuth, + trustedProxies, + allowRealIpFallback, + rateLimiter, + }), + ), }); } if (openResponsesEnabled && isOpenResponsesPath(scopedRequestPath)) { requestStages.push({ name: "openresponses", run: async () => - (await getOpenResponsesHttpModule()).handleOpenResponsesHttpRequest(req, res, { - auth: resolvedAuthValue, - config: openResponsesConfig, - trustedProxies, - allowRealIpFallback, - rateLimiter, - }), + await runWithGatewayHttpWorkAdmission(res, async () => + (await getOpenResponsesHttpModule()).handleOpenResponsesHttpRequest(req, res, { + auth: resolvedAuthValue, + config: openResponsesConfig, + trustedProxies, + allowRealIpFallback, + rateLimiter, + }), + ), }); } if (openAiChatCompletionsEnabled && isOpenAiChatCompletionsPath(scopedRequestPath)) { requestStages.push({ name: "openai", run: async () => - (await getOpenAiHttpModule()).handleOpenAiHttpRequest(req, res, { - auth: resolvedAuthValue, - config: openAiChatCompletionsConfig, - trustedProxies, - allowRealIpFallback, - rateLimiter, - }), + await runWithGatewayHttpWorkAdmission(res, async () => + (await getOpenAiHttpModule()).handleOpenAiHttpRequest(req, res, { + auth: resolvedAuthValue, + config: openAiChatCompletionsConfig, + trustedProxies, + allowRealIpFallback, + rateLimiter, + }), + ), }); } if ( @@ -843,6 +902,7 @@ export function attachGatewayUpgradeHandler(opts: { const configSnapshot = getRuntimeConfig(); const trustedProxies = configSnapshot.gateway?.trustedProxies ?? []; const allowRealIpFallback = configSnapshot.gateway?.allowRealIpFallback === true; + const requestClientIp = resolveRequestClientIp(req, trustedProxies, allowRealIpFallback); const scopedNodeCapability = normalizePluginNodeCapabilityScopedUrl(req.url ?? "/"); if (scopedNodeCapability.malformedScopedPath) { writeUpgradeAuthFailure(socket, { ok: false, reason: "unauthorized" }); @@ -916,12 +976,21 @@ export function attachGatewayUpgradeHandler(opts: { gatewayAuthSatisfied: pluginGatewayAuthSatisfied, gatewayRequestAuth: pluginGatewayRequestAuth, gatewayRequestOperatorScopes: pluginGatewayRequestOperatorScopes, + gatewayRequestClientIp: requestClientIp, }) ) { return; } } - const preauthBudgetKey = resolveRequestClientIp(req, trustedProxies, allowRealIpFallback); + // 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. + if (isGatewayWorkAdmissionClosed()) { + writeUpgradeServiceUnavailable(socket, "Gateway websocket admission closed"); + socket.destroy(); + return; + } + const preauthBudgetKey = requestClientIp; if (wss.listenerCount("connection") === 0) { writeUpgradeServiceUnavailable(socket, "Gateway websocket handlers unavailable"); socket.destroy(); diff --git a/src/gateway/server-methods.control-plane-rate-limit.test.ts b/src/gateway/server-methods.control-plane-rate-limit.test.ts index f35caa2782b8..62b5a8855383 100644 --- a/src/gateway/server-methods.control-plane-rate-limit.test.ts +++ b/src/gateway/server-methods.control-plane-rate-limit.test.ts @@ -3,6 +3,10 @@ */ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { isRetryableGatewayStartupUnavailableError } from "../../packages/gateway-protocol/src/startup-unavailable.js"; +import { + resetGatewayWorkAdmission, + tryBeginGatewaySuspendAdmission, +} from "../process/gateway-work-admission.js"; import { testing as controlPlaneRateLimitTesting, resolveControlPlaneRateLimitKey, @@ -16,12 +20,14 @@ const noWebchat = () => false; describe("gateway control-plane write rate limit", () => { beforeEach(() => { controlPlaneRateLimitTesting.resetControlPlaneRateLimitState(); + resetGatewayWorkAdmission(); vi.useFakeTimers(); vi.setSystemTime(new Date("2026-02-19T00:00:00.000Z")); }); afterEach(() => { vi.useRealTimers(); + resetGatewayWorkAdmission(); controlPlaneRateLimitTesting.resetControlPlaneRateLimitState(); }); @@ -145,6 +151,57 @@ describe("gateway control-plane write rate limit", () => { expect(handlerCalls).toHaveBeenCalledTimes(4); }); + it("does not consume the write budget for requests refused during suspension", async () => { + const handlerCalls = vi.fn(); + const handler: GatewayRequestHandler = (opts) => { + handlerCalls(opts); + opts.respond(true, undefined, undefined); + }; + const context = buildContext(); + const client = buildClient(); + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + + for (let attempt = 0; attempt < 4; attempt += 1) { + const refused = await runRequest({ method: "config.patch", context, client, handler }); + expect(respondCall(refused)[2]).toMatchObject({ + code: "UNAVAILABLE", + details: { reason: "gateway-suspending" }, + }); + } + expect(suspension?.release()).toBe(true); + + const allowed = await runRequest({ method: "config.patch", context, client, handler }); + expect(allowed).toHaveBeenCalledWith(true, undefined, undefined); + expect(handlerCalls).toHaveBeenCalledOnce(); + }); + + it("keeps suspension preparation rate-limited while admission is closed", async () => { + const handlerCalls = vi.fn(); + const handler: GatewayRequestHandler = (opts) => { + handlerCalls(opts); + opts.respond(true, undefined, undefined); + }; + const context = buildContext(); + const client = buildClient(); + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + + await runRequest({ method: "gateway.suspend.prepare", context, client, handler }); + await runRequest({ method: "gateway.suspend.prepare", context, client, handler }); + await runRequest({ method: "gateway.suspend.prepare", context, client, handler }); + const blocked = await runRequest({ + method: "gateway.suspend.prepare", + context, + client, + handler, + }); + + expect(handlerCalls).toHaveBeenCalledTimes(3); + expect(respondCall(blocked)[2]).toMatchObject({ code: "UNAVAILABLE", retryable: true }); + expect(suspension?.release()).toBe(true); + }); + it.each(STARTUP_UNAVAILABLE_GATEWAY_METHODS)( "blocks startup-gated method %s before dispatch with a retryable startup error", async (method) => { diff --git a/src/gateway/server-methods.suspension-admission.test.ts b/src/gateway/server-methods.suspension-admission.test.ts new file mode 100644 index 000000000000..da0fd110544c --- /dev/null +++ b/src/gateway/server-methods.suspension-admission.test.ts @@ -0,0 +1,273 @@ +// Proves dispatcher root-work accounting and fail-closed suspension behavior. +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { + resetGatewaySuspendCoordinatorForTest, + resumeGatewaySuspend, +} from "../infra/gateway-suspend-coordinator.js"; +import { + getActiveGatewayRootWorkCount, + resetGatewayWorkAdmission, + tryBeginGatewayRootWorkAdmission, + tryBeginGatewaySuspendAdmission, +} from "../process/gateway-work-admission.js"; +import { + createGatewayMethodRegistry, + createPluginGatewayMethodDescriptor, +} from "./methods/registry.js"; +import { handleGatewayRequest } from "./server-methods.js"; +import { suspendHandlers } from "./server-methods/suspend.js"; +import type { GatewayRequestHandler } from "./server-methods/types.js"; + +function deferred() { + let resolve = () => {}; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +function dispatch(params: { + method: string; + scope: "operator.read" | "operator.write" | "operator.admin"; + handler: GatewayRequestHandler; + requestParams?: Record; + context?: Parameters[0]["context"]; +}) { + const respond = vi.fn(); + const methodRegistry = createGatewayMethodRegistry([ + createPluginGatewayMethodDescriptor({ + pluginId: "suspend-proof", + name: params.method, + handler: params.handler, + scope: params.scope, + }), + ]); + const request = handleGatewayRequest({ + req: { + type: "req", + id: `request-${params.method}`, + method: params.method, + params: params.requestParams ?? {}, + }, + respond, + client: { + connId: "conn-suspend-proof", + connect: { + role: "operator", + scopes: [params.scope], + client: { id: "cli", version: "test", platform: "linux", mode: "cli" }, + minProtocol: 1, + maxProtocol: 1, + }, + }, + isWebchatConnect: () => false, + context: + params.context ?? + ({ logGateway: { warn: vi.fn() } } as unknown as Parameters< + typeof handleGatewayRequest + >[0]["context"]), + methodRegistry, + }); + return { request, respond }; +} + +beforeEach(() => { + resetGatewaySuspendCoordinatorForTest(); + resetGatewayWorkAdmission(); +}); + +afterEach(() => { + resetGatewaySuspendCoordinatorForTest(); + resetGatewayWorkAdmission(); +}); + +describe("gateway request suspension admission", () => { + it("keeps preparation busy while a previously admitted handler is active", async () => { + const started = deferred(); + const finish = deferred(); + const handler = vi.fn(async ({ respond }) => { + started.resolve(); + await finish.promise; + respond(true, { ok: true }); + }); + const active = dispatch({ + method: "suspend-proof.run", + scope: "operator.write", + handler, + }); + await started.promise; + expect(getActiveGatewayRootWorkCount()).toBe(1); + + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(getActiveGatewayRootWorkCount()).toBe(1); + expect(suspension?.rollback()).toBe(true); + + finish.resolve(); + await active.request; + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); + + it("reports a concurrent root as busy then excludes its own prepare request", async () => { + const started = deferred(); + const finish = deferred(); + const active = dispatch({ + method: "suspend-proof.concurrent", + scope: "operator.write", + handler: async ({ respond }) => { + started.resolve(); + await finish.promise; + respond(true, { ok: true }); + }, + }); + await started.promise; + + const prepareHandler = suspendHandlers["gateway.suspend.prepare"]; + expect(prepareHandler).toBeTypeOf("function"); + if (!prepareHandler) { + throw new Error("expected gateway suspension prepare handler"); + } + const cron = { + pauseScheduling: vi.fn(), + resumeScheduling: vi.fn(), + getSuspensionBlockerCount: vi.fn(() => 0), + }; + const context = { + cron, + logGateway: { warn: vi.fn() }, + chatAbortControllers: new Map(), + chatQueuedTurns: new Map(), + } as unknown as Parameters[0]["context"]; + const busy = dispatch({ + method: "gateway.suspend.prepare", + scope: "operator.admin", + handler: prepareHandler, + requestParams: { requestId: "request-concurrent-root" }, + context, + }); + await busy.request; + + expect(busy.respond).toHaveBeenCalledWith( + true, + expect.objectContaining({ + status: "busy", + reason: "active-work", + activeCount: 1, + blockers: expect.arrayContaining([ + expect.objectContaining({ kind: "root-request", count: 1 }), + ]), + }), + ); + + finish.resolve(); + await active.request; + const ready = dispatch({ + method: "gateway.suspend.prepare", + scope: "operator.admin", + handler: prepareHandler, + requestParams: { requestId: "request-own-root-excluded" }, + context, + }); + await ready.request; + + expect(ready.respond).toHaveBeenCalledWith( + true, + expect.objectContaining({ + status: "ready", + activeCount: 0, + blockers: [], + }), + ); + const readyPayload = ready.respond.mock.calls[0]?.[1] as { suspensionId?: string } | undefined; + expect(readyPayload?.suspensionId).toBeTypeOf("string"); + expect(resumeGatewaySuspend(readyPayload?.suspensionId ?? "missing")).toMatchObject({ + ok: true, + resumed: true, + }); + }); + + it("rejects new read and write handlers outside the suspension allowlist", async () => { + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + + const writeHandler = vi.fn(); + const blocked = dispatch({ + method: "suspend-proof.write", + scope: "operator.write", + handler: writeHandler, + }); + await blocked.request; + expect(writeHandler).not.toHaveBeenCalled(); + expect(blocked.respond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ + code: "UNAVAILABLE", + retryable: true, + details: expect.objectContaining({ reason: "gateway-suspending" }), + }), + ); + + const readHandler = vi.fn(({ respond }) => { + respond(true, { state: "visible" }); + }); + const allowed = dispatch({ + method: "suspend-proof.read", + scope: "operator.read", + handler: readHandler, + }); + await allowed.request; + expect(readHandler).not.toHaveBeenCalled(); + expect(allowed.respond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ code: "UNAVAILABLE", retryable: true }), + ); + suspension?.release(); + }); + + it("keeps suspension status reachable while prepared", async () => { + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + const handler = vi.fn(({ respond }) => { + respond(true, { ok: true }); + }); + + const status = dispatch({ + method: "gateway.suspend.status", + scope: "operator.read", + handler, + }); + await status.request; + + expect(handler).toHaveBeenCalledOnce(); + expect(status.respond).toHaveBeenCalledWith(true, { ok: true }); + suspension?.release(); + }); + + it("rejects suspension preparation nested inside another root request", async () => { + const root = tryBeginGatewayRootWorkAdmission(); + expect(root?.ownsRoot).toBe(true); + const handler = vi.fn(); + + await root?.run(async () => { + const nested = dispatch({ + method: "gateway.suspend.prepare", + scope: "operator.admin", + handler, + }); + await nested.request; + expect(nested.respond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ + code: "UNAVAILABLE", + retryable: true, + details: expect.objectContaining({ reason: "nested-gateway-request" }), + }), + ); + }); + + root?.release(); + expect(handler).not.toHaveBeenCalled(); + }); +}); diff --git a/src/gateway/server-methods.ts b/src/gateway/server-methods.ts index d9654d4e3683..b2f3b501b0de 100644 --- a/src/gateway/server-methods.ts +++ b/src/gateway/server-methods.ts @@ -7,6 +7,11 @@ import { } from "../../packages/gateway-protocol/src/startup-unavailable.js"; import { getPluginRegistryState } from "../plugins/runtime-state.js"; import { withPluginRuntimeGatewayRequestScope } from "../plugins/runtime/gateway-request-scope.js"; +import { + getGatewaySuspendAdmissionPhase, + isGatewayRestartDraining, + tryBeginGatewayRootWorkAdmission, +} from "../process/gateway-work-admission.js"; import { formatControlPlaneActor, resolveControlPlaneActor } from "./control-plane-audit.js"; import { consumeControlPlaneWriteBudget } from "./control-plane-rate-limit.js"; import { @@ -190,6 +195,10 @@ const loadRestartHandlers = lazyHandlerModule( () => import("./server-methods/restart.js"), (module) => module.restartHandlers, ); +const loadSuspendHandlers = lazyHandlerModule( + () => import("./server-methods/suspend.js"), + (module) => module.suspendHandlers, +); const loadSendHandlers = lazyHandlerModule( () => import("./server-methods/send.js"), (module) => module.sendHandlers, @@ -306,6 +315,16 @@ function authorizeGatewayMethod( return null; } +const SUSPEND_CONTROL_METHODS = new Set([ + "gateway.suspend.prepare", + "gateway.suspend.status", + "gateway.suspend.resume", +]); + +function isGatewayMethodAllowedDuringSuspension(method: string): boolean { + return SUSPEND_CONTROL_METHODS.has(method); +} + export const coreGatewayHandlers: GatewayRequestHandlers = { ...createLazyCoreHandlers({ methods: ["connect"], @@ -665,6 +684,10 @@ export const coreGatewayHandlers: GatewayRequestHandlers = { methods: ["gateway.restart.request", "gateway.restart.preflight"], loadHandlers: loadRestartHandlers, }), + ...createLazyCoreHandlers({ + methods: ["gateway.suspend.prepare", "gateway.suspend.status", "gateway.suspend.resume"], + loadHandlers: loadSuspendHandlers, + }), ...createLazyCoreHandlers({ methods: ["message.action", "send", "poll"], loadHandlers: loadSendHandlers, @@ -781,33 +804,40 @@ export async function handleGatewayRequest( ); return; } - if (methodRegistry.isControlPlaneWrite(req.method)) { - const budget = consumeControlPlaneWriteBudget({ client }); - if (!budget.allowed) { - // Control-plane writes mutate gateway-wide state; rate limit before handler lookup so - // plugin and aux write methods share the same protection. - const actor = resolveControlPlaneActor(client); - context.logGateway.warn( - `control-plane write rate-limited method=${req.method} ${formatControlPlaneActor(actor)} retryAfterMs=${budget.retryAfterMs} key=${budget.key}`, - ); - respond( - false, - undefined, - errorShape( - ErrorCodes.UNAVAILABLE, - `rate limit exceeded for ${req.method}; retry after ${Math.ceil(budget.retryAfterMs / 1000)}s`, - { - retryable: true, - retryAfterMs: budget.retryAfterMs, - details: { - method: req.method, - limit: "3 per 60s", - }, - }, - ), - ); - return; + const rejectRateLimitedControlPlaneWrite = (): boolean => { + if (!methodRegistry.isControlPlaneWrite(req.method)) { + return false; } + const budget = consumeControlPlaneWriteBudget({ client }); + if (budget.allowed) { + return false; + } + const actor = resolveControlPlaneActor(client); + context.logGateway.warn( + `control-plane write rate-limited method=${req.method} ${formatControlPlaneActor(actor)} retryAfterMs=${budget.retryAfterMs} key=${budget.key}`, + ); + respond( + false, + undefined, + errorShape( + ErrorCodes.UNAVAILABLE, + `rate limit exceeded for ${req.method}; retry after ${Math.ceil(budget.retryAfterMs / 1000)}s`, + { + retryable: true, + retryAfterMs: budget.retryAfterMs, + details: { + method: req.method, + limit: "3 per 60s", + }, + }, + ), + ); + return true; + }; + const isSuspendPrepare = req.method === "gateway.suspend.prepare"; + if (isSuspendPrepare && rejectRateLimitedControlPlaneWrite()) { + // Preparation must stay protected even before it owns the root admission that it closes. + return; } const handler = methodRegistry.getHandler(req.method) as GatewayRequestHandler | undefined; if (!handler) { @@ -818,6 +848,50 @@ export async function handleGatewayRequest( ); return; } + const rootWorkAdmission = tryBeginGatewayRootWorkAdmission(); + if ( + req.method === "gateway.suspend.prepare" && + rootWorkAdmission && + !rootWorkAdmission.ownsRoot + ) { + respond( + false, + undefined, + errorShape(ErrorCodes.UNAVAILABLE, "gateway suspension cannot begin from a nested request", { + retryable: true, + retryAfterMs: 1_000, + details: { method: req.method, reason: "nested-gateway-request" }, + }), + ); + return; + } + if (!rootWorkAdmission && !isGatewayMethodAllowedDuringSuspension(req.method)) { + const restartDraining = isGatewayRestartDraining(); + respond( + false, + undefined, + errorShape( + ErrorCodes.UNAVAILABLE, + `${req.method} unavailable during gateway ${restartDraining ? "restart" : "suspension"}`, + { + retryable: true, + retryAfterMs: 1_000, + details: { + method: req.method, + reason: restartDraining ? "gateway-restarting" : "gateway-suspending", + phase: getGatewaySuspendAdmissionPhase(), + }, + }, + ), + ); + return; + } + if (!isSuspendPrepare && rejectRateLimitedControlPlaneWrite()) { + // A closed admission must reject first so refused writes do not exhaust the controller's + // budget and strand it behind rate limiting after suspension resumes. + rootWorkAdmission?.release(); + return; + } const invokeHandler = () => handler({ req, @@ -831,5 +905,18 @@ export async function handleGatewayRequest( // subagent methods (e.g. context engine tools spawning sub-agents // during tool execution) can dispatch back into the gateway. // The scope also carries caller identity into plugin-owned gateway methods. - await withPluginRuntimeGatewayRequestScope({ context, client, isWebchatConnect }, invokeHandler); + const invokeWithRequestScope = async () => + await withPluginRuntimeGatewayRequestScope( + { context, client, isWebchatConnect }, + invokeHandler, + ); + if (!rootWorkAdmission) { + await invokeWithRequestScope(); + return; + } + try { + await rootWorkAdmission.run(invokeWithRequestScope); + } finally { + rootWorkAdmission.release(); + } } diff --git a/src/gateway/server-methods/agent.ts b/src/gateway/server-methods/agent.ts index 2d5c3a640214..43fe67f9eba2 100644 --- a/src/gateway/server-methods/agent.ts +++ b/src/gateway/server-methods/agent.ts @@ -116,6 +116,7 @@ import { } from "../../infra/voicewake-routing.js"; import type { PromptImageOrderEntry } from "../../media/prompt-image-order.js"; import type { PluginHookSessionEndReason } from "../../plugins/hook-types.js"; +import { retainGatewayRootWorkAdmissionContinuation } from "../../process/gateway-work-admission.js"; import { classifySessionKeyShape, isAcpSessionKey, @@ -3543,9 +3544,12 @@ export const agentHandlers: GatewayRequestHandlers = { if (!activeRunAbort.registered) { activeGatewayWorkAdmission.release(); } + let releaseGatewayRootContinuation: (() => void) | undefined; const cleanupAdmittedRun: typeof activeRunAbort.cleanup = (options) => { activeRunAbort.cleanup(options); activeGatewayWorkAdmission.release(); + releaseGatewayRootContinuation?.(); + releaseGatewayRootContinuation = undefined; }; if (activeRunAbort.registered) { retainEmbeddedAgentRunAbortabilityForRunId(runId); @@ -3644,6 +3648,9 @@ export const agentHandlers: GatewayRequestHandlers = { // is scheduled out of this request handler so immediate agent.wait calls // can reach the gateway before the pre-turn runner monopolizes the loop. gatewayAdmissionTransferred = true; + // Reserve the detached run before this request releases its root. Otherwise + // its inherited ALS context becomes retired and rejects subordinate work. + releaseGatewayRootContinuation = retainGatewayRootWorkAdmissionContinuation() ?? undefined; void activeGatewayWorkAdmission.run(async () => { await yieldAfterAgentAcceptedAck(); diff --git a/src/gateway/server-methods/chat.ts b/src/gateway/server-methods/chat.ts index 516dd2631958..c262f521fe1c 100644 --- a/src/gateway/server-methods/chat.ts +++ b/src/gateway/server-methods/chat.ts @@ -104,6 +104,10 @@ import { deleteMediaBuffer, MEDIA_MAX_BYTES, type SavedMedia } from "../../media import { createChannelMessageReplyPipeline } from "../../plugin-sdk/channel-outbound.js"; import type { ChannelRouteRef } from "../../plugin-sdk/channel-route.js"; import { isPluginOwnedSessionBindingRecord } from "../../plugins/conversation-binding.js"; +import { + retainGatewayRootWorkAdmissionContinuation, + runWithGatewayIndependentRootWorkContinuation, +} from "../../process/gateway-work-admission.js"; import { normalizeAgentId, scopeLegacySessionKeyToAgent } from "../../routing/session-key.js"; import { normalizeInputProvenance, type InputProvenance } from "../../sessions/input-provenance.js"; import { resolveSendPolicy } from "../../sessions/send-policy.js"; @@ -4237,9 +4241,12 @@ export const chatHandlers: GatewayRequestHandlers = { }); return; } + let releaseGatewayRootContinuation: (() => void) | undefined; const cleanupAdmittedRun: typeof activeRunAbort.cleanup = (options) => { activeRunAbort.cleanup(options); gatewayWorkAdmission?.release(); + releaseGatewayRootContinuation?.(); + releaseGatewayRootContinuation = undefined; }; claimAgentRunContext(clientRunId, { sessionKey, @@ -4407,7 +4414,7 @@ export const chatHandlers: GatewayRequestHandlers = { const chatSendAckedAtMs = chatSendTiming?.ackedAtMs ?? performance.now(); const titleSource = stripInlineDirectiveTagsForDisplay(rawMessage).text; if (isDashboardSessionTitleCandidate({ sessionKey, userMessage: titleSource })) { - void (async () => { + void runWithGatewayIndependentRootWorkContinuation(async () => { const titleEntry = entry?.sessionId === admittedSessionId ? entry @@ -4432,7 +4439,7 @@ export const chatHandlers: GatewayRequestHandlers = { reason: "chat.title", }); } - })().catch((err: unknown) => { + }).catch((err: unknown) => { context.logGateway.warn( `dashboard session title generation failed: ${formatForLog(err)}`, ); @@ -4790,6 +4797,9 @@ export const chatHandlers: GatewayRequestHandlers = { } emitServerTiming("first-assistant-event", undefined, dispatchStartedAtMs); }; + // Reserve the detached dispatch before this request releases its root. Otherwise + // its inherited ALS context becomes retired and rejects queued/session work. + releaseGatewayRootContinuation = retainGatewayRootWorkAdmissionContinuation() ?? undefined; void gatewayWorkAdmission .run(() => measureDiagnosticsTimelineSpan( @@ -5839,17 +5849,19 @@ export const chatHandlers: GatewayRequestHandlers = { }); }) .finally(() => { + const dispatchError = pendingDispatchLifecycleError; + // Reserve error projection before cleanup retires the dispatch root. Restart + // drain may already reject fresh roots, but this accepted request must finish. + const releaseDispatchErrorRoot = dispatchError + ? retainGatewayRootWorkAdmissionContinuation() + : null; cleanupAdmittedRun(); clearAgentRunContext(clientRunId, lifecycleGeneration); context.removeChatRun(clientRunId, clientRunId, sessionKey); - if (!pendingDispatchLifecycleError) { + if (!dispatchError) { return; } const persistDispatchLifecycleError = async () => { - const dispatchError = pendingDispatchLifecycleError; - if (!dispatchError) { - return; - } const hasActiveRun = hasTrackedActiveSessionRun({ context, requestedKey: rawSessionKey, @@ -5888,7 +5900,13 @@ export const chatHandlers: GatewayRequestHandlers = { ); } }; - void persistDispatchLifecycleError(); + void persistDispatchLifecycleError() + .catch((continuationErr: unknown) => { + context.logGateway.warn( + `webchat session lifecycle continuation failed: ${formatForLog(continuationErr)}`, + ); + }) + .finally(() => releaseDispatchErrorRoot?.()); }); } catch (err) { cleanupAdmittedRun({ force: true }); diff --git a/src/gateway/server-methods/nodes.helpers.ts b/src/gateway/server-methods/nodes.helpers.ts index cbba400abe01..115d5c9fe0bc 100644 --- a/src/gateway/server-methods/nodes.helpers.ts +++ b/src/gateway/server-methods/nodes.helpers.ts @@ -4,9 +4,11 @@ import { normalizeOptionalString } from "@openclaw/normalization-core/string-coe import { ErrorCodes, errorShape, +} from "../../../packages/gateway-protocol/src/schema/error-codes.js"; +import { formatValidationErrors, -} from "../../../packages/gateway-protocol/src/index.js"; -import type { ValidationError } from "../../../packages/gateway-protocol/src/index.js"; + type ValidationError, +} from "../../../packages/gateway-protocol/src/validation-errors.js"; export { safeParseJson } from "../server-json.js"; import { formatForLog } from "../ws-log.js"; import type { RespondFn } from "./types.js"; diff --git a/src/gateway/server-methods/shared-types.ts b/src/gateway/server-methods/shared-types.ts index 164a9d8e9bfc..2e618cec7b36 100644 --- a/src/gateway/server-methods/shared-types.ts +++ b/src/gateway/server-methods/shared-types.ts @@ -4,12 +4,11 @@ import type { ConnectParams, ErrorShape, RequestFrame, -} from "../../../packages/gateway-protocol/src/index.js"; +} from "../../../packages/gateway-protocol/src/schema/frames.js"; import type { ModelCatalogEntry } from "../../agents/model-catalog.types.js"; import type { CliDeps } from "../../cli/deps.types.js"; import type { HealthSummary } from "../../commands/health.types.js"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; -import type { CronServiceContract } from "../../cron/service-contract.js"; import type { PluginApprovalRequestPayload } from "../../infra/plugin-approvals.js"; import type { createSubsystemLogger } from "../../logging/subsystem.js"; import type { WizardSession } from "../../wizard/session.js"; @@ -28,6 +27,7 @@ import type { ChatRunEntry, ChatRunRegistration, } from "../server-chat-state.js"; +import type { GatewayCronServiceContract } from "../server-cron.js"; import type { DedupeEntry } from "../server-shared.js"; import type { GatewayEventLoopHealth } from "../server/event-loop-health.js"; import type { TerminalLaunchResolution } from "../terminal/launch.js"; @@ -68,7 +68,7 @@ export type RespondFn = ( /** Runtime services and mutable gateway state available to request handlers. */ export type GatewayRequestContext = { deps: CliDeps; - cron: CronServiceContract; + cron: GatewayCronServiceContract; cronStorePath: string; getRuntimeConfig: () => OpenClawConfig; resolveTerminalLaunchPolicy: (agentId?: string) => TerminalLaunchResolution; diff --git a/src/gateway/server-methods/suspend.test.ts b/src/gateway/server-methods/suspend.test.ts new file mode 100644 index 000000000000..16b22a0448bf --- /dev/null +++ b/src/gateway/server-methods/suspend.test.ts @@ -0,0 +1,176 @@ +// Covers suspension RPC validation and coordinator response mapping. +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { suspendHandlers } from "./suspend.js"; + +const coordinator = vi.hoisted(() => ({ + prepare: vi.fn(), + status: vi.fn(), + resume: vi.fn(), +})); + +vi.mock("../../infra/gateway-suspend-coordinator.js", () => ({ + prepareGatewaySuspend: coordinator.prepare, + getGatewaySuspendStatus: coordinator.status, + resumeGatewaySuspend: coordinator.resume, +})); + +vi.mock("../server-active-work.js", () => ({ + createGatewayServerActiveWorkInspectors: vi.fn(() => ({ getChatRuns: vi.fn(() => 0) })), +})); + +function invoke(method: keyof typeof suspendHandlers, params: unknown) { + const respond = vi.fn(); + const pauseScheduling = vi.fn(); + const resumeScheduling = vi.fn(); + const warn = vi.fn(); + const handler = suspendHandlers[method]; + return Promise.resolve( + handler({ + params, + respond, + context: { + cron: { pauseScheduling, resumeScheduling }, + logGateway: { warn }, + chatAbortControllers: new Map(), + chatQueuedTurns: new Map(), + }, + } as unknown as Parameters[0]), + ).then(() => ({ respond, pauseScheduling, resumeScheduling })); +} + +beforeEach(() => { + vi.clearAllMocks(); +}); + +describe("gateway suspend handlers", () => { + it("validates the closed prepare params shape", async () => { + const { respond } = await invoke("gateway.suspend.prepare", { + requestId: "request-1", + extra: true, + }); + + expect(coordinator.prepare).not.toHaveBeenCalled(); + expect(respond).toHaveBeenCalledWith(false, undefined, { + code: "INVALID_REQUEST", + message: "invalid gateway.suspend.prepare params", + }); + }); + + it("wires prepare to scheduler pause/resume and returns busy or ready", async () => { + coordinator.prepare.mockReturnValueOnce({ + status: "busy", + reason: "active-work", + activeCount: 1, + blockers: [{ kind: "queue", count: 1, message: "busy" }], + }); + const { respond, pauseScheduling, resumeScheduling } = await invoke("gateway.suspend.prepare", { + requestId: "request-1", + }); + + expect(coordinator.prepare).toHaveBeenCalledWith( + expect.objectContaining({ + requestId: "request-1", + pauseScheduling: expect.any(Function), + resumeScheduling: expect.any(Function), + }), + ); + const options = coordinator.prepare.mock.calls[0]?.[0]; + options.pauseScheduling(); + options.resumeScheduling(); + expect(pauseScheduling).toHaveBeenCalledOnce(); + expect(resumeScheduling).toHaveBeenCalledOnce(); + expect(respond).toHaveBeenCalledWith( + true, + expect.objectContaining({ status: "busy", reason: "active-work" }), + ); + }); + + it("maps a competing prepared lease to retryable unavailable", async () => { + coordinator.prepare.mockReturnValueOnce({ status: "conflict", expiresAtMs: Date.now() + 5000 }); + const { respond } = await invoke("gateway.suspend.prepare", { requestId: "request-2" }); + + expect(respond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ + code: "UNAVAILABLE", + details: expect.objectContaining({ reason: "gateway-suspension-conflict" }), + retryable: true, + }), + ); + }); + + it("maps prepare and status recovery to the same retryable unavailable error", async () => { + const recovering = { + status: "recovering", + reason: "scheduler-resume-failed", + retryAfterMs: 1_000, + }; + coordinator.prepare.mockReturnValueOnce(recovering); + coordinator.status.mockReturnValueOnce(recovering); + + const prepared = await invoke("gateway.suspend.prepare", { requestId: "request-recovery" }); + const status = await invoke("gateway.suspend.status", { suspensionId: "stale-id" }); + + for (const respond of [prepared.respond, status.respond]) { + expect(respond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ + code: "UNAVAILABLE", + message: "gateway scheduler recovery is pending", + retryable: true, + retryAfterMs: 1_000, + details: { reason: "scheduler-resume-failed" }, + }), + ); + } + }); + + it("keeps resume idempotent and rejects a mismatched active lease", async () => { + coordinator.resume.mockReturnValueOnce({ ok: false, reason: "suspension-mismatch" }); + const mismatch = await invoke("gateway.suspend.resume", { + suspensionId: "suspension-wrong", + }); + expect(mismatch.respond).toHaveBeenCalledWith(false, undefined, { + code: "INVALID_REQUEST", + message: "gateway suspension id does not match", + }); + + coordinator.resume.mockReturnValueOnce({ + ok: true, + status: "running", + resumed: false, + }); + const resumed = await invoke("gateway.suspend.resume", { suspensionId: "suspension-1" }); + expect(resumed.respond).toHaveBeenCalledWith(true, { + ok: true, + status: "running", + resumed: false, + }); + }); + + it("returns retryable unavailable when scheduler resume needs retry", async () => { + coordinator.resume.mockReturnValueOnce({ + ok: false, + reason: "scheduler-resume-failed", + retryAfterMs: 1_000, + }); + + const { respond } = await invoke("gateway.suspend.resume", { + suspensionId: "suspension-1", + }); + + expect(respond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ + code: "UNAVAILABLE", + message: "gateway scheduler recovery is pending", + retryable: true, + retryAfterMs: 1_000, + details: { reason: "scheduler-resume-failed" }, + }), + ); + }); +}); diff --git a/src/gateway/server-methods/suspend.ts b/src/gateway/server-methods/suspend.ts new file mode 100644 index 000000000000..1ae15996457e --- /dev/null +++ b/src/gateway/server-methods/suspend.ts @@ -0,0 +1,107 @@ +// Gateway RPC handlers for cooperative, host-neutral process suspension. +import { + ErrorCodes, + errorShape, + validateGatewaySuspendPrepareParams, + validateGatewaySuspendResumeParams, + validateGatewaySuspendStatusParams, +} from "../../../packages/gateway-protocol/src/index.js"; +import { + getGatewaySuspendStatus, + prepareGatewaySuspend, + resumeGatewaySuspend, +} from "../../infra/gateway-suspend-coordinator.js"; +import { createGatewayServerActiveWorkInspectors } from "../server-active-work.js"; +import type { GatewayRequestHandlers } from "./types.js"; + +function invalidParams(method: string) { + return errorShape(ErrorCodes.INVALID_REQUEST, `invalid ${method} params`); +} + +function schedulerRecoveryError(retryAfterMs: number) { + return errorShape(ErrorCodes.UNAVAILABLE, "gateway scheduler recovery is pending", { + retryable: true, + retryAfterMs, + details: { reason: "scheduler-resume-failed" }, + }); +} + +export const suspendHandlers: GatewayRequestHandlers = { + "gateway.suspend.prepare": async ({ respond, params, context }) => { + if (!validateGatewaySuspendPrepareParams(params)) { + respond(false, undefined, invalidParams("gateway.suspend.prepare")); + return; + } + const requestId = params.requestId.trim(); + const result = prepareGatewaySuspend({ + requestId, + pauseScheduling: () => context.cron.pauseScheduling(), + resumeScheduling: () => context.cron.resumeScheduling(), + inspect: createGatewayServerActiveWorkInspectors(context), + warn: (message) => context.logGateway.warn(message), + }); + if (result.status === "conflict") { + respond( + false, + undefined, + errorShape(ErrorCodes.UNAVAILABLE, "another gateway suspension is already prepared", { + retryable: true, + retryAfterMs: Math.max(0, result.expiresAtMs - Date.now()), + details: { reason: "gateway-suspension-conflict", expiresAtMs: result.expiresAtMs }, + }), + ); + return; + } + if (result.status === "recovering") { + respond(false, undefined, schedulerRecoveryError(result.retryAfterMs)); + return; + } + respond(true, result); + }, + "gateway.suspend.status": async ({ respond, params }) => { + if (!validateGatewaySuspendStatusParams(params)) { + respond(false, undefined, invalidParams("gateway.suspend.status")); + return; + } + const suspensionId = params.suspensionId.trim(); + const result = getGatewaySuspendStatus(suspensionId); + if (result.status === "conflict") { + respond( + false, + undefined, + errorShape(ErrorCodes.UNAVAILABLE, "a different gateway suspension is prepared", { + retryable: true, + retryAfterMs: Math.max(0, result.expiresAtMs - Date.now()), + details: { reason: "gateway-suspension-conflict", expiresAtMs: result.expiresAtMs }, + }), + ); + return; + } + if (result.status === "recovering") { + respond(false, undefined, schedulerRecoveryError(result.retryAfterMs)); + return; + } + respond(true, result); + }, + "gateway.suspend.resume": async ({ respond, params }) => { + if (!validateGatewaySuspendResumeParams(params)) { + respond(false, undefined, invalidParams("gateway.suspend.resume")); + return; + } + const suspensionId = params.suspensionId.trim(); + const result = resumeGatewaySuspend(suspensionId); + if (!result.ok) { + if (result.reason === "scheduler-resume-failed") { + respond(false, undefined, schedulerRecoveryError(result.retryAfterMs)); + return; + } + respond( + false, + undefined, + errorShape(ErrorCodes.INVALID_REQUEST, "gateway suspension id does not match"), + ); + return; + } + respond(true, result); + }, +}; diff --git a/src/gateway/server-plugins.ts b/src/gateway/server-plugins.ts index 656f9dfff29d..5da017521132 100644 --- a/src/gateway/server-plugins.ts +++ b/src/gateway/server-plugins.ts @@ -9,7 +9,7 @@ import { GATEWAY_CLIENT_IDS, GATEWAY_CLIENT_MODES, } from "../../packages/gateway-protocol/src/client-info.js"; -import type { ErrorShape } from "../../packages/gateway-protocol/src/index.js"; +import type { ErrorShape } from "../../packages/gateway-protocol/src/schema/frames.js"; import { PROTOCOL_VERSION } from "../../packages/gateway-protocol/src/version.js"; import { normalizeModelRef, parseModelRef } from "../agents/model-selection.js"; import { applyPluginAutoEnable } from "../config/plugin-auto-enable.js"; diff --git a/src/gateway/server-reload-handlers.test.ts b/src/gateway/server-reload-handlers.test.ts index 4ca81900ca8b..0b18bd6754d2 100644 --- a/src/gateway/server-reload-handlers.test.ts +++ b/src/gateway/server-reload-handlers.test.ts @@ -4,11 +4,20 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { ConfigWriteNotification } from "../config/config.js"; import type { OpenClawConfig } from "../config/types.openclaw.js"; -import { consumeGatewaySigusr1RestartIntent } from "../infra/restart.js"; +import { + consumeGatewaySigusr1RestartIntent, + markGatewaySigusr1RestartHandled, + testing as restartTesting, +} from "../infra/restart.js"; import { pinActivePluginChannelRegistry, releasePinnedPluginChannelRegistry, } from "../plugins/runtime.js"; +import { + isGatewayWorkAdmissionClosed, + resetGatewayWorkAdmission, + tryBeginGatewayRootWorkAdmission, +} from "../process/gateway-work-admission.js"; import { createEmptyRuntimeWebToolsMetadata } from "../secrets/runtime-fast-path.js"; import { activateSecretsRuntimeSnapshot, clearSecretsRuntimeSnapshot } from "../secrets/runtime.js"; import { createChannelTestPluginBase, createTestRegistry } from "../test-utils/channel-plugins.js"; @@ -449,6 +458,48 @@ describe("gateway hot reload model state", () => { }); describe("gateway restart deferral preflight", () => { + it("holds root admission across an immediate config-reload restart signal", () => { + restartTesting.resetSigusr1State(); + resetGatewayWorkAdmission(); + const signalSpy = vi.fn(); + process.once("SIGUSR1", signalSpy); + const { requestGatewayRestart } = createReloadHandlersForTest(); + + try { + expect( + requestGatewayRestart( + { + changedPaths: ["gateway.port"], + restartGateway: true, + restartReasons: ["gateway.port"], + hotReasons: [], + reloadHooks: false, + restartGmailWatcher: false, + restartCron: false, + restartHeartbeat: false, + restartHealthMonitor: false, + reloadPlugins: false, + restartChannels: new Set(), + disposeMcpRuntimes: false, + noopPaths: [], + }, + {}, + ), + ).toBe(true); + + expect(signalSpy).toHaveBeenCalledOnce(); + expect(isGatewayWorkAdmissionClosed()).toBe(true); + expect(tryBeginGatewayRootWorkAdmission()).toBeNull(); + + markGatewaySigusr1RestartHandled(); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + } finally { + process.removeListener("SIGUSR1", signalSpy); + restartTesting.resetSigusr1State(); + resetGatewayWorkAdmission(); + } + }); + it("defers channel hot reload until active embedded work drains", async () => { const previousSkipChannels = process.env.OPENCLAW_SKIP_CHANNELS; const previousSkipProviders = process.env.OPENCLAW_SKIP_PROVIDERS; @@ -815,7 +866,6 @@ describe("gateway restart deferral preflight", () => { }); it("logs active task run ids before waiting and when forcing after timeout", async () => { - const restartTesting = (await import("../infra/restart.js")).testing; restartTesting.resetSigusr1State(); const logReload = { info: vi.fn(), warn: vi.fn() }; const { requestGatewayRestart } = createReloadHandlersForTest(logReload); @@ -902,7 +952,6 @@ describe("gateway restart deferral preflight", () => { }); it("uses the default restart deferral timeout when config omits deferralTimeoutMs", async () => { - const restartTesting = (await import("../infra/restart.js")).testing; restartTesting.resetSigusr1State(); const { requestGatewayRestart } = createReloadHandlersForTest(); hoisted.activeTaskCount.value = 1; diff --git a/src/gateway/server-reload-handlers.ts b/src/gateway/server-reload-handlers.ts index a5456f16284f..5e0e142cba09 100644 --- a/src/gateway/server-reload-handlers.ts +++ b/src/gateway/server-reload-handlers.ts @@ -23,11 +23,12 @@ import type { HeartbeatRunner } from "../infra/heartbeat-runner.js"; import { resetDirectoryCache } from "../infra/outbound/target-resolver.js"; import { deferGatewayRestartUntilIdle, - emitGatewayRestart, + emitGatewayRestartWithSignalAdmission, resolveGatewayRestartDeferralTimeoutMs, setGatewaySigusr1RestartPolicy, } from "../infra/restart.js"; import { getTotalQueueSize } from "../process/command-queue.js"; +import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js"; import { clearSecretsRuntimeSnapshot, getActiveSecretsRuntimeSnapshot, @@ -713,7 +714,10 @@ export function createGatewayReloadHandlers(params: GatewayReloadHandlerParams) } // No active operations or pending replies, restart immediately params.logReload.warn(`config change requires gateway restart (${reasons})`); - const emitted = emitGatewayRestart(); + // The managed reloader owns independent root admission until onRestart + // returns. Extend that fence across signal delivery until the run loop + // atomically promotes it to one-way restart drain. + const emitted = emitGatewayRestartWithSignalAdmission(); if (!emitted) { params.logReload.info("gateway restart already scheduled; skipping duplicate signal"); } @@ -778,6 +782,7 @@ export function startManagedGatewayConfigReloader( initialConfig: params.initialConfig, initialCompareConfig: params.initialCompareConfig, initialInternalWriteHash: params.initialInternalWriteHash, + runTransaction: runWithGatewayIndependentRootWorkAdmission, readSnapshot: params.readSnapshot, promoteSnapshot: async (snapshot, _reason) => await params.promoteSnapshot(snapshot), subscribeToWrites: params.subscribeToWrites, diff --git a/src/gateway/server-restart-sentinel.test.ts b/src/gateway/server-restart-sentinel.test.ts index 9d569be0f017..ca158678fb46 100644 --- a/src/gateway/server-restart-sentinel.test.ts +++ b/src/gateway/server-restart-sentinel.test.ts @@ -158,6 +158,7 @@ const mocks = vi.hoisted(() => { recordInboundSessionAndDispatchReply: vi.fn( async (_params: RecordInboundSessionAndDispatchReplyParams) => {}, ), + logDebug: vi.fn(), logInfo: vi.fn(), logWarn: vi.fn(), logError: vi.fn(), @@ -283,6 +284,7 @@ vi.mock("../infra/heartbeat-wake.js", async () => { vi.mock("../logging/subsystem.js", () => { const logger = { + debug: mocks.logDebug, info: mocks.logInfo, warn: mocks.logWarn, error: mocks.logError, @@ -305,6 +307,12 @@ const { refreshLatestUpdateRestartSentinel, scheduleRestartSentinelWake, } = await import("./server-restart-sentinel.js"); +const { + getActiveGatewayRootWorkCount, + getGatewaySuspendAdmissionPhase, + resetGatewayWorkAdmission, + tryBeginGatewaySuspendAdmission, +} = await import("../process/gateway-work-admission.js"); function expectRecordFields( record: unknown, @@ -371,12 +379,15 @@ function expectContinuationDispatchFields( describe("scheduleRestartSentinelWake", () => { afterEach(() => { + resetGatewayWorkAdmission(); vi.useRealTimers(); }); beforeEach(() => { + resetGatewayWorkAdmission(); vi.useRealTimers(); mocks.queuedSessionDelivery = null; + mocks.readRestartSentinel.mockReset(); mocks.readRestartSentinel.mockResolvedValue({ version: 1, payload: { @@ -476,6 +487,47 @@ describe("scheduleRestartSentinelWake", () => { expect(mocks.logWarn).not.toHaveBeenCalled(); }); + it("defers pending update retries until suspension resumes", async () => { + vi.useFakeTimers(); + const pendingPayload: RestartSentinelPayload = { + kind: "update", + status: "skipped", + ts: 123, + stats: { + mode: "git", + handoffId: "handoff-1", + reason: "managed-service-handoff-started", + }, + }; + let finishRetryRead: ((value: RestartSentinel | null) => void) | undefined; + const retryRead = new Promise((resolve) => { + finishRetryRead = resolve; + }); + mocks.readRestartSentinel + .mockResolvedValueOnce({ version: 1, payload: pendingPayload }) + .mockImplementationOnce(async () => (await retryRead) as RestartSentinel); + + await scheduleRestartSentinelWake({ deps: {} as never }); + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + + await vi.advanceTimersByTimeAsync(1); + expect(getGatewaySuspendAdmissionPhase()).toBe("prepared"); + expect(mocks.readRestartSentinel).toHaveBeenCalledTimes(1); + expect(getActiveGatewayRootWorkCount()).toBe(0); + + expect(suspension?.release()).toBe(true); + await vi.waitFor(() => { + expect(mocks.readRestartSentinel).toHaveBeenCalledTimes(2); + }); + expect(getActiveGatewayRootWorkCount()).toBe(1); + + finishRetryRead?.(null); + await vi.waitFor(() => { + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); + }); + it("retries outbound delivery once and logs a warning without dropping the agent wake", async () => { vi.useFakeTimers(); mocks.deliverOutboundPayloads diff --git a/src/gateway/server-restart-sentinel.ts b/src/gateway/server-restart-sentinel.ts index 46e8975b7ed7..ed48c6015483 100644 --- a/src/gateway/server-restart-sentinel.ts +++ b/src/gateway/server-restart-sentinel.ts @@ -41,6 +41,7 @@ import { isPendingControlPlaneUpdateRestartSentinel } from "../infra/update-cont import { createSubsystemLogger } from "../logging/subsystem.js"; import { stringifyRouteThreadId } from "../plugin-sdk/channel-route.js"; import type { OutboundReplyPayload } from "../plugin-sdk/reply-payload.js"; +import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js"; import { deliveryContextFromSession, mergeDeliveryContext, @@ -468,9 +469,11 @@ async function loadRestartSentinelStartupTask(params: { const attempt = params.attempt ?? 0; if (attempt < CONTROL_PLANE_UPDATE_PENDING_MAX_ATTEMPTS) { const timer = setTimeout(() => { - void scheduleRestartSentinelWakeAttempt({ - deps: params.deps, - attempt: attempt + 1, + void runWithGatewayIndependentRootWorkAdmission(async () => { + await scheduleRestartSentinelWakeAttempt({ + deps: params.deps, + attempt: attempt + 1, + }); }).catch((err: unknown) => { log.warn(`restart sentinel pending update retry failed: ${formatErrorMessage(err)}`); }); diff --git a/src/gateway/server-runtime-services.ts b/src/gateway/server-runtime-services.ts index a1cfc9ebea3b..d83f632cc2db 100644 --- a/src/gateway/server-runtime-services.ts +++ b/src/gateway/server-runtime-services.ts @@ -5,6 +5,7 @@ import type { OpenClawConfig } from "../config/types.openclaw.js"; import { isVitestRuntimeEnv } from "../infra/env.js"; import { startHeartbeatRunner, type HeartbeatRunner } from "../infra/heartbeat-runner.js"; import type { PluginMetadataRegistryView } from "../plugins/plugin-metadata-snapshot.types.js"; +import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js"; import { isGatewayModelPricingEnabled } from "./model-pricing-config.js"; import type { startGatewayMaintenanceTimers } from "./server-maintenance.js"; import { @@ -30,10 +31,10 @@ export function startGatewayCronWithLogging(params: { afterStart?: () => Promise; logCron: { error: (message: string) => void }; }): void { - void params.cron - .start() - .then(() => params.afterStart?.()) - .catch((err: unknown) => params.logCron.error(`failed to start: ${String(err)}`)); + void runWithGatewayIndependentRootWorkAdmission(async () => { + await params.cron.start(); + await params.afterStart?.(); + }).catch((err: unknown) => params.logCron.error(`failed to start: ${String(err)}`)); } function clearGatewayMaintenanceHandles(maintenance: GatewayMaintenanceHandles | null): void { @@ -100,38 +101,42 @@ export function scheduleGatewayPostReadyMaintenance(params: { if (params.isClosing()) { return; } - void runGatewayPostReadyMaintenance({ - startMaintenance: async () => { - if (params.isClosing()) { - return null; - } - const maintenance = await params.startMaintenance(); - if (params.isClosing()) { - // Maintenance can allocate intervals before shutdown is observed; clear them here - // instead of handing live timers to a closing gateway. - clearGatewayMaintenanceHandles(maintenance); - return null; - } - return maintenance; - }, - applyMaintenance: (maintenance) => { - if (params.isClosing()) { - clearGatewayMaintenanceHandles(maintenance); - return; - } - params.applyMaintenance(maintenance); - }, - shouldStartCron: () => !params.isClosing() && params.shouldStartCron(), - markCronStartHandled: params.markCronStartHandled, - cron: params.cron, - logCron: params.logCron, - log: params.log, - recordPostReadyMemory: () => { - if (!params.isClosing()) { - params.recordPostReadyMemory(); - } - }, - }); + void runWithGatewayIndependentRootWorkAdmission(async () => + runGatewayPostReadyMaintenance({ + startMaintenance: async () => { + if (params.isClosing()) { + return null; + } + const maintenance = await params.startMaintenance(); + if (params.isClosing()) { + // Maintenance can allocate intervals before shutdown is observed; clear them here + // instead of handing live timers to a closing gateway. + clearGatewayMaintenanceHandles(maintenance); + return null; + } + return maintenance; + }, + applyMaintenance: (maintenance) => { + if (params.isClosing()) { + clearGatewayMaintenanceHandles(maintenance); + return; + } + params.applyMaintenance(maintenance); + }, + shouldStartCron: () => !params.isClosing() && params.shouldStartCron(), + markCronStartHandled: params.markCronStartHandled, + cron: params.cron, + logCron: params.logCron, + log: params.log, + recordPostReadyMemory: () => { + if (!params.isClosing()) { + params.recordPostReadyMemory(); + } + }, + }), + ).catch((err: unknown) => + params.log.warn(`gateway post-ready maintenance deferred task failed: ${String(err)}`), + ); }, params.delayMs); timer.unref?.(); return timer; @@ -143,7 +148,7 @@ function recoverPendingOutboundDeliveries(params: { }): void { // Recovery is best-effort background work; startup must continue even if outbound modules fail // to import or queued delivery replay fails. - void (async () => { + void runWithGatewayIndependentRootWorkAdmission(async () => { const { recoverPendingDeliveries } = await import("../infra/outbound/delivery-queue.js"); const { deliverOutboundPayloadsInternal } = await import("../infra/outbound/deliver.js"); const logRecovery = params.log.child("delivery-recovery"); @@ -152,7 +157,7 @@ function recoverPendingOutboundDeliveries(params: { log: logRecovery, cfg: params.cfg, }); - })().catch((err: unknown) => params.log.error(`Delivery recovery failed: ${String(err)}`)); + }).catch((err: unknown) => params.log.error(`Delivery recovery failed: ${String(err)}`)); } function recoverPendingSessionDeliveries(params: { @@ -163,7 +168,7 @@ function recoverPendingSessionDeliveries(params: { // Delay session continuation recovery so the gateway has time to publish ready state and // request routing before replaying restart-sentinel deliveries. const timer = setTimeout(() => { - void (async () => { + void runWithGatewayIndependentRootWorkAdmission(async () => { const { recoverPendingRestartContinuationDeliveries } = await import("./server-restart-sentinel.js"); const logRecovery = params.log.child("session-delivery-recovery"); @@ -172,7 +177,7 @@ function recoverPendingSessionDeliveries(params: { log: logRecovery, maxEnqueuedAt: params.maxEnqueuedAt, }); - })().catch((err: unknown) => + }).catch((err: unknown) => params.log.error(`Session delivery recovery failed: ${String(err)}`), ); }, 1_250); @@ -191,7 +196,7 @@ function startGatewayModelPricingRefreshOnDemand(params: { let stopRefresh: (() => void) | undefined; // Import pricing refresh lazily; many gateway starts never use model-pricing metadata. // The stopped flag closes the race where shutdown happens before the import resolves. - void (async () => { + void runWithGatewayIndependentRootWorkAdmission(async () => { const { startGatewayModelPricingRefresh } = await import("./model-pricing-cache.js"); if (stopped) { return; @@ -204,7 +209,7 @@ function startGatewayModelPricingRefreshOnDemand(params: { stopRefresh(); stopRefresh = undefined; } - })().catch((err: unknown) => + }).catch((err: unknown) => params.log.error(`Model pricing refresh failed to start: ${String(err)}`), ); return () => { diff --git a/src/gateway/server-runtime-state.ts b/src/gateway/server-runtime-state.ts index 4fe771cd6ff1..01dc029363f2 100644 --- a/src/gateway/server-runtime-state.ts +++ b/src/gateway/server-runtime-state.ts @@ -30,7 +30,11 @@ import { createToolEventRecipientRegistry, } from "./server-chat-state.js"; import { MAX_PREAUTH_PAYLOAD_BYTES } from "./server-constants.js"; -import { attachGatewayUpgradeHandler, createGatewayHttpServer } from "./server-http.js"; +import { + attachGatewayUpgradeHandler, + createGatewayHttpServer, + runWithGatewayHttpWorkAdmission, +} from "./server-http.js"; import type { GatewayRequestContext } from "./server-methods/types.js"; import type { DedupeEntry } from "./server-shared.js"; import type { HookClientIpConfig, HooksRequestHandler } from "./server/hooks-request-handler.js"; @@ -153,20 +157,22 @@ export async function createGatewayRuntimeState(params: { if (url.pathname !== basePath && !url.pathname.startsWith(`${basePath}/`)) { return false; } - if (!loadedHooksRequestHandler) { - // Hooks are cold for most gateway starts; create the handler only after a request - // matches the configured base path so startup avoids importing hook runtime code. - const { createGatewayHooksRequestHandler } = await import("./server/hooks.js"); - loadedHooksRequestHandler = createGatewayHooksRequestHandler({ - deps: params.deps, - getHooksConfig: params.hooksConfig, - getClientIpConfig: params.getHookClientIpConfig, - bindHost: params.bindHost, - port: params.port, - logHooks: params.logHooks, - }); - } - return await loadedHooksRequestHandler(req, res); + return await runWithGatewayHttpWorkAdmission(res, async () => { + if (!loadedHooksRequestHandler) { + // Hooks are cold for most gateway starts; create the handler only after a request + // matches the configured base path so startup avoids importing hook runtime code. + const { createGatewayHooksRequestHandler } = await import("./server/hooks.js"); + loadedHooksRequestHandler = createGatewayHooksRequestHandler({ + deps: params.deps, + getHooksConfig: params.hooksConfig, + getClientIpConfig: params.getHookClientIpConfig, + bindHost: params.bindHost, + port: params.port, + logHooks: params.logHooks, + }); + } + return await loadedHooksRequestHandler(req, res); + }); }; let loadedPluginRequestHandler: GatewayPluginRequestHandler | null = null; diff --git a/src/gateway/server-shared.ts b/src/gateway/server-shared.ts index ea915108c99f..434d68fb7fdf 100644 --- a/src/gateway/server-shared.ts +++ b/src/gateway/server-shared.ts @@ -1,6 +1,6 @@ // Gateway shared request-state types. // Defines cached dedupe entries for idempotent Gateway method calls. -import type { ErrorShape } from "../../packages/gateway-protocol/src/index.js"; +import type { ErrorShape } from "../../packages/gateway-protocol/src/schema/frames.js"; export const PENDING_CHAT_SEND_DEDUPE_PREFIX = "pending-chat:"; diff --git a/src/gateway/server-startup-post-attach.test.ts b/src/gateway/server-startup-post-attach.test.ts index f712ac645159..8856d24f9bef 100644 --- a/src/gateway/server-startup-post-attach.test.ts +++ b/src/gateway/server-startup-post-attach.test.ts @@ -11,6 +11,10 @@ import type { PluginHookGatewayStartEvent, } from "../plugins/hook-types.js"; import type { PluginServicesHandle } from "../plugins/services.js"; +import { + getActiveGatewayRootWorkCount, + resetGatewayWorkAdmission, +} from "../process/gateway-work-admission.js"; import { closeOpenClawStateDatabaseForTest } from "../state/openclaw-state-db.js"; import { withEnvAsync } from "../test-utils/env.js"; @@ -40,7 +44,8 @@ const hoisted = vi.hoisted(() => { skipped: 0, })); const scheduleRestartAbortedMainSessionRecovery = vi.fn(); - const scheduleRestartSentinelWake = vi.fn(); + const scheduleRestartSentinelWake = + vi.fn(); const refreshLatestUpdateRestartSentinel = vi.fn< typeof import("./server-restart-sentinel.js").refreshLatestUpdateRestartSentinel >(async () => null); @@ -301,6 +306,7 @@ function firstGatewayStartCall( describe("startGatewayPostAttachRuntime", () => { beforeEach(() => { + resetGatewayWorkAdmission(); closeOpenClawStateDatabaseForTest(); vi.stubEnv("OPENCLAW_SKIP_CHANNELS", "0"); vi.stubEnv("OPENCLAW_SKIP_PROVIDERS", "0"); @@ -364,6 +370,7 @@ describe("startGatewayPostAttachRuntime", () => { }); afterEach(() => { + resetGatewayWorkAdmission(); closeOpenClawStateDatabaseForTest(); vi.useRealTimers(); vi.unstubAllEnvs(); @@ -371,6 +378,10 @@ describe("startGatewayPostAttachRuntime", () => { it("re-enables startup-gated methods after post-attach sidecars start", async () => { const unavailableGatewayMethods = new Set(["chat.history", "models.list"]); + const methodsAtRecoveryRegistration: string[][] = []; + hoisted.scheduleRestartAbortedMainSessionRecovery.mockImplementationOnce(() => { + methodsAtRecoveryRegistration.push([...unavailableGatewayMethods]); + }); const onSidecarsReady = vi.fn(); const log = { info: vi.fn(), warn: vi.fn() }; @@ -399,6 +410,7 @@ describe("startGatewayPostAttachRuntime", () => { expect(hoisted.scheduleRestartAbortedMainSessionRecovery).toHaveBeenCalledWith({ cfg: { hooks: { internal: { enabled: false } } }, }); + expect(methodsAtRecoveryRegistration).toStrictEqual([["chat.history", "models.list"]]); expect(hoisted.startGatewayMemoryBackend).not.toHaveBeenCalled(); }); @@ -430,6 +442,29 @@ describe("startGatewayPostAttachRuntime", () => { expect(events).toEqual(["sidecars", "returned", "sentinel"]); }); + it("keeps delayed restart sentinel recovery admitted until wake work completes", async () => { + vi.useFakeTimers(); + let finishWake: (() => void) | undefined; + const wake = new Promise((resolve) => { + finishWake = resolve; + }); + hoisted.scheduleRestartSentinelWake.mockReturnValueOnce(wake); + + testing.scheduleRestartSentinelWakeAfterReady({ + deps: {} as never, + log: { warn: vi.fn() }, + }); + await vi.advanceTimersByTimeAsync(750); + + expect(hoisted.scheduleRestartSentinelWake).toHaveBeenCalledOnce(); + expect(getActiveGatewayRootWorkCount()).toBe(1); + + finishWake?.(); + await vi.waitFor(() => { + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); + }); + it("starts sidecars while startup logging is still pending", async () => { const events: string[] = []; let finishStartupLog: (() => void) | undefined; @@ -2050,6 +2085,13 @@ describe("startGatewayPostAttachRuntime", () => { it("dispatches registered gateway startup internal hooks without configured hook packs", async () => { vi.useFakeTimers(); hoisted.hasInternalHookListeners.mockReturnValue(true); + let releaseHook = () => {}; + hoisted.triggerInternalHook.mockImplementationOnce( + () => + new Promise((resolve) => { + releaseHook = resolve; + }), + ); const cfg = {} as never; const deps = {} as never; @@ -2088,7 +2130,11 @@ describe("startGatewayPostAttachRuntime", () => { }, ); expect(hoisted.triggerInternalHook).toHaveBeenCalledWith(hoisted.startupHookEvent); + expect(getActiveGatewayRootWorkCount()).toBe(1); + releaseHook(); + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); } finally { + releaseHook(); vi.useRealTimers(); } }); diff --git a/src/gateway/server-startup-post-attach.ts b/src/gateway/server-startup-post-attach.ts index c4dec73015fe..7ea2d74435d5 100644 --- a/src/gateway/server-startup-post-attach.ts +++ b/src/gateway/server-startup-post-attach.ts @@ -16,6 +16,7 @@ import type { loadOpenClawPlugins } from "../plugins/loader.js"; import { getPluginModuleLoaderStats } from "../plugins/plugin-module-loader-cache.js"; import type { PluginRegistry } from "../plugins/registry.js"; import type { PluginServicesHandle } from "../plugins/services.js"; +import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js"; import { createLazyRuntimeModule } from "../shared/lazy-runtime.js"; import { GATEWAY_EVENT_UPDATE_AVAILABLE, @@ -157,13 +158,12 @@ function scheduleGatewayMemoryBackend(params: { return; } const start = () => { - void import("./server-startup-memory.js") - .then(({ startGatewayMemoryBackend }) => - startGatewayMemoryBackend({ cfg: params.cfg, log: params.log }), - ) - .catch((err: unknown) => { - params.log.warn(`qmd memory startup initialization failed: ${String(err)}`); - }); + void runWithGatewayIndependentRootWorkAdmission(async () => { + const { startGatewayMemoryBackend } = await import("./server-startup-memory.js"); + await startGatewayMemoryBackend({ cfg: params.cfg, log: params.log }); + }).catch((err: unknown) => { + params.log.warn(`qmd memory startup initialization failed: ${String(err)}`); + }); }; if (params.policy.mode === "immediate") { setImmediate(start); @@ -181,12 +181,10 @@ function schedulePostAttachUpdateSentinelRefresh(params: { >; }): void { const handle = setImmediate(() => { - void measureStartup(params.startupTrace, "post-attach.update-sentinel", async () => { - try { + void runWithGatewayIndependentRootWorkAdmission(async () => { + await measureStartup(params.startupTrace, "post-attach.update-sentinel", async () => { await params.refreshLatestUpdateRestartSentinel(); - } catch (err) { - params.log.warn(`restart sentinel refresh failed: ${String(err)}`); - } + }); }).catch((err: unknown) => { params.log.warn(`restart sentinel refresh failed: ${String(err)}`); }); @@ -210,39 +208,41 @@ function scheduleProviderAuthStatePrewarm(params: { let pendingRewarmReason: string | undefined; const isStopped = () => stopped; const delayMs = params.delayMs ?? PROVIDER_AUTH_PREWARM_START_DELAY_MS; - void (async () => { + void runWithGatewayIndependentRootWorkAdmission(async () => { const [{ setAuthProfileFailureHook }, { clearCurrentProviderAuthState }] = await Promise.all([ import("../agents/auth-profiles/failure-hook.js"), import("../agents/model-provider-auth-state.js"), ]); const loadProviderAuthWarmModule = () => import("../agents/model-provider-auth.js"); const runRewarm = async (reason: string) => { - if (isStopped()) { - return; - } - const cfg = params.getConfig(); - rewarmInFlight = true; - try { - const { warmCurrentProviderAuthStateOffMainThread } = await loadProviderAuthWarmModule(); - const metrics = await measureProviderAuthWarm(() => - warmCurrentProviderAuthStateOffMainThread(cfg, { isCancelled: isStopped }), - ); + await runWithGatewayIndependentRootWorkAdmission(async () => { if (isStopped()) { return; } - params.log.info( - `provider auth state re-warmed (${reason}) ${formatProviderAuthWarmMetrics(metrics)}`, - ); - } catch (err) { - params.log.warn(`provider auth state rewarm failed: ${String(err)}`); - } finally { - rewarmInFlight = false; - const nextReason = pendingRewarmReason; - pendingRewarmReason = undefined; - if (nextReason && !isStopped()) { - scheduleAuthMapRewarm(nextReason); + const cfg = params.getConfig(); + rewarmInFlight = true; + try { + const { warmCurrentProviderAuthStateOffMainThread } = await loadProviderAuthWarmModule(); + const metrics = await measureProviderAuthWarm(() => + warmCurrentProviderAuthStateOffMainThread(cfg, { isCancelled: isStopped }), + ); + if (isStopped()) { + return; + } + params.log.info( + `provider auth state re-warmed (${reason}) ${formatProviderAuthWarmMetrics(metrics)}`, + ); + } catch (err) { + params.log.warn(`provider auth state rewarm failed: ${String(err)}`); + } finally { + rewarmInFlight = false; + const nextReason = pendingRewarmReason; + pendingRewarmReason = undefined; + if (nextReason && !isStopped()) { + scheduleAuthMapRewarm(nextReason); + } } - } + }); }; const scheduleAuthMapRewarm = (reason: string) => { // Collapse repeated auth-profile failures into one rewarm turn while a @@ -279,7 +279,7 @@ function scheduleProviderAuthStatePrewarm(params: { } startupTimer = setTimeout( () => { - void (async () => { + void runWithGatewayIndependentRootWorkAdmission(async () => { if (isStopped()) { return; } @@ -294,14 +294,14 @@ function scheduleProviderAuthStatePrewarm(params: { params.log.info( `provider auth state pre-warmed ${formatProviderAuthWarmMetrics(metrics)}`, ); - })().catch((err: unknown) => { + }).catch((err: unknown) => { params.log.warn(`provider auth state pre-warm failed: ${String(err)}`); }); }, Math.max(0, delayMs), ); startupTimer.unref?.(); - })().catch((err: unknown) => { + }).catch((err: unknown) => { params.log.warn(`provider auth state pre-warm setup failed: ${String(err)}`); }); return { @@ -335,26 +335,28 @@ function scheduleAgentRuntimePluginPrewarm(params: { timer = setTimeout( () => { timer = undefined; - void measureStartup(params.startupTrace, "post-ready.agent-runtime-plugins", async () => { - if (isStopped()) { - return; - } - const started = performance.now(); - const { ensureRuntimePluginsLoaded } = await import("../agents/runtime-plugins.js"); - const cfg = params.getConfig(); - if (isStopped()) { - return; - } - ensureRuntimePluginsLoaded({ - config: cfg, - workspaceDir: params.workspaceDir, - allowGatewaySubagentBinding: true, + void runWithGatewayIndependentRootWorkAdmission(async () => { + await measureStartup(params.startupTrace, "post-ready.agent-runtime-plugins", async () => { + if (isStopped()) { + return; + } + const started = performance.now(); + const { ensureRuntimePluginsLoaded } = await import("../agents/runtime-plugins.js"); + const cfg = params.getConfig(); + if (isStopped()) { + return; + } + ensureRuntimePluginsLoaded({ + config: cfg, + workspaceDir: params.workspaceDir, + allowGatewaySubagentBinding: true, + }); + if (!isStopped()) { + params.log.info( + `agent runtime plugins pre-warmed in ${(performance.now() - started).toFixed(0)}ms`, + ); + } }); - if (!isStopped()) { - params.log.info( - `agent runtime plugins pre-warmed in ${(performance.now() - started).toFixed(0)}ms`, - ); - } }).catch((err: unknown) => { params.log.warn(`agent runtime plugin pre-warm failed: ${String(err)}`); }); @@ -387,9 +389,11 @@ function schedulePostReadySidecarTask(params: { if (isStopped()) { return; } - void measureStartup(params.startupTrace, params.name, () => - params.run(isStopped, abortController.signal), - ).catch((err: unknown) => { + void runWithGatewayIndependentRootWorkAdmission(async () => { + await measureStartup(params.startupTrace, params.name, () => + params.run(isStopped, abortController.signal), + ); + }).catch((err: unknown) => { params.log.warn(`${params.name} failed after gateway ready: ${String(err)}`); }); }); @@ -406,6 +410,20 @@ function schedulePostReadySidecarTask(params: { }; } +function scheduleRestartSentinelWakeAfterReady(params: { + deps: CliDeps; + log: { warn: (msg: string) => void }; +}): void { + setTimeout(() => { + void runWithGatewayIndependentRootWorkAdmission(async () => { + const { scheduleRestartSentinelWake } = await loadGatewayRestartSentinelModule(); + await scheduleRestartSentinelWake({ deps: params.deps }); + }).catch((err: unknown) => { + params.log.warn(`restart sentinel wake failed to schedule: ${String(err)}`); + }); + }, 750); +} + type CleanStaleLockFiles = typeof import("../agents/session-write-lock.js").cleanStaleLockFiles; type MarkRestartAbortedMainSessionsFromLocks = typeof import("../agents/main-session-restart-recovery.js").markRestartAbortedMainSessionsFromLocks; @@ -779,19 +797,22 @@ export async function startGatewaySidecars(params: { // Run startup hooks after sidecar startup has yielded once so gateway bind // and channel startup are not delayed by hook handlers. setTimeout(() => { - void loadInternalHooksModule().then(({ createInternalHookEvent, triggerInternalHook }) => { + void runWithGatewayIndependentRootWorkAdmission(async () => { + const { createInternalHookEvent, triggerInternalHook } = await loadInternalHooksModule(); const hookEvent = createInternalHookEvent("gateway", "startup", "gateway:startup", { cfg: params.cfg, deps: params.deps, workspaceDir: params.defaultWorkspaceDir, }); - void triggerInternalHook(hookEvent); + await triggerInternalHook(hookEvent); + }).catch((err: unknown) => { + params.logHooks.warn(`gateway startup hook failed: ${String(err)}`); }); }, 250); } if (params.cfg.acp?.enabled) { - void (async () => { + void runWithGatewayIndependentRootWorkAdmission(async () => { const ready = await measureStartup(params.startupTrace, "sidecars.acp.runtime-ready", () => waitForAcpRuntimeBackendReady({ backendId: params.cfg.acp?.backend }), ); @@ -815,7 +836,7 @@ export async function startGatewaySidecars(params: { `acp startup identity reconcile (renderer=${ACP_SESSION_IDENTITY_RENDERER_VERSION}): checked=${result.checked} resolved=${result.resolved} failed=${result.failed}`, ); }); - })().catch((err: unknown) => { + }).catch((err: unknown) => { params.log.warn(`acp startup identity reconcile failed: ${String(err)}`); }); } @@ -864,15 +885,7 @@ export async function startGatewaySidecars(params: { if (!(await hasRestartSentinelFast())) { return; } - setTimeout(() => { - void loadGatewayRestartSentinelModule() - .then(({ scheduleRestartSentinelWake }) => - scheduleRestartSentinelWake({ deps: params.deps }), - ) - .catch((err: unknown) => { - params.log.warn(`restart sentinel wake failed to schedule: ${String(err)}`); - }); - }, 750); + scheduleRestartSentinelWakeAfterReady({ deps: params.deps, log: params.log }); }, }); @@ -1025,16 +1038,19 @@ function createDeferredGatewayUpdateCheck(params: { if (stopped) { return; } - void measureStartup(params.startupTrace, "post-attach.update-check", () => - params.runtimeDeps.scheduleGatewayUpdateCheck({ - cfg: params.cfg, - log: params.log, - isNixMode: params.isNixMode, - onUpdateAvailableChange: (updateAvailable) => { - const payload: GatewayUpdateAvailableEventPayload = { updateAvailable }; - params.broadcast(GATEWAY_EVENT_UPDATE_AVAILABLE, payload, { dropIfSlow: true }); - }, - }), + void runWithGatewayIndependentRootWorkAdmission( + async () => + await measureStartup(params.startupTrace, "post-attach.update-check", () => + params.runtimeDeps.scheduleGatewayUpdateCheck({ + cfg: params.cfg, + log: params.log, + isNixMode: params.isNixMode, + onUpdateAvailableChange: (updateAvailable) => { + const payload: GatewayUpdateAvailableEventPayload = { updateAvailable }; + params.broadcast(GATEWAY_EVENT_UPDATE_AVAILABLE, payload, { dropIfSlow: true }); + }, + }), + ), ) .then((nextStop) => { if (stopped) { @@ -1259,9 +1275,6 @@ export async function startGatewayPostAttachRuntime( loaderStatsAfter.sourceTransformFallbacks - loaderStatsBefore.sourceTransformFallbacks, ], ]); - for (const method of STARTUP_UNAVAILABLE_GATEWAY_METHODS) { - params.unavailableGatewayMethods.delete(method); - } try { const { scheduleRestartAbortedMainSessionRecovery } = await loadMainSessionRestartRecoveryModule(); @@ -1269,6 +1282,11 @@ export async function startGatewayPostAttachRuntime( } catch (err) { params.log.warn(`main-session restart recovery failed to schedule: ${String(err)}`); } + // Capture the orphan-recovery cutoff before new startup-gated agent + // work can create sessions that the recovery scan must leave alone. + for (const method of STARTUP_UNAVAILABLE_GATEWAY_METHODS) { + params.unavailableGatewayMethods.delete(method); + } if (!pluginServicesReported) { reportPluginServices(result.pluginServices); } @@ -1343,19 +1361,21 @@ export async function startGatewayPostAttachRuntime( const hookRunner = await runtimeDeps.getGlobalHookRunner(); if (hookRunner?.hasHooks("gateway_start")) { const { withPluginHttpRouteRegistry } = await import("../plugins/http-registry.js"); - void withPluginHttpRouteRegistry(sidecarsResult.pluginRegistry, () => - hookRunner.runGatewayStart( - { port: params.port }, - { - port: params.port, - config: params.gatewayPluginConfigAtStart, - workspaceDir: params.defaultWorkspaceDir, - getCron: () => - params.getCronService?.() ?? - (params.deps.cron as PluginHookGatewayCronService | undefined), - }, - ), - ).catch((err: unknown) => { + void runWithGatewayIndependentRootWorkAdmission(async () => { + await withPluginHttpRouteRegistry(sidecarsResult.pluginRegistry, () => + hookRunner.runGatewayStart( + { port: params.port }, + { + port: params.port, + config: params.gatewayPluginConfigAtStart, + workspaceDir: params.defaultWorkspaceDir, + getCron: () => + params.getCronService?.() ?? + (params.deps.cron as PluginHookGatewayCronService | undefined), + }, + ), + ); + }).catch((err: unknown) => { params.log.warn(`gateway_start hook failed: ${String(err)}`); }); } @@ -1398,6 +1418,7 @@ export const testing = { cleanupStaleSessionLocks, scheduleProviderAuthStatePrewarm, schedulePrimaryModelPrewarm, + scheduleRestartSentinelWakeAfterReady, shouldSkipStartupModelPrewarm, stopPostReadySidecarsAfterCloseStarted, }; diff --git a/src/gateway/server.agent.gateway-server-agent-a.test.ts b/src/gateway/server.agent.gateway-server-agent-a.test.ts index dbd864525746..c11763726658 100644 --- a/src/gateway/server.agent.gateway-server-agent-a.test.ts +++ b/src/gateway/server.agent.gateway-server-agent-a.test.ts @@ -4,6 +4,11 @@ import fs from "node:fs/promises"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import type { ChannelPlugin } from "../channels/plugins/types.js"; +import { + getActiveGatewayRootWorkCount, + isGatewaySubordinateWorkAdmissionClosed, + tryBeginGatewaySuspendAdmission, +} from "../process/gateway-work-admission.js"; import { createChannelTestPluginBase, createDirectOutboundTestAdapter, @@ -215,6 +220,42 @@ describe("gateway server agent", () => { testState.allowFrom = undefined; }); + test("keeps accepted detached agent work on its retained request root", async () => { + await setTestSessionStore({ + entries: { + main: { + sessionId: "sess-agent-detached-root", + updatedAt: Date.now(), + }, + }, + }); + let subordinateAdmissionClosed: boolean | undefined; + vi.mocked(agentCommand).mockImplementationOnce(async () => { + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension).not.toBeNull(); + try { + subordinateAdmissionClosed = isGatewaySubordinateWorkAdmissionClosed(); + } finally { + suspension?.rollback(); + } + }); + + const res = await rpcReq(gatewaySuite.ws, "agent", { + message: "prove detached root transfer", + sessionKey: "main", + idempotencyKey: "idem-agent-detached-root", + }); + + expect(res.ok).toBe(true); + expect(res.payload?.status).toBe("accepted"); + await vi.waitFor(() => { + expect(subordinateAdmissionClosed).toBe(false); + }); + await vi.waitFor(() => { + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); + }); + test("agent marks implicit delivery when lastTo is stale", async () => { testState.allowFrom = ["+436769770569"]; await setTestSessionStore({ diff --git a/src/gateway/server.chat.gateway-server-chat.test.ts b/src/gateway/server.chat.gateway-server-chat.test.ts index 2fc833adc0fb..f4a35e09a684 100644 --- a/src/gateway/server.chat.gateway-server-chat.test.ts +++ b/src/gateway/server.chat.gateway-server-chat.test.ts @@ -6,9 +6,18 @@ import path from "node:path"; import { beforeEach, describe, expect, test, vi } from "vitest"; import { WebSocket } from "ws"; import { emitAgentEvent, registerAgentRunContext } from "../infra/agent-events.js"; +import { + getActiveGatewayRootWorkCount, + isGatewaySubordinateWorkAdmissionClosed, + markGatewayRestartDraining, + resetGatewayWorkAdmission, + tryBeginGatewaySuspendAdmission, +} from "../process/gateway-work-admission.js"; import { extractFirstTextBlock } from "../shared/chat-message-content.js"; +import { createDeferred } from "../test-utils/deferred.js"; import { captureEnv, setTestEnvValue } from "../test-utils/env.js"; import { GATEWAY_CLIENT_MODES, GATEWAY_CLIENT_NAMES } from "../utils/message-channel.js"; +import * as sessionLifecycleState from "./session-lifecycle-state.js"; import { connectOk, dispatchInboundMessageMock, @@ -201,6 +210,64 @@ describe("gateway server chat", () => { ); }); + test("keeps started chat dispatch on its retained request root", async () => { + await withMainSessionStore(async () => { + let subordinateAdmissionClosed: boolean | undefined; + dispatchInboundMessageMock.mockImplementationOnce(async (...args: unknown[]) => { + await new Promise((resolve) => setTimeout(resolve, 0)); + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension).not.toBeNull(); + try { + subordinateAdmissionClosed = isGatewaySubordinateWorkAdmissionClosed(); + } finally { + suspension?.rollback(); + } + const [params] = args as [ + { + dispatcher: { + sendFinalReply: (payload: { text: string }) => boolean; + markComplete: () => void; + waitForIdle: () => Promise; + getQueuedCounts: () => { final: number; block: number; tool: number }; + }; + }, + ]; + params.dispatcher.sendFinalReply({ text: "detached root stayed live" }); + params.dispatcher.markComplete(); + await params.dispatcher.waitForIdle(); + return { + queuedFinal: true, + counts: params.dispatcher.getQueuedCounts(), + }; + }); + const finalPromise = onceMessage( + ws, + (message) => + message.type === "event" && + message.event === "chat" && + message.payload?.state === "final" && + message.payload?.runId === "idem-chat-detached-root", + 8_000, + ); + + const res = await rpcReq(ws, "chat.send", { + sessionKey: "main", + message: "prove detached root transfer", + idempotencyKey: "idem-chat-detached-root", + }); + + expect(res.ok).toBe(true); + expect(res.payload?.status).toBe("started"); + await vi.waitFor(() => { + expect(subordinateAdmissionClosed).toBe(false); + }); + await finalPromise; + await vi.waitFor(() => { + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); + }); + }); + const waitForAgentRunOk = async (runId: string, timeoutMs = 1_000) => { const res = await rpcReq(ws, "agent.wait", { runId, @@ -714,7 +781,7 @@ describe("gateway server chat", () => { }); }); - test("marks a running webchat session failed when dispatch rejects before a reply", async () => { + test("marks a running webchat session failed when restart drain overlaps dispatch rejection", async () => { await withMainSessionStore(async (dir) => { await writeSessionStore({ entries: { @@ -729,34 +796,72 @@ describe("gateway server chat", () => { }); const subscribeRes = await rpcReq(ws, "sessions.subscribe", {}); expect(subscribeRes.ok).toBe(true); - dispatchInboundMessageMock.mockRejectedValueOnce(new Error("provider rejected request")); - - const errorPromise = onceMessage( - ws, - (o) => - o.type === "event" && - o.event === "chat" && - o.payload?.state === "error" && - o.payload?.runId === "idem-dispatch-error-1", - 8_000, - ); - const sessionChangedPromise = onceMessage( - ws, - (o) => - o.type === "event" && - o.event === "sessions.changed" && - o.payload?.reason === "chat.dispatch-error" && - o.payload?.sessionKey === "agent:main:main", - 8_000, - ); - const res = await rpcReq(ws, "chat.send", { - sessionKey: "main", - message: "run: pwd", - idempotencyKey: "idem-dispatch-error-1", - }); - expect(res.ok).toBe(true); - await errorPromise; - const sessionChanged = await sessionChangedPromise; + const rejectDispatch = createDeferred(); + const releasePersistence = createDeferred(); + let dispatchStarted = false; + let persistenceEntered = false; + const persistLifecycleEvent = sessionLifecycleState.persistGatewaySessionLifecycleEvent; + const persistSpy = vi + .spyOn(sessionLifecycleState, "persistGatewaySessionLifecycleEvent") + .mockImplementation(async (params) => { + persistenceEntered = true; + await releasePersistence.promise; + await persistLifecycleEvent(params); + }); + const sessionChanged = await (async () => { + try { + dispatchInboundMessageMock.mockImplementationOnce(async () => { + dispatchStarted = true; + await rejectDispatch.promise; + throw new Error("provider rejected request"); + }); + const errorPromise = onceMessage( + ws, + (o) => + o.type === "event" && + o.event === "chat" && + o.payload?.state === "error" && + o.payload?.runId === "idem-dispatch-error-1", + 8_000, + ); + const sessionChangedPromise = onceMessage( + ws, + (o) => + o.type === "event" && + o.event === "sessions.changed" && + o.payload?.reason === "chat.dispatch-error" && + o.payload?.sessionKey === "agent:main:main", + 8_000, + ); + const res = await rpcReq(ws, "chat.send", { + sessionKey: "main", + message: "run: pwd", + idempotencyKey: "idem-dispatch-error-1", + }); + expect(res.ok).toBe(true); + await vi.waitFor(() => { + expect(dispatchStarted).toBe(true); + }); + markGatewayRestartDraining(); + rejectDispatch.resolve(); + await errorPromise; + await vi.waitFor(() => { + expect(persistenceEntered).toBe(true); + }); + expect(getActiveGatewayRootWorkCount()).toBe(1); + releasePersistence.resolve(); + const changed = await sessionChangedPromise; + await vi.waitFor(() => { + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); + return changed; + } finally { + rejectDispatch.resolve(); + releasePersistence.resolve(); + persistSpy.mockRestore(); + resetGatewayWorkAdmission(); + } + })(); expectRecordFields(sessionChanged.payload, { sessionId: "sess-main", status: "failed", diff --git a/src/gateway/server.preauth-hardening.test.ts b/src/gateway/server.preauth-hardening.test.ts index 6ab8624f93e9..ff3ab781a693 100644 --- a/src/gateway/server.preauth-hardening.test.ts +++ b/src/gateway/server.preauth-hardening.test.ts @@ -10,6 +10,7 @@ import { resetDiagnosticEventsForTest, type DiagnosticEventPayload, } from "../infra/diagnostic-events.js"; +import { 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"; @@ -144,6 +145,22 @@ describe("gateway pre-auth hardening", () => { } }); + it("rejects core websocket upgrades while suspension admission is closed", async () => { + const harness = await createGatewaySuiteHarness(); + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + + try { + await expect(requestUpgradeRejection(harness.port)).resolves.toEqual({ + status: 503, + body: "Gateway websocket admission closed", + }); + } finally { + suspension?.release(); + await harness.close(); + } + }); + it("closes idle unauthenticated sockets after the handshake timeout", async () => { setEnvForTest("OPENCLAW_TEST_HANDSHAKE_TIMEOUT_MS", "200"); diff --git a/src/gateway/server.roles-allowlist-update.test.ts b/src/gateway/server.roles-allowlist-update.test.ts index 0a6e85f9c29c..541a09fb4130 100644 --- a/src/gateway/server.roles-allowlist-update.test.ts +++ b/src/gateway/server.roles-allowlist-update.test.ts @@ -445,6 +445,9 @@ describe("gateway update.run", () => { await vi.waitFor(() => { expect(updateMock).toHaveBeenCalledOnce(); }, FAST_WAIT_OPTS); + await vi.waitFor(() => { + expect(sigusr1).toHaveBeenCalled(); + }, FAST_WAIT_OPTS); } finally { process.off("SIGUSR1", sigusr1); } diff --git a/src/gateway/server/plugins-http.runtime-scopes.test.ts b/src/gateway/server/plugins-http.runtime-scopes.test.ts index e4d0efb40266..492848a5d7e6 100644 --- a/src/gateway/server/plugins-http.runtime-scopes.test.ts +++ b/src/gateway/server/plugins-http.runtime-scopes.test.ts @@ -11,6 +11,7 @@ import { setActivePluginRegistry, } from "../../plugins/runtime.js"; import { getPluginRuntimeGatewayRequestScope } from "../../plugins/runtime/gateway-request-scope.js"; +import { resolveControlPlaneRateLimitKey } from "../control-plane-rate-limit.js"; import { ExecApprovalManager } from "../exec-approval-manager.js"; import type { AuthorizedGatewayHttpRequest } from "../http-utils.js"; import { authorizeOperatorScopesForMethod, CLI_DEFAULT_OPERATOR_SCOPES } from "../method-scopes.js"; @@ -248,6 +249,60 @@ describe("plugin HTTP route runtime scopes", () => { }); }); + it("threads resolved HTTP client IP into the shared control-plane rate identity", async () => { + const observed: Array<{ clientIp?: string; connId?: string; rateLimitKey: string }> = []; + const handler = createPluginRequestHandler({ + routes: [ + createRoute({ + path: SECURE_HOOK_PATH, + auth: "gateway", + gatewayRuntimeScopeSurface: "trusted-operator", + gatewayMethodDispatchAllowed: true, + handler: async () => { + const client = getPluginRuntimeGatewayRequestScope()?.client; + observed.push({ + ...(client?.clientIp ? { clientIp: client.clientIp } : {}), + ...(client?.connId ? { connId: client.connId } : {}), + rateLimitKey: resolveControlPlaneRateLimitKey(client ?? null), + }); + return true; + }, + }), + ], + }); + + for (const clientIp of ["203.0.113.10", "203.0.113.10", "203.0.113.11"]) { + const { handled, res } = await dispatchPluginRequest(handler, { + path: SECURE_HOOK_PATH, + authContext: { + gatewayAuthSatisfied: true, + gatewayRequestAuth: { authMethod: "token", trustDeclaredOperatorScopes: false }, + gatewayRequestClientIp: clientIp, + }, + }); + expect(handled).toBe(true); + expect(res.statusCode).toBe(200); + } + + expect(observed).toEqual([ + { + clientIp: "203.0.113.10", + connId: "plugin-http:203.0.113.10", + rateLimitKey: "unknown-device|203.0.113.10", + }, + { + clientIp: "203.0.113.10", + connId: "plugin-http:203.0.113.10", + rateLimitKey: "unknown-device|203.0.113.10", + }, + { + clientIp: "203.0.113.11", + connId: "plugin-http:203.0.113.11", + rateLimitKey: "unknown-device|203.0.113.11", + }, + ]); + }); + it("uses server-local routes and gateway context when the active registry belongs to another gateway", async () => { const serverAContext = { label: "server-a" } as unknown as GatewayRequestContext; const serverBContext = { label: "server-b" } as unknown as GatewayRequestContext; diff --git a/src/gateway/server/plugins-http.ts b/src/gateway/server/plugins-http.ts index 2b8746fce66d..06969dc81fa3 100644 --- a/src/gateway/server/plugins-http.ts +++ b/src/gateway/server/plugins-http.ts @@ -46,8 +46,11 @@ function resolvePluginRoutePathContextForRequest( function createPluginRouteRuntimeClient( scopes: readonly string[], + clientIp: string | undefined, ): GatewayRequestOptions["client"] { return { + connId: `plugin-http:${clientIp ?? "unknown"}`, + ...(clientIp ? { clientIp } : {}), connect: { minProtocol: PROTOCOL_VERSION, maxProtocol: PROTOCOL_VERSION, @@ -71,6 +74,7 @@ function writeUpgradeUnauthorized(socket: Duplex) { type PluginRouteRuntimeDispatchContext = { gatewayRequestAuth?: AuthorizedGatewayHttpRequest; gatewayRequestOperatorScopes?: readonly string[]; + gatewayRequestClientIp?: string; }; function getMissingPluginRouteRuntimeContext( @@ -92,6 +96,7 @@ function createPluginRouteRuntimeScope(params: { gatewayRequestContext?: GatewayRequestContext; gatewayRequestAuth?: AuthorizedGatewayHttpRequest; gatewayRequestOperatorScopes?: readonly string[]; + gatewayRequestClientIp?: string; }): PluginRouteRuntimeScope { const runtimeScopes = params.route.auth !== "gateway" @@ -103,7 +108,10 @@ function createPluginRouteRuntimeScope(params: { "trusted-operator", ) : params.gatewayRequestOperatorScopes!; - const runtimeClient = createPluginRouteRuntimeClient(runtimeScopes); + const runtimeClient = createPluginRouteRuntimeClient( + runtimeScopes, + params.gatewayRequestClientIp, + ); return { ...(params.gatewayRequestContext ? { context: params.gatewayRequestContext } : {}), client: runtimeClient, @@ -120,6 +128,7 @@ export type PluginRouteDispatchContext = { gatewayAuthSatisfied?: boolean; gatewayRequestAuth?: AuthorizedGatewayHttpRequest; gatewayRequestOperatorScopes?: readonly string[]; + gatewayRequestClientIp?: string; }; export type PluginHttpRequestHandler = ( @@ -189,6 +198,7 @@ export function createGatewayPluginRequestHandler(params: { gatewayRequestContext, gatewayRequestAuth, gatewayRequestOperatorScopes, + gatewayRequestClientIp: dispatchContext?.gatewayRequestClientIp, }), async () => route.handler(req, res), ); @@ -265,6 +275,7 @@ export function createGatewayPluginUpgradeHandler(params: { gatewayRequestContext, gatewayRequestAuth, gatewayRequestOperatorScopes, + gatewayRequestClientIp: dispatchContext?.gatewayRequestClientIp, }), async () => route.handleUpgrade?.(req, socket, head), ); 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 new file mode 100644 index 000000000000..84521024a81e --- /dev/null +++ b/src/gateway/server/ws-connection/message-handler.suspension-admission.test.ts @@ -0,0 +1,229 @@ +// WebSocket connect suspension tests cover root admission before handshake mutations. +import type { IncomingMessage } from "node:http"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import type { WebSocket } from "ws"; +import { PROTOCOL_VERSION } from "../../../../packages/gateway-protocol/src/index.js"; +import { + getActiveGatewayRootWorkCount, + resetGatewayWorkAdmission, + tryBeginGatewaySuspendAdmission, +} from "../../../process/gateway-work-admission.js"; +import type { GatewayRequestContext } from "../../server-methods/types.js"; + +const { incrementPresenceVersionMock, loadConfigMock, upsertPresenceMock } = vi.hoisted(() => ({ + incrementPresenceVersionMock: vi.fn(() => 2), + loadConfigMock: vi.fn(() => ({ gateway: { auth: { mode: "none" } } })), + upsertPresenceMock: vi.fn(), +})); + +vi.mock("../../../config/config.js", () => ({ + getRuntimeConfig: loadConfigMock, + loadConfig: loadConfigMock, +})); +vi.mock("../../../config/io.js", () => ({ + getRuntimeConfig: loadConfigMock, +})); +vi.mock("../../../infra/system-presence.js", () => ({ + upsertPresence: upsertPresenceMock, +})); +vi.mock("../health-state.js", () => ({ + buildGatewaySnapshot: vi.fn(() => ({ + presence: [], + health: {}, + stateVersion: { presence: 1, health: 1 }, + uptimeMs: 1, + sessionDefaults: { + defaultAgentId: "main", + mainKey: "main", + mainSessionKey: "main", + scope: "per-sender", + }, + })), + getHealthCache: vi.fn(() => null), + getHealthVersion: vi.fn(() => 1), + incrementPresenceVersion: incrementPresenceVersionMock, +})); + +import { attachGatewayWsMessageHandler } from "./message-handler.js"; + +function createLogger() { + return { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; +} + +function attachHarness(params: { deferSocketSend?: boolean } = {}) { + let onMessage: ((data: string) => void) | undefined; + let finishSocketSend: (() => void) | undefined; + let client: unknown = null; + const socketSend = vi.fn((_payload: string, callback?: (error?: Error) => void) => { + if (params.deferSocketSend) { + finishSocketSend = () => callback?.(); + return; + } + callback?.(); + }); + const socket = { + _receiver: {}, + send: socketSend, + on: vi.fn((event: string, handler: (data: string) => void) => { + if (event === "message") { + onMessage = handler; + } + return socket; + }), + } as unknown as WebSocket; + const close = vi.fn(); + const setClient = vi.fn((next: unknown) => { + client = next; + return true; + }); + + attachGatewayWsMessageHandler({ + socket, + upgradeReq: { + headers: { host: "127.0.0.1:19001" }, + socket: { localAddress: "127.0.0.1", remoteAddress: "127.0.0.1" }, + } as unknown as IncomingMessage, + connId: "suspension-connect", + remoteAddr: "127.0.0.1", + localAddr: "127.0.0.1", + requestHost: "127.0.0.1:19001", + connectNonce: "suspension-connect-nonce", + getResolvedAuth: () => ({ mode: "none", allowTailscale: false }), + gatewayMethods: [], + events: [], + extraHandlers: {}, + buildRequestContext: () => ({}) as GatewayRequestContext, + refreshHealthSnapshot: vi.fn(async () => ({}) as never), + send: vi.fn(), + close, + isClosed: vi.fn(() => false), + clearHandshakeTimer: vi.fn(), + getClient: () => client as never, + setClient: setClient as never, + setHandshakeState: vi.fn(), + advanceHandshakePhase: vi.fn(), + setCloseCause: vi.fn(), + setLastFrameMeta: vi.fn(), + originCheckMetrics: { hostHeaderFallbackAccepted: 0 }, + logGateway: createLogger() as never, + logHealth: createLogger() as never, + logWsControl: createLogger() as never, + }); + if (!onMessage) { + throw new Error("expected websocket message handler"); + } + + return { + close, + finishSocketSend: () => finishSocketSend?.(), + get client() { + return client; + }, + sendConnect: () => + onMessage?.( + JSON.stringify({ + type: "req", + id: "connect-1", + method: "connect", + params: { + minProtocol: PROTOCOL_VERSION, + maxProtocol: PROTOCOL_VERSION, + client: { + id: "gateway-client", + version: "dev", + platform: "test", + mode: "backend", + }, + role: "operator", + scopes: [], + caps: [], + }, + }), + ), + setClient, + socketSend, + }; +} + +beforeEach(() => { + resetGatewayWorkAdmission(); + vi.clearAllMocks(); +}); + +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(); + + 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; + }; + }; + 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"); + }); + + if (phase === "prepared") { + suspension?.release(); + } else { + suspension?.rollback(); + } + }, + ); + + it("keeps an accepted handshake visible as root work until hello is sent", async () => { + const harness = attachHarness({ deferSocketSend: true }); + + harness.sendConnect(); + + await vi.waitFor(() => { + expect(harness.socketSend).toHaveBeenCalledOnce(); + }); + expect(getActiveGatewayRootWorkCount()).toBe(1); + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension).not.toBeNull(); + expect(getActiveGatewayRootWorkCount()).toBe(1); + expect(suspension?.rollback()).toBe(true); + + harness.finishSocketSend(); + await vi.waitFor(() => { + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); + }); +}); diff --git a/src/gateway/server/ws-connection/message-handler.ts b/src/gateway/server/ws-connection/message-handler.ts index e7ccf3d13955..6742c62aa847 100644 --- a/src/gateway/server/ws-connection/message-handler.ts +++ b/src/gateway/server/ws-connection/message-handler.ts @@ -92,6 +92,12 @@ import { loadVoiceWakeConfig } from "../../../infra/voicewake.js"; import { rawDataToString } from "../../../infra/ws.js"; import { logRejectedLargePayload } from "../../../logging/diagnostic-payload.js"; import type { createSubsystemLogger } from "../../../logging/subsystem.js"; +import { + getGatewaySuspendAdmissionPhase, + isGatewayRestartDraining, + runWithGatewayIndependentRootWorkAdmission, + tryBeginGatewayRootWorkAdmission, +} from "../../../process/gateway-work-admission.js"; import { BOOTSTRAP_HANDOFF_OPERATOR_SCOPES, isPairingSetupBootstrapProfile, @@ -188,6 +194,8 @@ import { isUnauthorizedRoleError, UnauthorizedFloodGuard } from "./unauthorized- type SubsystemLogger = ReturnType; const DEVICE_SIGNATURE_SKEW_MS = 2 * 60 * 1000; +const GATEWAY_WORK_ADMISSION_RETRY_AFTER_MS = 1_000; +const GATEWAY_WORK_ADMISSION_CLOSE_CODE = 1013; const DEVICE_CREDENTIAL_INVALIDATING_METHODS = new Set([ "device.pair.remove", "device.token.rotate", @@ -662,6 +670,11 @@ export function attachGatewayWsMessageHandler(params: GatewayWsMessageHandlerPar close(4001, `client invalidated: ${reason}`); return true; }; + const runDetachedConnectWork = (run: () => Promise, onError: (error: unknown) => void) => { + // Connect-triggered mutations outlive hello-ok. Give each tail its own + // root lease so suspension cannot report ready while one is still active. + void runWithGatewayIndependentRootWorkAdmission(run).catch(onError); + }; const handleMessage = async (data: RawData) => { if (isClosed()) { @@ -2179,10 +2192,16 @@ export function attachGatewayWsMessageHandler(params: GatewayWsMessageHandlerPar nodeIdsForPairing.add(instanceIdLocal); } for (const nodeId of nodeIdsForPairing) { - void updatePairedNodeMetadata(nodeId, { - lastConnectedAtMs: nodeSession.connectedAtMs, - }).catch((err: unknown) => - logGateway.warn(`failed to record last connect for ${nodeId}: ${formatForLog(err)}`), + runDetachedConnectWork( + async () => { + await updatePairedNodeMetadata(nodeId, { + lastConnectedAtMs: nodeSession.connectedAtMs, + }); + }, + (err) => + logGateway.warn( + `failed to record last connect for ${nodeId}: ${formatForLog(err)}`, + ), ); } recordRemoteNodeInfo({ @@ -2194,42 +2213,48 @@ export function attachGatewayWsMessageHandler(params: GatewayWsMessageHandlerPar commands: nodeSession.commands, remoteIp: nodeSession.remoteIp, }); - void refreshRemoteNodeBins({ - nodeId: nodeSession.nodeId, - platform: nodeSession.platform, - deviceFamily: nodeSession.deviceFamily, - commands: nodeSession.commands, - cfg: getRuntimeConfig(), - // The node socket is registered before macOS app command handlers finish warming. - // Delay only the connect-time probe; later skill refreshes use the live session. - readinessDelayMs: 5_000, - }).catch((err: unknown) => - logGateway.warn( - `remote bin probe failed for ${nodeSession.nodeId}: ${formatForLog(err)}`, - ), + runDetachedConnectWork( + async () => { + await refreshRemoteNodeBins({ + nodeId: nodeSession.nodeId, + platform: nodeSession.platform, + deviceFamily: nodeSession.deviceFamily, + commands: nodeSession.commands, + cfg: getRuntimeConfig(), + // The node socket is registered before macOS app command handlers finish warming. + // Delay only the connect-time probe; later skill refreshes use the live session. + readinessDelayMs: 5_000, + }); + }, + (err) => + logGateway.warn( + `remote bin probe failed for ${nodeSession.nodeId}: ${formatForLog(err)}`, + ), ); - void loadVoiceWakeConfig() - .then((cfg) => { + runDetachedConnectWork( + async () => { + const cfg = await loadVoiceWakeConfig(); context.nodeRegistry.sendEvent(nodeSession.nodeId, "voicewake.changed", { triggers: cfg.triggers, }); - }) - .catch((err: unknown) => + }, + (err) => logGateway.warn( `voicewake snapshot failed for ${nodeSession.nodeId}: ${formatForLog(err)}`, ), - ); - void loadVoiceWakeRoutingConfig() - .then((routing) => { + ); + runDetachedConnectWork( + async () => { + const routing = await loadVoiceWakeRoutingConfig(); context.nodeRegistry.sendEvent(nodeSession.nodeId, "voicewake.routing.changed", { config: routing, }); - }) - .catch((err: unknown) => + }, + (err) => logGateway.warn( `voicewake routing snapshot failed for ${nodeSession.nodeId}: ${formatForLog(err)}`, ), - ); + ); } const snapshot = buildGatewaySnapshot({ @@ -2510,8 +2535,80 @@ export function attachGatewayWsMessageHandler(params: GatewayWsMessageHandlerPar } }; + const rejectConnectForClosedAdmission = async (data: RawData): Promise => { + if (isClosed() || getRawDataByteLength(data) > MAX_PREAUTH_PAYLOAD_BYTES) { + return false; + } + let parsed: unknown; + try { + parsed = JSON.parse(rawDataToString(data)); + } catch { + return false; + } + if ( + !validateRequestFrame(parsed) || + parsed.method !== "connect" || + !validateConnectParams(parsed.params) + ) { + return false; + } + + const restartDraining = isGatewayRestartDraining(); + const reason = restartDraining ? "gateway-restarting" : "gateway-suspending"; + const operation = restartDraining ? "restart" : "suspension"; + const phase = getGatewaySuspendAdmissionPhase(); + setLastFrameMeta({ type: "req", method: "connect", id: parsed.id }); + setHandshakeState("failed"); + setCloseCause(reason, { + method: "connect", + phase, + }); + await sendFrame({ + type: "res", + id: parsed.id, + ok: false, + error: errorShape(ErrorCodes.UNAVAILABLE, `connect unavailable during gateway ${operation}`, { + retryable: true, + retryAfterMs: GATEWAY_WORK_ADMISSION_RETRY_AFTER_MS, + details: { + method: "connect", + reason, + phase, + }, + }), + }).catch(() => {}); + queueMicrotask(() => + close(GATEWAY_WORK_ADMISSION_CLOSE_CODE, `gateway ${operation} in progress`), + ); + return true; + }; + + const handleIncomingMessage = async (data: RawData) => { + if (getClient()) { + await handleMessage(data); + return; + } + const admission = tryBeginGatewayRootWorkAdmission(); + if (!admission) { + if (await rejectConnectForClosedAdmission(data)) { + return; + } + // Malformed pre-auth frames still use the established validation and + // close path; only a validated connect can cross into mutable work. + await handleMessage(data); + return; + } + try { + await admission.run(() => handleMessage(data)); + } finally { + admission.release(); + } + }; + socket.on("message", (data) => { - void runWithDiagnosticTraceContext(createDiagnosticTraceContext(), () => handleMessage(data)); + void runWithDiagnosticTraceContext(createDiagnosticTraceContext(), () => + handleIncomingMessage(data), + ); }); } diff --git a/src/gateway/server/ws-types.ts b/src/gateway/server/ws-types.ts index b8c314c4cf9f..e27641177c66 100644 --- a/src/gateway/server/ws-types.ts +++ b/src/gateway/server/ws-types.ts @@ -1,6 +1,6 @@ // Gateway WebSocket client types describe authenticated client state retained by the server. import type { WebSocket } from "ws"; -import type { ConnectParams } from "../../../packages/gateway-protocol/src/index.js"; +import type { ConnectParams } from "../../../packages/gateway-protocol/src/schema/frames.js"; import type { AgentRuntimeIdentity } from "../agent-runtime-identity-token.js"; import type { PluginNodeCapabilityClient } from "../plugin-node-capability.js"; diff --git a/src/gateway/session-reset-service.ts b/src/gateway/session-reset-service.ts index b5c1d66414dd..894483de001b 100644 --- a/src/gateway/session-reset-service.ts +++ b/src/gateway/session-reset-service.ts @@ -50,6 +50,7 @@ import { getSessionBindingService } from "../infra/outbound/session-binding-serv import { getGlobalHookRunner } from "../plugins/hook-runner-global.js"; import { runPluginHostCleanup } from "../plugins/host-hook-cleanup.js"; import { getActivePluginRegistry } from "../plugins/runtime.js"; +import { runWithGatewayIndependentRootWorkContinuation } from "../process/gateway-work-admission.js"; import { isSubagentSessionKey, normalizeAgentId, @@ -188,7 +189,9 @@ export function emitGatewaySessionEndPluginHook(params: { nextSessionId: params.nextSessionId, nextSessionKey: params.nextSessionKey, }); - void hookRunner.runSessionEnd(payload.event, payload.context).catch((err: unknown) => { + void runWithGatewayIndependentRootWorkContinuation(async () => { + await hookRunner.runSessionEnd(payload.event, payload.context); + }).catch((err: unknown) => { logVerbose(`session_end hook failed: ${String(err)}`); }); } @@ -231,7 +234,9 @@ export function emitGatewaySessionStartPluginHook(params: { cfg: params.cfg, resumedFrom: params.resumedFrom, }); - void hookRunner.runSessionStart(payload.event, payload.context).catch((err: unknown) => { + void runWithGatewayIndependentRootWorkContinuation(async () => { + await hookRunner.runSessionStart(payload.event, payload.context); + }).catch((err: unknown) => { logVerbose(`session_start hook failed: ${String(err)}`); }); } diff --git a/src/gateway/test-helpers.server.ts b/src/gateway/test-helpers.server.ts index a6eec3501cf6..686d4244e4fe 100644 --- a/src/gateway/test-helpers.server.ts +++ b/src/gateway/test-helpers.server.ts @@ -26,10 +26,13 @@ import { getPairedDevice, requestDevicePairing, } from "../infra/device-pairing.js"; +import { resetGatewaySuspendCoordinatorForTest } from "../infra/gateway-suspend-coordinator.js"; +import { __testing as restartTesting } from "../infra/restart.js"; import { drainSystemEvents, peekSystemEvents } from "../infra/system-events.js"; import { rawDataToString } from "../infra/ws.js"; import { resetLogger, setLoggerOverride } from "../logging.js"; import { clearGatewaySubagentRuntime } from "../plugins/runtime/gateway-bindings.js"; +import { resetGatewayWorkAdmission } from "../process/gateway-work-admission.js"; import { DEFAULT_AGENT_ID, normalizeMainKey, @@ -250,9 +253,22 @@ function applyGatewaySkipEnv() { : "openclaw-test-no-bundled-extensions"; } +function resetGatewayLifecycleTestState(options: { preserveRuntimeBindings: boolean }): void { + // Resume a held scheduler before hard admission reset invalidates and forgets its + // lease. Then cancel restart timers and retire their module-local signal lease. + resetGatewaySuspendCoordinatorForTest(); + if (options.preserveRuntimeBindings) { + restartTesting.resetSigusr1TransientState(); + } else { + restartTesting.resetSigusr1State(); + } + resetGatewayWorkAdmission(); +} + async function resetGatewayTestState(options: { uniqueConfigRoot: boolean }) { // Some tests intentionally use fake timers; ensure they don't leak into gateway suites. vi.useRealTimers(); + resetGatewayLifecycleTestState({ preserveRuntimeBindings: false }); setLoggerOverride({ level: "silent", consoleLevel: "silent" }); if (!tempHome) { throw new Error("resetGatewayTestState called before temp home was initialized"); @@ -369,6 +385,7 @@ async function resetGatewayTestState(options: { uniqueConfigRoot: boolean }) { async function cleanupGatewayTestHome(options: { restoreEnv: boolean }) { vi.useRealTimers(); + resetGatewayLifecycleTestState({ preserveRuntimeBindings: activeSuiteGatewayServerCount > 0 }); clearGatewaySubagentRuntime(); resetLogger(); resetTaskRegistryForTests({ persist: false }); @@ -395,6 +412,7 @@ async function cleanupGatewayTestHome(options: { restoreEnv: boolean }) { async function resetGatewayTestRuntimeOnly() { vi.useRealTimers(); + resetGatewayLifecycleTestState({ preserveRuntimeBindings: true }); setLoggerOverride({ level: "silent", consoleLevel: "silent" }); applyGatewaySkipEnv(); delete process.env.OPENCLAW_GATEWAY_TOKEN; diff --git a/src/hooks/internal-hooks.ts b/src/hooks/internal-hooks.ts index 47f9d9f1ab17..0fa3c1fe9d8d 100644 --- a/src/hooks/internal-hooks.ts +++ b/src/hooks/internal-hooks.ts @@ -5,7 +5,7 @@ * like command processing, session lifecycle, etc. */ -import type { SessionsPatchParams } from "../../packages/gateway-protocol/src/schema.js"; +import type { SessionsPatchParams } from "../../packages/gateway-protocol/src/schema/sessions.js"; import type { WorkspaceBootstrapFile } from "../agents/workspace.js"; import type { CliDeps } from "../cli/outbound-send-deps.js"; import type { SessionEntry } from "../config/sessions/types.js"; diff --git a/src/infra/exec-approval-channel-runtime.ts b/src/infra/exec-approval-channel-runtime.ts index 3845e253150d..4e76f43fbf93 100644 --- a/src/infra/exec-approval-channel-runtime.ts +++ b/src/infra/exec-approval-channel-runtime.ts @@ -1,6 +1,6 @@ // Runs the gateway-backed runtime that delivers native approval events. import { readConnectErrorDetailCode } from "../../packages/gateway-protocol/src/connect-error-details.js"; -import type { EventFrame } from "../../packages/gateway-protocol/src/index.js"; +import type { EventFrame } from "../../packages/gateway-protocol/src/schema/frames.js"; import { startGatewayClientWhenEventLoopReady } from "../gateway/client-start-readiness.js"; import type { GatewayClient, GatewayReconnectPausedInfo } from "../gateway/client.js"; import { isApprovalMethod } from "../gateway/method-scopes.js"; diff --git a/src/infra/gateway-active-work.ts b/src/infra/gateway-active-work.ts new file mode 100644 index 000000000000..a72f6e632fe3 --- /dev/null +++ b/src/infra/gateway-active-work.ts @@ -0,0 +1,188 @@ +// Collects process activity shared by restart and host-suspension decisions. +import { getActiveEmbeddedRunCount } from "../agents/embedded-agent-runner/run-state.js"; +import { getTotalPendingReplies } from "../auto-reply/reply/dispatcher-registry.js"; +import { getActiveCronJobCount } from "../cron/active-jobs.js"; +import { getTotalQueueSize } from "../process/command-queue.js"; +import { getActiveGatewayRootWorkCount } from "../process/gateway-work-admission.js"; +import { + getActiveSessionLifecycleMutationCount, + getActiveSessionWorkAdmissionCount, +} from "../sessions/session-lifecycle-admission.js"; +import { getSuspensionVisibleCronTaskRunCount } from "../tasks/cron-task-cancel.js"; +import { getInspectableActiveTaskRestartBlockers } from "../tasks/task-registry.maintenance.js"; +import { + type ActiveTaskRestartBlocker, + formatActiveTaskRestartBlocker, +} from "../tasks/task-restart-blocker.js"; + +export type GatewayActiveWorkCounts = { + queueSize: number; + pendingReplies: number; + embeddedRuns: number; + cronRuns: number; + activeTasks: number; + rootRequests: number; + sessionAdmissions: number; + sessionMutations: number; + chatRuns: number; + queuedTurns: number; + terminalPersistence: number; + terminalSessions: number; + /** Compatibility aggregate. Categories can overlap; use individual counts for diagnostics. */ + totalActive: number; +}; + +export type GatewayActiveWorkBlocker = { + kind: + | "queue" + | "reply" + | "embedded-run" + | "cron-run" + | "task" + | "root-request" + | "session-admission" + | "session-mutation" + | "chat-run" + | "queued-turn" + | "terminal-persistence" + | "terminal-session"; + count: number; + message: string; + task?: ActiveTaskRestartBlocker; +}; + +export type GatewayActiveWorkSnapshot = { + idle: boolean; + counts: GatewayActiveWorkCounts; + blockers: GatewayActiveWorkBlocker[]; +}; + +export type GatewayActiveWorkInspectors = { + getQueueSize: () => number; + getPendingReplies: () => number; + getEmbeddedRuns: () => number; + getCronRuns: () => number; + getActiveTasks: () => number; + getTaskBlockers: () => ActiveTaskRestartBlocker[]; + getRootRequests: () => number; + getSessionAdmissions: () => number; + getSessionMutations: () => number; + getChatRuns: () => number; + getQueuedTurns: () => number; + getTerminalPersistence: () => number; + getTerminalSessions: () => number; +}; + +const defaultInspectors: GatewayActiveWorkInspectors = { + getQueueSize: getTotalQueueSize, + getPendingReplies: getTotalPendingReplies, + getEmbeddedRuns: getActiveEmbeddedRunCount, + getCronRuns: () => Math.max(getActiveCronJobCount(), getSuspensionVisibleCronTaskRunCount()), + getActiveTasks: () => getInspectableActiveTaskRestartBlockers().length, + getTaskBlockers: getInspectableActiveTaskRestartBlockers, + getRootRequests: () => getActiveGatewayRootWorkCount({ excludeCurrent: true }), + getSessionAdmissions: getActiveSessionWorkAdmissionCount, + getSessionMutations: getActiveSessionLifecycleMutationCount, + getChatRuns: () => 0, + getQueuedTurns: () => 0, + getTerminalPersistence: () => 0, + getTerminalSessions: () => 0, +}; + +function normalizeCount(value: number): number { + return Number.isFinite(value) && value > 0 ? Math.floor(value) : 0; +} + +export function createGatewayActiveWorkSnapshot( + inspectors: Partial = {}, +): GatewayActiveWorkSnapshot { + const resolved = { ...defaultInspectors, ...inspectors }; + const counts: GatewayActiveWorkCounts = { + queueSize: normalizeCount(resolved.getQueueSize()), + pendingReplies: normalizeCount(resolved.getPendingReplies()), + embeddedRuns: normalizeCount(resolved.getEmbeddedRuns()), + cronRuns: normalizeCount(resolved.getCronRuns()), + activeTasks: normalizeCount(resolved.getActiveTasks()), + rootRequests: normalizeCount(resolved.getRootRequests()), + sessionAdmissions: normalizeCount(resolved.getSessionAdmissions()), + sessionMutations: normalizeCount(resolved.getSessionMutations()), + chatRuns: normalizeCount(resolved.getChatRuns()), + queuedTurns: normalizeCount(resolved.getQueuedTurns()), + terminalPersistence: normalizeCount(resolved.getTerminalPersistence()), + terminalSessions: normalizeCount(resolved.getTerminalSessions()), + totalActive: 0, + }; + counts.totalActive = Object.entries(counts).reduce( + (total, [key, count]) => (key === "totalActive" ? total : total + count), + 0, + ); + + const blockers: GatewayActiveWorkBlocker[] = []; + const add = (count: number, kind: GatewayActiveWorkBlocker["kind"], message: string) => { + if (count > 0) { + blockers.push({ kind, count, message }); + } + }; + add(counts.queueSize, "queue", `${counts.queueSize} queued or active operation(s)`); + add( + counts.pendingReplies, + "reply", + `${counts.pendingReplies} pending reply delivery operation(s)`, + ); + add(counts.embeddedRuns, "embedded-run", `${counts.embeddedRuns} active embedded run(s)`); + add(counts.cronRuns, "cron-run", `${counts.cronRuns} active cron run(s)`); + add(counts.rootRequests, "root-request", `${counts.rootRequests} active gateway request(s)`); + add( + counts.sessionAdmissions, + "session-admission", + `${counts.sessionAdmissions} admitted session turn(s)`, + ); + add( + counts.sessionMutations, + "session-mutation", + `${counts.sessionMutations} active session lifecycle mutation(s)`, + ); + add(counts.chatRuns, "chat-run", `${counts.chatRuns} active chat run(s)`); + add(counts.queuedTurns, "queued-turn", `${counts.queuedTurns} queued chat turn(s)`); + add( + counts.terminalPersistence, + "terminal-persistence", + `${counts.terminalPersistence} pending terminal session write(s)`, + ); + add( + counts.terminalSessions, + "terminal-session", + `${counts.terminalSessions} open terminal session(s)`, + ); + + if (counts.activeTasks > 0) { + const taskBlockers = resolved.getTaskBlockers(); + if (taskBlockers.length === 0) { + blockers.push({ + kind: "task", + count: counts.activeTasks, + message: `${counts.activeTasks} active background task run(s)`, + }); + } else { + const shownTaskBlockers = taskBlockers.slice(0, 8); + for (const task of shownTaskBlockers) { + blockers.push({ + kind: "task", + count: 1, + message: formatActiveTaskRestartBlocker(task), + task, + }); + } + const omitted = counts.activeTasks - shownTaskBlockers.length; + if (omitted > 0) { + blockers.push({ + kind: "task", + count: omitted, + message: `${omitted} additional active background task run(s)`, + }); + } + } + } + + return { idle: counts.totalActive === 0, counts, blockers }; +} diff --git a/src/infra/gateway-suspend-coordinator.test.ts b/src/infra/gateway-suspend-coordinator.test.ts new file mode 100644 index 000000000000..a0b087a4ce8a --- /dev/null +++ b/src/infra/gateway-suspend-coordinator.test.ts @@ -0,0 +1,404 @@ +// Covers atomic refuse-only suspension preparation, renewal, and release. +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { + isGatewayWorkAdmissionClosed, + markGatewayRestartDraining, + resetGatewayWorkAdmission, +} from "../process/gateway-work-admission.js"; +import type { GatewayActiveWorkInspectors } from "./gateway-active-work.js"; +import { + GATEWAY_SUSPEND_TTL_MS, + getGatewaySuspendStatus, + prepareGatewaySuspend, + resetGatewaySuspendCoordinatorForTest, + resumeGatewaySuspend, +} from "./gateway-suspend-coordinator.js"; + +function inspectors( + overrides: Partial = {}, +): GatewayActiveWorkInspectors { + return { + getQueueSize: () => 0, + getPendingReplies: () => 0, + getEmbeddedRuns: () => 0, + getCronRuns: () => 0, + getActiveTasks: () => 0, + getTaskBlockers: () => [], + getRootRequests: () => 0, + getSessionAdmissions: () => 0, + getSessionMutations: () => 0, + getChatRuns: () => 0, + getQueuedTurns: () => 0, + getTerminalPersistence: () => 0, + getTerminalSessions: () => 0, + ...overrides, + }; +} + +beforeEach(() => { + resetGatewaySuspendCoordinatorForTest(); + resetGatewayWorkAdmission(); +}); + +afterEach(() => { + resetGatewaySuspendCoordinatorForTest(); + resetGatewayWorkAdmission(); +}); + +describe("gateway suspend coordinator", () => { + it("test reset resumes a held scheduler before admission is cleared", () => { + const resumeScheduling = vi.fn(() => { + expect(isGatewayWorkAdmissionClosed()).toBe(true); + }); + expect( + prepareGatewaySuspend({ + requestId: "request-test-reset", + pauseScheduling: vi.fn(), + resumeScheduling, + inspect: inspectors(), + }), + ).toMatchObject({ status: "ready" }); + + resetGatewaySuspendCoordinatorForTest(); + resetGatewayWorkAdmission(); + + expect(resumeScheduling).toHaveBeenCalledOnce(); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + }); + + it("reopens admission in the same turn when active work refuses preparation", () => { + const events: string[] = []; + const result = prepareGatewaySuspend({ + requestId: "request-busy", + pauseScheduling: () => events.push("pause"), + resumeScheduling: () => events.push("resume"), + inspect: inspectors({ + getQueueSize: () => { + events.push("inspect"); + return 1; + }, + }), + }); + + expect(result.status).toBe("busy"); + expect(events).toEqual(["pause", "inspect", "resume"]); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + }); + + it("keeps admission closed until a failed busy rollback resumes scheduling", () => { + vi.useFakeTimers(); + try { + const resumeScheduling = vi + .fn() + .mockImplementationOnce(() => { + throw new Error("timer unavailable"); + }) + .mockImplementationOnce(() => {}); + const first = prepareGatewaySuspend({ + requestId: "request-busy-resume-retry", + pauseScheduling: vi.fn(), + resumeScheduling, + inspect: inspectors({ getQueueSize: () => 1 }), + }); + + expect(first).toEqual({ + status: "recovering", + reason: "scheduler-resume-failed", + retryAfterMs: 1_000, + }); + expect(isGatewayWorkAdmissionClosed()).toBe(true); + expect(getGatewaySuspendStatus("stale-id")).toEqual(first); + expect(resumeGatewaySuspend("stale-id")).toEqual({ + ok: false, + reason: "scheduler-resume-failed", + retryAfterMs: 1_000, + }); + expect( + prepareGatewaySuspend({ + requestId: "request-before-scheduler-resume", + pauseScheduling: vi.fn(), + resumeScheduling, + inspect: inspectors(), + }), + ).toEqual(first); + + vi.advanceTimersByTime(1_000); + expect(resumeScheduling).toHaveBeenCalledTimes(2); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + expect(getGatewaySuspendStatus("stale-id")).toEqual({ status: "running" }); + + expect( + prepareGatewaySuspend({ + requestId: "request-after-scheduler-resume", + pauseScheduling: vi.fn(), + resumeScheduling, + inspect: inspectors(), + createSuspensionId: () => "suspension-after-scheduler-resume", + }), + ).toMatchObject({ + status: "ready", + suspensionId: "suspension-after-scheduler-resume", + }); + vi.advanceTimersByTime(1_000); + expect(resumeScheduling).toHaveBeenCalledTimes(2); + } finally { + vi.useRealTimers(); + } + }); + + it("cancels scheduler recovery when restart supersedes suspension", () => { + vi.useFakeTimers(); + try { + const resumeScheduling = vi.fn(() => { + throw new Error("timer unavailable"); + }); + expect( + prepareGatewaySuspend({ + requestId: "request-recovery-restart", + pauseScheduling: vi.fn(), + resumeScheduling, + inspect: inspectors({ getQueueSize: () => 1 }), + }), + ).toMatchObject({ status: "recovering" }); + + markGatewayRestartDraining(); + vi.advanceTimersByTime(1_000); + + expect(resumeScheduling).toHaveBeenCalledOnce(); + expect(isGatewayWorkAdmissionClosed()).toBe(true); + expect(getGatewaySuspendStatus("stale-id")).toEqual({ status: "running" }); + } finally { + vi.useRealTimers(); + } + }); + + it("owns recovery when inspection fails before admission commits", () => { + vi.useFakeTimers(); + try { + const resumeScheduling = vi + .fn() + .mockImplementationOnce(() => { + throw new Error("timer unavailable"); + }) + .mockImplementationOnce(() => {}); + const result = prepareGatewaySuspend({ + requestId: "request-inspection-failure", + pauseScheduling: vi.fn(), + resumeScheduling, + inspect: inspectors({ + getQueueSize: () => { + throw new Error("inspection failed"); + }, + }), + }); + + expect(result).toMatchObject({ status: "recovering" }); + expect(isGatewayWorkAdmissionClosed()).toBe(true); + vi.advanceTimersByTime(1_000); + expect(resumeScheduling).toHaveBeenCalledTimes(2); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + } finally { + vi.useRealTimers(); + } + }); + + it("renews one ready lease and resumes only with the matching id", () => { + const resumeScheduling = vi.fn(); + const first = prepareGatewaySuspend({ + requestId: "request-ready", + pauseScheduling: vi.fn(), + resumeScheduling, + inspect: inspectors(), + nowMs: () => 1_000, + createSuspensionId: () => "suspension-1", + }); + expect(first).toMatchObject({ + status: "ready", + suspensionId: "suspension-1", + expiresAtMs: 1_000 + GATEWAY_SUSPEND_TTL_MS, + }); + expect(isGatewayWorkAdmissionClosed()).toBe(true); + expect(getGatewaySuspendStatus("suspension-1")).toEqual({ + status: "ready", + expiresAtMs: 1_000 + GATEWAY_SUSPEND_TTL_MS, + }); + + const renewed = prepareGatewaySuspend({ + requestId: "request-ready", + pauseScheduling: vi.fn(), + resumeScheduling, + inspect: inspectors({ getQueueSize: () => 99 }), + nowMs: () => 2_000, + }); + expect(renewed).toMatchObject({ + status: "ready", + suspensionId: "suspension-1", + expiresAtMs: 2_000 + GATEWAY_SUSPEND_TTL_MS, + }); + expect( + prepareGatewaySuspend({ + requestId: "request-other", + pauseScheduling: vi.fn(), + resumeScheduling, + }).status, + ).toBe("conflict"); + + expect(resumeGatewaySuspend("wrong-id")).toEqual({ + ok: false, + reason: "suspension-mismatch", + }); + expect(isGatewayWorkAdmissionClosed()).toBe(true); + expect(resumeGatewaySuspend("suspension-1")).toEqual({ + ok: true, + status: "running", + resumed: true, + }); + expect(resumeScheduling).toHaveBeenCalledOnce(); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + expect(resumeGatewaySuspend("suspension-1")).toEqual({ + ok: true, + status: "running", + resumed: false, + }); + }); + + it("lets restart supersede a suspension without reopening its scheduler", () => { + const resumeScheduling = vi.fn(); + const result = prepareGatewaySuspend({ + requestId: "request-restart", + pauseScheduling: vi.fn(), + resumeScheduling, + inspect: inspectors(), + createSuspensionId: () => "suspension-restart", + }); + expect(result.status).toBe("ready"); + + markGatewayRestartDraining(); + + expect(getGatewaySuspendStatus("suspension-restart")).toEqual({ status: "running" }); + expect(resumeScheduling).not.toHaveBeenCalled(); + expect(isGatewayWorkAdmissionClosed()).toBe(true); + }); + + it("exposes scheduler recovery after a ready lease cannot resume", () => { + vi.useFakeTimers(); + try { + const resumeScheduling = vi + .fn() + .mockImplementationOnce(() => { + throw new Error("timer unavailable"); + }) + .mockImplementationOnce(() => {}); + prepareGatewaySuspend({ + requestId: "request-resume-retry", + pauseScheduling: vi.fn(), + resumeScheduling, + inspect: inspectors(), + createSuspensionId: () => "suspension-resume-retry", + }); + + expect(resumeGatewaySuspend("suspension-resume-retry")).toMatchObject({ + ok: false, + reason: "scheduler-resume-failed", + }); + expect(isGatewayWorkAdmissionClosed()).toBe(true); + expect(getGatewaySuspendStatus("suspension-resume-retry")).toMatchObject({ + status: "recovering", + }); + expect( + prepareGatewaySuspend({ + requestId: "request-resume-retry", + pauseScheduling: vi.fn(), + resumeScheduling, + inspect: inspectors(), + }), + ).toMatchObject({ status: "recovering" }); + expect(resumeGatewaySuspend("suspension-resume-retry")).toMatchObject({ + ok: false, + reason: "scheduler-resume-failed", + }); + + vi.advanceTimersByTime(1_000); + expect(resumeScheduling).toHaveBeenCalledTimes(2); + expect(getGatewaySuspendStatus("suspension-resume-retry")).toEqual({ status: "running" }); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + } finally { + vi.useRealTimers(); + } + }); + + it("auto-resumes an abandoned ready lease at expiry", () => { + vi.useFakeTimers(); + try { + const resumeScheduling = vi.fn(); + prepareGatewaySuspend({ + requestId: "request-expiry", + pauseScheduling: vi.fn(), + resumeScheduling, + inspect: inspectors(), + createSuspensionId: () => "suspension-expiry", + }); + + vi.advanceTimersByTime(GATEWAY_SUSPEND_TTL_MS); + + expect(getGatewaySuspendStatus("suspension-expiry")).toEqual({ status: "running" }); + expect(resumeScheduling).toHaveBeenCalledOnce(); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + } finally { + vi.useRealTimers(); + } + }); + + it("enters recovery when lease expiry cannot resume the scheduler", () => { + vi.useFakeTimers(); + try { + const resumeScheduling = vi + .fn() + .mockImplementationOnce(() => { + throw new Error("timer unavailable"); + }) + .mockImplementationOnce(() => {}); + prepareGatewaySuspend({ + requestId: "request-expiry-recovery", + pauseScheduling: vi.fn(), + resumeScheduling, + inspect: inspectors(), + createSuspensionId: () => "suspension-expiry-recovery", + }); + + vi.advanceTimersByTime(GATEWAY_SUSPEND_TTL_MS); + expect(getGatewaySuspendStatus("suspension-expiry-recovery")).toMatchObject({ + status: "recovering", + }); + expect(isGatewayWorkAdmissionClosed()).toBe(true); + + vi.advanceTimersByTime(1_000); + expect(resumeScheduling).toHaveBeenCalledTimes(2); + expect(getGatewaySuspendStatus("suspension-expiry-recovery")).toEqual({ + status: "running", + }); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + } finally { + vi.useRealTimers(); + } + }); + + it("expires synchronously when timer delivery is delayed", () => { + let nowMs = 10_000; + const resumeScheduling = vi.fn(); + prepareGatewaySuspend({ + requestId: "request-delayed-expiry", + pauseScheduling: vi.fn(), + resumeScheduling, + inspect: inspectors(), + nowMs: () => nowMs, + createSuspensionId: () => "suspension-delayed-expiry", + }); + + nowMs += GATEWAY_SUSPEND_TTL_MS; + + expect(getGatewaySuspendStatus("suspension-delayed-expiry")).toEqual({ status: "running" }); + expect(resumeScheduling).toHaveBeenCalledOnce(); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + }); +}); diff --git a/src/infra/gateway-suspend-coordinator.ts b/src/infra/gateway-suspend-coordinator.ts new file mode 100644 index 000000000000..264d3c2992fc --- /dev/null +++ b/src/infra/gateway-suspend-coordinator.ts @@ -0,0 +1,415 @@ +// Coordinates an atomic, refuse-only host suspension preparation lease. +import { randomUUID } from "node:crypto"; +import type { + GatewaySuspendPrepareResult as GatewaySuspendPrepareWireResult, + GatewaySuspendResumeResult as GatewaySuspendResumeWireResult, + GatewaySuspendStatusResult as GatewaySuspendStatusWireResult, +} from "../../packages/gateway-protocol/src/index.js"; +import { tryBeginGatewaySuspendAdmission } from "../process/gateway-work-admission.js"; +import { resolveGlobalSingleton } from "../shared/global-singleton.js"; +import { + createGatewayActiveWorkSnapshot, + type GatewayActiveWorkInspectors, + type GatewayActiveWorkSnapshot, +} from "./gateway-active-work.js"; + +export const GATEWAY_SUSPEND_TTL_MS = 2 * 60_000; +export const GATEWAY_SUSPEND_RETRY_AFTER_MS = 20_000; +const GATEWAY_SCHEDULER_RECOVERY_RETRY_MS = 1_000; + +type GatewaySchedulerRecoveryResult = { + status: "recovering"; + reason: "scheduler-resume-failed"; + retryAfterMs: number; +}; + +export type GatewaySuspendPrepareResult = + | GatewaySuspendPrepareWireResult + | { status: "conflict"; expiresAtMs: number } + | GatewaySchedulerRecoveryResult; + +export type GatewaySuspendStatusResult = + | GatewaySuspendStatusWireResult + | { status: "conflict"; expiresAtMs: number } + | GatewaySchedulerRecoveryResult; + +export type GatewaySuspendResumeResult = + | GatewaySuspendResumeWireResult + | { ok: false; reason: "suspension-mismatch" } + | { ok: false; reason: "scheduler-resume-failed"; retryAfterMs: number }; + +type GatewaySuspendCoordinatorEntryBase = { + owner: object; + resumeScheduling: () => void; + reopenAdmission: () => boolean; + warn?: (message: string) => void; + timer?: ReturnType; +}; + +type HeldGatewaySuspension = GatewaySuspendCoordinatorEntryBase & { + kind: "held"; + requestId: string; + suspensionId: string; + expiresAtMs: number; + snapshot: GatewayActiveWorkSnapshot; + nowMs: () => number; +}; + +type GatewaySchedulerRecovery = GatewaySuspendCoordinatorEntryBase & { + kind: "recovering"; +}; + +type GatewaySuspendCoordinatorEntry = HeldGatewaySuspension | GatewaySchedulerRecovery; + +type GatewaySuspendCoordinatorState = { + current: GatewaySuspendCoordinatorEntry | null; +}; + +const COORDINATOR_STATE = resolveGlobalSingleton( + Symbol.for("openclaw.gatewaySuspendCoordinatorState"), + (): GatewaySuspendCoordinatorState => ({ + current: null, + }), +); + +function schedulerRecoveryResult(): GatewaySchedulerRecoveryResult { + return { + status: "recovering", + reason: "scheduler-resume-failed", + retryAfterMs: GATEWAY_SCHEDULER_RECOVERY_RETRY_MS, + }; +} + +function clearEntryTimer(entry: GatewaySuspendCoordinatorEntry): void { + if (entry.timer) { + clearTimeout(entry.timer); + entry.timer = undefined; + } +} + +function scheduleEntry( + entry: GatewaySuspendCoordinatorEntry, + delayMs: number, + callback: () => void, +): void { + clearEntryTimer(entry); + entry.timer = setTimeout(callback, delayMs); + entry.timer.unref?.(); +} + +function resumeAndReopen(entry: GatewaySuspendCoordinatorEntry): boolean { + try { + entry.resumeScheduling(); + } catch (err) { + entry.warn?.(`gateway scheduler recovery failed: ${String(err)}`); + enterSchedulerRecovery(entry); + return false; + } + if (COORDINATOR_STATE.current !== entry) { + return true; + } + if (!entry.reopenAdmission()) { + entry.warn?.("gateway scheduler recovery could not reopen admission"); + enterSchedulerRecovery(entry); + return false; + } + clearEntryTimer(entry); + COORDINATOR_STATE.current = null; + return true; +} + +function enterSchedulerRecovery(entry: GatewaySuspendCoordinatorEntry): void { + if (COORDINATOR_STATE.current !== entry) { + return; + } + if (entry.kind === "recovering") { + scheduleRecoveryRetry(entry); + return; + } + clearEntryTimer(entry); + const recovery: GatewaySchedulerRecovery = { + kind: "recovering", + owner: entry.owner, + resumeScheduling: entry.resumeScheduling, + reopenAdmission: entry.reopenAdmission, + warn: entry.warn, + }; + COORDINATOR_STATE.current = recovery; + scheduleRecoveryRetry(recovery); +} + +function scheduleRecoveryRetry(entry: GatewaySuspendCoordinatorEntry): void { + scheduleEntry(entry, GATEWAY_SCHEDULER_RECOVERY_RETRY_MS, () => { + if (COORDINATOR_STATE.current === entry) { + resumeAndReopen(entry); + } + }); +} + +function normalizeExpiredHeldSuspension( + held: HeldGatewaySuspension, +): GatewaySuspendCoordinatorEntry | null { + if (held.nowMs() < held.expiresAtMs) { + return held; + } + resumeAndReopen(held); + return COORDINATOR_STATE.current; +} + +function armSchedulerRecovery( + recovery: Omit, +): GatewaySchedulerRecovery { + const entry: GatewaySchedulerRecovery = { kind: "recovering", ...recovery }; + scheduleRecoveryRetry(entry); + return entry; +} + +// Rollback stays fail-closed: scheduler recovery must finish before admission +// reopens, otherwise an old retry can resume scheduling under a newer lease. +function resumeSchedulingBeforeReopen(params: { + owner: object; + resumeScheduling: () => void; + reopenAdmission: () => boolean; + isInvalidated: () => boolean; + warn?: (message: string) => void; +}): boolean { + if (params.isInvalidated()) { + return true; + } + try { + params.resumeScheduling(); + } catch (err) { + params.warn?.(`gateway scheduler resume failed during suspension rollback: ${String(err)}`); + COORDINATOR_STATE.current = armSchedulerRecovery({ + owner: params.owner, + resumeScheduling: params.resumeScheduling, + reopenAdmission: params.reopenAdmission, + warn: params.warn, + }); + return false; + } + if (!params.isInvalidated()) { + params.reopenAdmission(); + } + return true; +} + +function armExpiry(held: Omit): HeldGatewaySuspension { + const entry: HeldGatewaySuspension = { kind: "held", ...held }; + scheduleEntry(entry, GATEWAY_SUSPEND_TTL_MS, () => { + if (COORDINATOR_STATE.current === entry) { + resumeAndReopen(entry); + } + }); + return entry; +} + +function renewHeldSuspension(held: HeldGatewaySuspension, nowMs: number): void { + held.expiresAtMs = nowMs + GATEWAY_SUSPEND_TTL_MS; + scheduleEntry(held, GATEWAY_SUSPEND_TTL_MS, () => { + if (COORDINATOR_STATE.current === held) { + resumeAndReopen(held); + } + }); +} + +/** Acquire, inspect, and either roll back immediately or hold an idle fence. */ +export function prepareGatewaySuspend(params: { + requestId: string; + pauseScheduling: () => void; + resumeScheduling: () => void; + inspect?: Partial; + nowMs?: () => number; + createSuspensionId?: () => string; + warn?: (message: string) => void; +}): GatewaySuspendPrepareResult { + const nowMs = (params.nowMs ?? Date.now)(); + const current = COORDINATOR_STATE.current; + if (current?.kind === "recovering") { + return schedulerRecoveryResult(); + } + const existing = current ? normalizeExpiredHeldSuspension(current) : null; + if (existing?.kind === "recovering") { + return schedulerRecoveryResult(); + } + if (existing) { + if (existing.requestId !== params.requestId) { + return { status: "conflict", expiresAtMs: existing.expiresAtMs }; + } + existing.nowMs = params.nowMs ?? Date.now; + renewHeldSuspension(existing, nowMs); + return { + status: "ready", + suspensionId: existing.suspensionId, + expiresAtMs: existing.expiresAtMs, + activeCount: existing.snapshot.counts.totalActive, + blockers: existing.snapshot.blockers, + }; + } + + const owner = {}; + let suspensionInvalidated = false; + const admission = tryBeginGatewaySuspendAdmission(() => { + suspensionInvalidated = true; + const activeEntry = COORDINATOR_STATE.current; + if (activeEntry?.owner !== owner) { + return; + } + clearEntryTimer(activeEntry); + COORDINATOR_STATE.current = null; + }); + if (!admission) { + const snapshot = createGatewayActiveWorkSnapshot(params.inspect); + return { + status: "busy", + reason: "gateway-draining", + retryAfterMs: GATEWAY_SUSPEND_RETRY_AFTER_MS, + activeCount: snapshot.counts.totalActive, + blockers: snapshot.blockers, + }; + } + + let schedulingPaused = false; + let admissionCommitted = false; + try { + params.pauseScheduling(); + schedulingPaused = true; + const snapshot = createGatewayActiveWorkSnapshot(params.inspect); + if (!snapshot.idle) { + const resumed = resumeSchedulingBeforeReopen({ + owner, + resumeScheduling: params.resumeScheduling, + reopenAdmission: admission.rollback, + isInvalidated: () => suspensionInvalidated, + warn: params.warn, + }); + schedulingPaused = false; + if (!resumed) { + return schedulerRecoveryResult(); + } + return { + status: "busy", + reason: "active-work", + retryAfterMs: GATEWAY_SUSPEND_RETRY_AFTER_MS, + activeCount: snapshot.counts.totalActive, + blockers: snapshot.blockers, + }; + } + if (!admission.commit()) { + throw new Error("gateway suspension admission changed during preparation"); + } + admissionCommitted = true; + const suspensionId = (params.createSuspensionId ?? randomUUID)(); + const expiresAtMs = nowMs + GATEWAY_SUSPEND_TTL_MS; + const held = armExpiry({ + owner, + requestId: params.requestId, + suspensionId, + expiresAtMs, + snapshot, + reopenAdmission: admission.release, + resumeScheduling: params.resumeScheduling, + nowMs: params.nowMs ?? Date.now, + warn: params.warn, + }); + COORDINATOR_STATE.current = held; + return { + status: "ready", + suspensionId, + expiresAtMs, + activeCount: snapshot.counts.totalActive, + blockers: snapshot.blockers, + }; + } catch (err) { + if (schedulingPaused) { + const resumed = resumeSchedulingBeforeReopen({ + owner, + resumeScheduling: params.resumeScheduling, + reopenAdmission: admissionCommitted ? admission.release : admission.rollback, + isInvalidated: () => suspensionInvalidated, + warn: params.warn, + }); + if (!resumed) { + return schedulerRecoveryResult(); + } + } else if (admissionCommitted) { + admission.release(); + } else { + admission.rollback(); + } + throw err; + } +} + +export function getGatewaySuspendStatus(suspensionId: string): GatewaySuspendStatusResult { + const current = COORDINATOR_STATE.current; + if (current?.kind === "recovering") { + return schedulerRecoveryResult(); + } + const held = current ? normalizeExpiredHeldSuspension(current) : null; + if (held?.kind === "recovering") { + return schedulerRecoveryResult(); + } + if (!held) { + return { status: "running" }; + } + if (held.suspensionId !== suspensionId) { + return { status: "conflict", expiresAtMs: held.expiresAtMs }; + } + return { status: "ready", expiresAtMs: held.expiresAtMs }; +} + +export function resumeGatewaySuspend(suspensionId: string): GatewaySuspendResumeResult { + const current = COORDINATOR_STATE.current; + if (current?.kind === "recovering") { + return { + ok: false, + reason: "scheduler-resume-failed", + retryAfterMs: GATEWAY_SCHEDULER_RECOVERY_RETRY_MS, + }; + } + const held = current ? normalizeExpiredHeldSuspension(current) : null; + if (held?.kind === "recovering") { + return { + ok: false, + reason: "scheduler-resume-failed", + retryAfterMs: GATEWAY_SCHEDULER_RECOVERY_RETRY_MS, + }; + } + if (!held) { + return { + ok: true, + status: "running", + resumed: false, + }; + } + if (held.suspensionId !== suspensionId) { + return { ok: false, reason: "suspension-mismatch" }; + } + if (!resumeAndReopen(held)) { + return { + ok: false, + reason: "scheduler-resume-failed", + retryAfterMs: GATEWAY_SCHEDULER_RECOVERY_RETRY_MS, + }; + } + return { + ok: true, + status: "running", + resumed: true, + }; +} + +export function resetGatewaySuspendCoordinatorForTest(): void { + const current = COORDINATOR_STATE.current; + if (current) { + clearEntryTimer(current); + try { + current.resumeScheduling(); + } catch (err) { + current.warn?.(`gateway scheduler resume failed during test reset: ${String(err)}`); + } + current.reopenAdmission(); + COORDINATOR_STATE.current = null; + } +} diff --git a/src/infra/heartbeat-wake.test.ts b/src/infra/heartbeat-wake.test.ts index 0330c5364ddd..122828f59cab 100644 --- a/src/infra/heartbeat-wake.test.ts +++ b/src/infra/heartbeat-wake.test.ts @@ -1,5 +1,10 @@ // Exercises heartbeat wake coalescing, retries, and skip handling. import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { + getActiveGatewayRootWorkCount, + resetGatewayWorkAdmission, + tryBeginGatewaySuspendAdmission, +} from "../process/gateway-work-admission.js"; import { HEARTBEAT_SKIP_CRON_IN_PROGRESS, HEARTBEAT_SKIP_LANES_BUSY, @@ -71,15 +76,69 @@ describe("heartbeat-wake", () => { } beforeEach(() => { + resetGatewayWorkAdmission(); resetHeartbeatWakeStateForTests(); }); afterEach(() => { resetHeartbeatWakeStateForTests(); + resetGatewayWorkAdmission(); vi.useRealTimers(); vi.restoreAllMocks(); }); + it("defers a full wake while gateway suspension is prepared", async () => { + vi.useFakeTimers(); + const activeRootCounts: number[] = []; + const handler = vi.fn(async () => { + activeRootCounts.push(getActiveGatewayRootWorkCount()); + return { status: "ran" as const, durationMs: 1 }; + }); + setHeartbeatWakeHandler(handler); + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + + requestHeartbeat(wake("interval", { coalesceMs: 0 })); + await vi.advanceTimersByTimeAsync(1); + + expect(handler).not.toHaveBeenCalled(); + expect(getActiveGatewayRootWorkCount()).toBe(0); + + expect(suspension?.release()).toBe(true); + await vi.advanceTimersByTimeAsync(0); + + expect(handler).toHaveBeenCalledOnce(); + expect(activeRootCounts).toEqual([1]); + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); + + it("counts an in-flight wake until the whole handler settles", async () => { + vi.useFakeTimers(); + let finishWake: (() => void) | undefined; + const wakeFinished = new Promise((resolve) => { + finishWake = resolve; + }); + const handler = vi.fn(async () => { + await wakeFinished; + return { status: "ran" as const, durationMs: 1 }; + }); + setHeartbeatWakeHandler(handler); + + requestHeartbeat(wake("manual", { coalesceMs: 0 })); + await vi.advanceTimersByTimeAsync(1); + + expect(handler).toHaveBeenCalledOnce(); + expect(getActiveGatewayRootWorkCount()).toBe(1); + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension).not.toBeNull(); + expect(getActiveGatewayRootWorkCount()).toBe(1); + expect(suspension?.rollback()).toBe(true); + + finishWake?.(); + await vi.advanceTimersByTimeAsync(0); + expect(getActiveGatewayRootWorkCount()).toBe(0); + }); + it("coalesces multiple wake requests into one run", async () => { vi.useFakeTimers(); const handler = vi.fn().mockResolvedValue({ status: "skipped", reason: "disabled" }); diff --git a/src/infra/heartbeat-wake.ts b/src/infra/heartbeat-wake.ts index 4f7189b69c68..a6c8e29b14e2 100644 --- a/src/infra/heartbeat-wake.ts +++ b/src/infra/heartbeat-wake.ts @@ -1,5 +1,6 @@ // Tracks heartbeat wake requests, busy skips, and retry timing. import { normalizeOptionalString } from "@openclaw/normalization-core/string-coerce"; +import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js"; import { resolveTimerTimeoutMs } from "../shared/number-coercion.js"; import { normalizeHeartbeatWakeReason } from "./heartbeat-reason.js"; @@ -235,7 +236,11 @@ function schedule(coalesceMs: number, kind: WakeTimerKind = "normal") { ...(pendingWake.sessionKey ? { sessionKey: pendingWake.sessionKey } : {}), ...(pendingWake.heartbeat ? { heartbeat: pendingWake.heartbeat } : {}), }; - const res = await active(wakeOpts); + // Each wake is detached process work: admit the whole handler before + // it can mutate sessions or commitments, and keep it visible until done. + const res = await runWithGatewayIndependentRootWorkAdmission(async () => + active(wakeOpts), + ); if (res.status === "skipped" && isRetryableHeartbeatBusySkipReason(res.reason)) { // The target runtime is busy; retry this wake target soon. queuePendingWakeReason({ diff --git a/src/infra/infra-runtime.test.ts b/src/infra/infra-runtime.test.ts index 7b96418755a5..0ee8641a94e0 100644 --- a/src/infra/infra-runtime.test.ts +++ b/src/infra/infra-runtime.test.ts @@ -6,6 +6,10 @@ import { clearRuntimeConfigSnapshot, setRuntimeConfigSnapshot, } from "../config/config.js"; +import { + isGatewayWorkAdmissionClosed, + tryBeginGatewayRootWorkAdmission, +} from "../process/gateway-work-admission.js"; import { makeNetworkInterfacesSnapshot } from "../test-helpers/network-interfaces.js"; import { testing, @@ -136,6 +140,26 @@ describe("infra runtime", () => { await vi.runAllTimersAsync(); }); + it("holds root admission from scheduled emission until the signal is handled", async () => { + const handler = () => {}; + process.on("SIGUSR1", handler); + try { + scheduleGatewaySigusr1Restart({ delayMs: 0 }); + await vi.advanceTimersByTimeAsync(0); + + expect(isGatewayWorkAdmissionClosed()).toBe(true); + expect(tryBeginGatewayRootWorkAdmission()).toBeNull(); + + markGatewaySigusr1RestartHandled(); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + const root = tryBeginGatewayRootWorkAdmission(); + expect(root).not.toBeNull(); + root?.release(); + } finally { + process.removeListener("SIGUSR1", handler); + } + }); + it("backs off before an emoji that crosses the restart reason limit", () => { const restart = scheduleGatewaySigusr1Restart({ delayMs: 0, @@ -632,6 +656,7 @@ describe("infra runtime", () => { expect(beforeEmit).toHaveBeenCalledTimes(1); expect(afterEmitRejected).toHaveBeenCalledTimes(1); + expect(isGatewayWorkAdmissionClosed()).toBe(false); }); it("still emits restart when preparation fails", async () => { diff --git a/src/infra/restart-coordinator.ts b/src/infra/restart-coordinator.ts index d0af52f42872..d477406cc39d 100644 --- a/src/infra/restart-coordinator.ts +++ b/src/infra/restart-coordinator.ts @@ -1,13 +1,8 @@ -// Coordinates restart requests around active embedded agent runs. -import { getActiveEmbeddedRunCount } from "../agents/embedded-agent-runner/run-state.js"; -import { getTotalPendingReplies } from "../auto-reply/reply/dispatcher-registry.js"; -import { getActiveCronJobCount } from "../cron/active-jobs.js"; -import { getTotalQueueSize } from "../process/command-queue.js"; -import { getInspectableActiveTaskRestartBlockers } from "../tasks/task-registry.maintenance.js"; import { - type ActiveTaskRestartBlocker, - formatActiveTaskRestartBlocker, -} from "../tasks/task-restart-blocker.js"; + createGatewayActiveWorkSnapshot, + type GatewayActiveWorkBlocker, + type GatewayActiveWorkInspectors, +} from "./gateway-active-work.js"; import { scheduleGatewaySigusr1Restart, type ScheduledRestart } from "./restart.js"; // Safe restart coordination checks active local work before scheduling SIGUSR1 @@ -20,14 +15,20 @@ export type SafeGatewayRestartCounts = { activeTasks: number; totalActive: number; }; - -export type SafeGatewayRestartBlocker = { +export type SafeGatewayRestartBlocker = Omit & { kind: "queue" | "reply" | "embedded-run" | "cron-run" | "task"; - count: number; - message: string; - task?: ActiveTaskRestartBlocker; }; +type SafeRestartInspectors = Pick< + GatewayActiveWorkInspectors, + | "getQueueSize" + | "getPendingReplies" + | "getEmbeddedRuns" + | "getCronRuns" + | "getActiveTasks" + | "getTaskBlockers" +>; + export type SafeGatewayRestartPreflight = { safe: boolean; counts: SafeGatewayRestartCounts; @@ -42,105 +43,33 @@ export type SafeGatewayRestartRequestResult = { restart: ScheduledRestart; }; -type SafeRestartInspectors = { - getQueueSize: () => number; - getPendingReplies: () => number; - getEmbeddedRuns: () => number; - getCronRuns: () => number; - getActiveTasks: () => number; - getTaskBlockers: () => ActiveTaskRestartBlocker[]; -}; - -const defaultInspectors: SafeRestartInspectors = { - getQueueSize: getTotalQueueSize, - getPendingReplies: getTotalPendingReplies, - getEmbeddedRuns: getActiveEmbeddedRunCount, - getCronRuns: getActiveCronJobCount, - getActiveTasks: () => getInspectableActiveTaskRestartBlockers().length, - getTaskBlockers: getInspectableActiveTaskRestartBlockers, -}; - -function normalizeCount(value: number): number { - return Number.isFinite(value) && value > 0 ? Math.floor(value) : 0; -} - -function createFallbackTaskBlocker(count: number): SafeGatewayRestartBlocker { - return { - kind: "task", - count, - message: `${count} active background task run(s)`, - }; -} - export function createSafeGatewayRestartPreflight( inspectors: Partial = {}, ): SafeGatewayRestartPreflight { - const resolved = { ...defaultInspectors, ...inspectors }; + const snapshot = createGatewayActiveWorkSnapshot({ + ...inspectors, + getRootRequests: () => 0, + getSessionAdmissions: () => 0, + getSessionMutations: () => 0, + getChatRuns: () => 0, + getQueuedTurns: () => 0, + getTerminalPersistence: () => 0, + getTerminalSessions: () => 0, + }); const counts: SafeGatewayRestartCounts = { - queueSize: normalizeCount(resolved.getQueueSize()), - pendingReplies: normalizeCount(resolved.getPendingReplies()), - embeddedRuns: normalizeCount(resolved.getEmbeddedRuns()), - cronRuns: normalizeCount(resolved.getCronRuns()), - activeTasks: normalizeCount(resolved.getActiveTasks()), - totalActive: 0, + queueSize: snapshot.counts.queueSize, + pendingReplies: snapshot.counts.pendingReplies, + embeddedRuns: snapshot.counts.embeddedRuns, + cronRuns: snapshot.counts.cronRuns, + activeTasks: snapshot.counts.activeTasks, + totalActive: + snapshot.counts.queueSize + + snapshot.counts.pendingReplies + + snapshot.counts.embeddedRuns + + snapshot.counts.cronRuns + + snapshot.counts.activeTasks, }; - counts.totalActive = - counts.queueSize + - counts.pendingReplies + - counts.embeddedRuns + - counts.cronRuns + - counts.activeTasks; - - const blockers: SafeGatewayRestartBlocker[] = []; - if (counts.queueSize > 0) { - blockers.push({ - kind: "queue", - count: counts.queueSize, - message: `${counts.queueSize} queued or active operation(s)`, - }); - } - if (counts.pendingReplies > 0) { - blockers.push({ - kind: "reply", - count: counts.pendingReplies, - message: `${counts.pendingReplies} pending reply delivery operation(s)`, - }); - } - if (counts.embeddedRuns > 0) { - blockers.push({ - kind: "embedded-run", - count: counts.embeddedRuns, - message: `${counts.embeddedRuns} active embedded run(s)`, - }); - } - if (counts.cronRuns > 0) { - blockers.push({ - kind: "cron-run", - count: counts.cronRuns, - message: `${counts.cronRuns} active cron run(s)`, - }); - } - if (counts.activeTasks > 0) { - const taskBlockers = resolved.getTaskBlockers(); - if (taskBlockers.length === 0) { - blockers.push(createFallbackTaskBlocker(counts.activeTasks)); - } else { - // Cap task details so restart diagnostics stay bounded even during a - // backlog; counts still preserve the total active-task signal. - for (const task of taskBlockers.slice(0, 8)) { - blockers.push({ - kind: "task", - count: 1, - message: formatActiveTaskRestartBlocker(task), - task, - }); - } - const omitted = counts.activeTasks - taskBlockers.length; - if (omitted > 0) { - blockers.push(createFallbackTaskBlocker(omitted)); - } - } - } + const blockers = snapshot.blockers as SafeGatewayRestartBlocker[]; const summary = blockers.length === 0 diff --git a/src/infra/restart-suspension.test.ts b/src/infra/restart-suspension.test.ts new file mode 100644 index 000000000000..ee8012a198ec --- /dev/null +++ b/src/infra/restart-suspension.test.ts @@ -0,0 +1,225 @@ +// Pins scheduled restart ordering against the reversible host-suspension fence. +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { + getActiveGatewayRootWorkCount, + isGatewayWorkAdmissionClosed, + resetGatewayWorkAdmission, +} from "../process/gateway-work-admission.js"; +import type { GatewayActiveWorkInspectors } from "./gateway-active-work.js"; +import { + prepareGatewaySuspend, + resetGatewaySuspendCoordinatorForTest, + resumeGatewaySuspend, +} from "./gateway-suspend-coordinator.js"; +import { + isGatewaySigusr1RestartExternallyAllowed, + scheduleGatewaySigusr1Restart, + setGatewaySigusr1RestartPolicy, + setPreRestartDeferralCheck, + testing, +} from "./restart.js"; + +function inspectors(): GatewayActiveWorkInspectors { + return { + getQueueSize: () => 0, + getPendingReplies: () => 0, + getEmbeddedRuns: () => 0, + getCronRuns: () => 0, + getActiveTasks: () => 0, + getTaskBlockers: () => [], + getRootRequests: () => getActiveGatewayRootWorkCount(), + getSessionAdmissions: () => 0, + getSessionMutations: () => 0, + getChatRuns: () => 0, + getQueuedTurns: () => 0, + getTerminalPersistence: () => 0, + getTerminalSessions: () => 0, + }; +} + +function countSigusr1Emits(calls: readonly unknown[][]): number { + return calls.filter((args) => args[0] === "SIGUSR1").length; +} + +describe("scheduled restart during gateway suspension", () => { + const sigusr1Handler = () => {}; + + beforeEach(() => { + testing.resetSigusr1State(); + resetGatewaySuspendCoordinatorForTest(); + resetGatewayWorkAdmission(); + vi.useFakeTimers(); + process.on("SIGUSR1", sigusr1Handler); + }); + + afterEach(() => { + process.removeListener("SIGUSR1", sigusr1Handler); + testing.resetSigusr1State(); + resetGatewaySuspendCoordinatorForTest(); + resetGatewayWorkAdmission(); + vi.useRealTimers(); + vi.restoreAllMocks(); + }); + + it("defers a previously scheduled restart until a ready suspension resumes", async () => { + const emitSpy = vi.spyOn(process, "emit"); + scheduleGatewaySigusr1Restart({ + delayMs: 1_000, + reason: "config.patch", + skipCooldown: true, + }); + + const prepared = prepareGatewaySuspend({ + requestId: "request-restart-delay", + pauseScheduling: vi.fn(), + resumeScheduling: vi.fn(), + inspect: inspectors(), + createSuspensionId: () => "suspension-restart-delay", + }); + expect(prepared.status).toBe("ready"); + + await vi.advanceTimersByTimeAsync(1_000); + expect(countSigusr1Emits(emitSpy.mock.calls)).toBe(0); + + expect(resumeGatewaySuspend("suspension-restart-delay")).toMatchObject({ + ok: true, + resumed: true, + }); + await vi.advanceTimersByTimeAsync(0); + expect(countSigusr1Emits(emitSpy.mock.calls)).toBe(1); + }); + + it("reports active work while a due restart is preparing to emit", async () => { + const emitSpy = vi.spyOn(process, "emit"); + let releasePreparation: () => void = () => {}; + const preparation = new Promise((resolve) => { + releasePreparation = resolve; + }); + scheduleGatewaySigusr1Restart({ + delayMs: 0, + reason: "config.patch", + skipCooldown: true, + emitHooks: { + beforeEmit: async () => preparation, + }, + }); + await vi.advanceTimersByTimeAsync(0); + + const prepared = prepareGatewaySuspend({ + requestId: "request-restart-preparing", + pauseScheduling: vi.fn(), + resumeScheduling: vi.fn(), + inspect: inspectors(), + }); + expect(prepared).toMatchObject({ + status: "busy", + reason: "active-work", + activeCount: 1, + }); + expect(countSigusr1Emits(emitSpy.mock.calls)).toBe(0); + + releasePreparation(); + await vi.advanceTimersByTimeAsync(0); + expect(countSigusr1Emits(emitSpy.mock.calls)).toBe(1); + + expect( + prepareGatewaySuspend({ + requestId: "request-after-restart-signal", + pauseScheduling: vi.fn(), + resumeScheduling: vi.fn(), + inspect: inspectors(), + }), + ).toMatchObject({ + status: "busy", + reason: "gateway-draining", + }); + }); + + it("resets transient restart state without dropping live runtime bindings", async () => { + const emitSpy = vi.spyOn(process, "emit"); + const preRestartCheck = vi.fn(() => 0); + setPreRestartDeferralCheck(preRestartCheck); + setGatewaySigusr1RestartPolicy({ allowExternal: true }); + + scheduleGatewaySigusr1Restart({ delayMs: 0, skipCooldown: true }); + await vi.advanceTimersByTimeAsync(0); + expect(countSigusr1Emits(emitSpy.mock.calls)).toBe(1); + expect(preRestartCheck).toHaveBeenCalledOnce(); + expect(isGatewayWorkAdmissionClosed()).toBe(true); + + testing.resetSigusr1TransientState(); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + expect(isGatewaySigusr1RestartExternallyAllowed()).toBe(true); + + scheduleGatewaySigusr1Restart({ delayMs: 0, skipCooldown: true }); + await vi.advanceTimersByTimeAsync(0); + expect(countSigusr1Emits(emitSpy.mock.calls)).toBe(2); + expect(preRestartCheck).toHaveBeenCalledTimes(2); + }); + + it("cancels delayed restart work during a transient reset", async () => { + const emitSpy = vi.spyOn(process, "emit"); + scheduleGatewaySigusr1Restart({ delayMs: 1_000, skipCooldown: true }); + + testing.resetSigusr1TransientState(); + await vi.advanceTimersByTimeAsync(1_000); + + expect(countSigusr1Emits(emitSpy.mock.calls)).toBe(0); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + }); + + it("cancels a due restart waiting behind a prepared suspension", async () => { + const emitSpy = vi.spyOn(process, "emit"); + expect( + prepareGatewaySuspend({ + requestId: "request-reset-waiting-restart", + pauseScheduling: vi.fn(), + resumeScheduling: vi.fn(), + inspect: inspectors(), + }), + ).toMatchObject({ status: "ready" }); + scheduleGatewaySigusr1Restart({ delayMs: 0, skipCooldown: true }); + await vi.advanceTimersByTimeAsync(0); + expect(countSigusr1Emits(emitSpy.mock.calls)).toBe(0); + + resetGatewaySuspendCoordinatorForTest(); + testing.resetSigusr1TransientState(); + resetGatewayWorkAdmission(); + await vi.advanceTimersByTimeAsync(0); + + expect(countSigusr1Emits(emitSpy.mock.calls)).toBe(0); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + }); + + it("rejects prepared hooks that finish after a transient reset", async () => { + const emitSpy = vi.spyOn(process, "emit"); + const preparationStarted = vi.fn(); + const afterEmitRejected = vi.fn(); + let releasePreparation = () => {}; + const preparation = new Promise((resolve) => { + releasePreparation = resolve; + }); + scheduleGatewaySigusr1Restart({ + delayMs: 0, + skipCooldown: true, + emitHooks: { + beforeEmit: async () => { + preparationStarted(); + await preparation; + }, + afterEmitRejected, + }, + }); + await vi.advanceTimersByTimeAsync(0); + expect(preparationStarted).toHaveBeenCalledOnce(); + + testing.resetSigusr1TransientState(); + resetGatewayWorkAdmission(); + releasePreparation(); + await vi.advanceTimersByTimeAsync(0); + + expect(afterEmitRejected).toHaveBeenCalledOnce(); + expect(countSigusr1Emits(emitSpy.mock.calls)).toBe(0); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + }); +}); diff --git a/src/infra/restart.ts b/src/infra/restart.ts index ac95c3b72174..2576226c16bd 100644 --- a/src/infra/restart.ts +++ b/src/infra/restart.ts @@ -9,6 +9,12 @@ import { resolveGatewaySystemdServiceName, } from "../daemon/constants.js"; import { createSubsystemLogger } from "../logging/subsystem.js"; +import { + beginGatewayRestartSignalAdmission, + isGatewayRestartDraining, + runWithGatewayIndependentRootWorkAdmission, + type GatewayRestartSignalAdmissionLease, +} from "../process/gateway-work-admission.js"; import { resolveTimerTimeoutMs } from "../shared/number-coercion.js"; import type { DB as OpenClawStateKyselyDatabase } from "../state/openclaw-state-db.generated.js"; import { @@ -58,6 +64,8 @@ let pendingRestartEmitHooks: RestartEmitHooks | undefined; let pendingRestartSessionKey: string | undefined; let pendingRestartSkipDeferral = false; let pendingRestartPreparing = false; +let pendingRestartSignalAdmission: GatewayRestartSignalAdmissionLease | null = null; +let restartTransientGeneration = 0; const activeDeferralPolls = new Set>(); function shouldPreferRestartReason(next?: string, current?: string): boolean { @@ -82,6 +90,17 @@ function clearPendingScheduledRestart(): void { pendingRestartPreparing = false; } +function clearPendingRestartSignalAdmission(): boolean { + const rolledBack = pendingRestartSignalAdmission?.rollback() ?? false; + pendingRestartSignalAdmission = null; + return rolledBack; +} + +/** Releases a signal fence when the run loop rejects or fails to handle the signal. */ +export function rollbackGatewayRestartSignalAdmission(): boolean { + return clearPendingRestartSignalAdmission(); +} + function armPendingRestartTimer(requestedDueAt: number, nowMs: number): void { pendingRestartTimer = setTimeout( () => { @@ -122,6 +141,7 @@ function clearActiveDeferralPolls(): void { export function resetGatewayRestartStateForInProcessRestart(): void { clearActiveDeferralPolls(); clearPendingScheduledRestart(); + clearPendingRestartSignalAdmission(); // Cancel any in-progress deferred channel reload so it doesn't race with // the restart to start the same channel (e.g. telegram double-spawn). void import("../gateway/server-reload-handlers.js") @@ -358,8 +378,8 @@ export function setPreRestartDeferralCheck(fn: () => number): void { /** * Emit an authorized SIGUSR1 gateway restart, guarded against duplicate emissions. * Returns true if SIGUSR1 was emitted, false if a restart was already emitted. - * Both scheduleGatewaySigusr1Restart and the config watcher should use this - * to ensure only one restart fires. + * Runtime callers use emitGatewayRestartWithSignalAdmission so the signal-to-drain + * handoff stays fenced; this lower-level primitive remains available to tests. */ export function emitGatewayRestart( reasonOverride?: string, @@ -407,6 +427,26 @@ export function emitGatewayRestart( return true; } +/** + * Emits while holding the signal-to-drain admission fence. + * + * The caller must already own root-work admission. Scheduled restarts use the + * independent-root wrapper below; config reloads run inside their reload root. + */ +export function emitGatewayRestartWithSignalAdmission( + reasonOverride?: string, + intent?: GatewayRestartIntent, +): boolean { + const signalAdmission = pendingRestartSignalAdmission ?? beginGatewayRestartSignalAdmission(); + pendingRestartSignalAdmission = signalAdmission; + const hadUnconsumedRestartSignal = hasUnconsumedRestartSignal(); + const emitted = emitGatewayRestart(reasonOverride, intent); + if (!emitted && !hadUnconsumedRestartSignal) { + clearPendingRestartSignalAdmission(); + } + return emitted; +} + function resetSigusr1AuthorizationIfExpired(now = Date.now()) { if (sigusr1AuthorizedCount <= 0) { return; @@ -475,6 +515,10 @@ export function markGatewaySigusr1RestartHandled(): void { emittedRestartReason = undefined; emittedRestartIntent = undefined; } + // Accepted handlers first promote the fence to one-way restart drain, so + // this rollback becomes a no-op there. Rejected or test-only handlers must + // reopen admission or the next restart/root would wait forever. + clearPendingRestartSignalAdmission(); } function rollBackGatewayRestartEmission(): void { @@ -537,11 +581,21 @@ function updatePendingRestartEmitHooks( return true; } -async function emitPreparedGatewayRestart( +async function rejectPreparedRestartHook(hooks: RestartEmitHooks | undefined): Promise { + try { + await hooks?.afterEmitRejected?.(); + } catch {} +} + +async function emitPreparedGatewayRestartUnderAdmission( hooks?: RestartEmitHooks, reasonOverride?: string, intent?: GatewayRestartIntent, + transientGeneration = restartTransientGeneration, ): Promise { + if (transientGeneration !== restartTransientGeneration) { + return; + } let nextHooks = hooks ?? pendingRestartEmitHooks; // Keep pendingRestartSessionKey alive across the await beforeEmit() window: // a different-session caller that coalesces while preparation runs would @@ -553,8 +607,11 @@ async function emitPreparedGatewayRestart( let preparedHooks: RestartEmitHooks | undefined; while (nextHooks) { if (preparedHooks) { - await preparedHooks.afterEmitRejected?.().catch(() => undefined); + await rejectPreparedRestartHook(preparedHooks); preparedHooks = undefined; + if (transientGeneration !== restartTransientGeneration) { + return; + } } try { await nextHooks.beforeEmit?.(); @@ -564,6 +621,10 @@ async function emitPreparedGatewayRestart( `restart preparation failed; restart will continue without it: ${String(err)}`, ); } + if (transientGeneration !== restartTransientGeneration) { + await rejectPreparedRestartHook(preparedHooks); + return; + } if (hooks) { break; } @@ -579,12 +640,40 @@ async function emitPreparedGatewayRestart( const preferredReason = shouldPreferRestartReason(pendingRestartReason, reasonOverride) ? pendingRestartReason : undefined; - const emitted = emitGatewayRestart( + const emitted = emitGatewayRestartWithSignalAdmission( preferredReason ?? reasonOverride, preferredReason && intent ? { ...intent, reason: preferredReason } : intent, ); if (!emitted) { - await preparedHooks?.afterEmitRejected?.().catch(() => undefined); + await rejectPreparedRestartHook(preparedHooks); + } +} + +async function emitPreparedGatewayRestart( + hooks?: RestartEmitHooks, + reasonOverride?: string, + intent?: GatewayRestartIntent, +): Promise { + const transientGeneration = restartTransientGeneration; + try { + // A delayed restart can become due after host suspension prepared. Independent + // root admission makes the transition atomic: due restarts block preparation, + // while a prepared suspension defers emission until it resumes. + await runWithGatewayIndependentRootWorkAdmission(async () => { + if (transientGeneration !== restartTransientGeneration) { + return; + } + await emitPreparedGatewayRestartUnderAdmission( + hooks, + reasonOverride, + intent, + transientGeneration, + ); + }); + } catch (err) { + if (!isGatewayRestartDraining()) { + throw err; + } } } @@ -1005,20 +1094,27 @@ export function scheduleGatewaySigusr1Restart(opts?: { }; } +function resetSigusr1TransientStateForTest(): void { + restartTransientGeneration += 1; + sigusr1AuthorizedCount = 0; + sigusr1AuthorizedUntil = 0; + restartCycleToken = 0; + emittedRestartToken = 0; + consumedRestartToken = 0; + emittedRestartReason = undefined; + emittedRestartIntent = undefined; + lastRestartEmittedAt = 0; + clearActiveDeferralPolls(); + clearPendingScheduledRestart(); + clearPendingRestartSignalAdmission(); +} + export const testing = { + resetSigusr1TransientState: resetSigusr1TransientStateForTest, resetSigusr1State() { - sigusr1AuthorizedCount = 0; - sigusr1AuthorizedUntil = 0; + resetSigusr1TransientStateForTest(); sigusr1ExternalAllowed = false; preRestartCheck = null; - restartCycleToken = 0; - emittedRestartToken = 0; - consumedRestartToken = 0; - emittedRestartReason = undefined; - emittedRestartIntent = undefined; - lastRestartEmittedAt = 0; - clearActiveDeferralPolls(); - clearPendingScheduledRestart(); }, }; export { testing as __testing }; diff --git a/src/plugin-sdk/gateway-runtime.ts b/src/plugin-sdk/gateway-runtime.ts index 10b4e84f80e8..712f305da1b7 100644 --- a/src/plugin-sdk/gateway-runtime.ts +++ b/src/plugin-sdk/gateway-runtime.ts @@ -42,6 +42,6 @@ export { createOperatorApprovalsGatewayClient, withOperatorApprovalsGatewayClient, } from "../gateway/operator-approvals-client.js"; -export { ErrorCodes, errorShape } from "../../packages/gateway-protocol/src/index.js"; -export type { EventFrame } from "../../packages/gateway-protocol/src/index.js"; +export { ErrorCodes, errorShape } from "../../packages/gateway-protocol/src/schema/error-codes.js"; +export type { EventFrame } from "../../packages/gateway-protocol/src/schema/frames.js"; export type { GatewayRequestHandlerOptions } from "../gateway/server-methods/types.js"; diff --git a/src/process/command-queue.test.ts b/src/process/command-queue.test.ts index 59565891c881..bebe8b99b4a0 100644 --- a/src/process/command-queue.test.ts +++ b/src/process/command-queue.test.ts @@ -2,6 +2,10 @@ import { MAX_TIMER_TIMEOUT_MS } from "@openclaw/normalization-core/number-coercion"; import { importFreshModule } from "openclaw/plugin-sdk/test-fixtures"; import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; +import { + tryBeginGatewayRootWorkAdmission, + tryBeginGatewaySuspendAdmission, +} from "./gateway-work-admission.js"; import { CommandLane } from "./lanes.js"; const diagnosticMocks = vi.hoisted(() => ({ @@ -852,6 +856,62 @@ describe("command queue", () => { await expect(task).resolves.toBe("ok"); }); + it("reversibly fences new enqueues without disturbing an active task", async () => { + const { task, release } = enqueueBlockedMainTask(async () => "active-finished"); + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + await expect( + enqueueCommandInLane(CommandLane.Main, async () => "blocked"), + ).rejects.toBeInstanceOf(GatewayDrainingError); + + release(); + await expect(task).resolves.toBe("active-finished"); + expect(suspension?.release()).toBe(true); + await expect(enqueueCommandInLane(CommandLane.Main, async () => "resumed")).resolves.toBe( + "resumed", + ); + }); + + it("lets an admitted root enqueue while suspension preparation refuses new work", async () => { + const continueRoot = createDeferred(); + const root = tryBeginGatewayRootWorkAdmission(); + expect(root).not.toBeNull(); + const result = root?.run(async () => { + await continueRoot.promise; + return await enqueueCommandInLane(CommandLane.Main, async () => "continued"); + }); + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + + try { + continueRoot.resolve(); + await expect(result).resolves.toBe("continued"); + await expect( + enqueueCommandInLane(CommandLane.Main, async () => "blocked"), + ).rejects.toBeInstanceOf(GatewayDrainingError); + } finally { + suspension?.rollback(); + root?.release(); + } + }); + + it("rejects subordinate enqueues from an admitted root after restart drain", async () => { + const continueRoot = createDeferred(); + const root = tryBeginGatewayRootWorkAdmission(); + expect(root).not.toBeNull(); + const result = root?.run(async () => { + await continueRoot.promise; + return await enqueueCommandInLane(CommandLane.Main, async () => "blocked"); + }); + + try { + markGatewayDraining(); + continueRoot.resolve(); + await expect(result).rejects.toBeInstanceOf(GatewayDrainingError); + } finally { + root?.release(); + } + }); + it("resetAllLanes clears gateway draining flag and re-allows enqueue", async () => { markGatewayDraining(); resetAllLanes(); diff --git a/src/process/command-queue.ts b/src/process/command-queue.ts index 63ce5a091452..0524a3fe0825 100644 --- a/src/process/command-queue.ts +++ b/src/process/command-queue.ts @@ -7,6 +7,14 @@ import { import { resolveGlobalSingleton } from "../shared/global-singleton.js"; import { clampPositiveTimerTimeoutMs } from "../shared/number-coercion.js"; import type { CommandQueueEnqueueOptions } from "./command-queue.types.js"; +import { + GatewayDrainingError, + isGatewaySubordinateWorkAdmissionClosed, + isGatewayWorkAdmissionClosed, + markGatewayRestartDraining, + resetGatewayWorkAdmission, +} from "./gateway-work-admission.js"; +export { GatewayDrainingError } from "./gateway-work-admission.js"; import { CommandLane } from "./lanes.js"; /** * Dedicated error type thrown when a queued command is rejected because @@ -42,17 +50,6 @@ export function isCommandLaneTaskTimeoutError(err: unknown, lane?: string): bool return lane === undefined || err.message.includes(`Command lane "${lane}" task timed out`); } -/** - * Dedicated error type thrown when a new command is rejected because the - * gateway is currently draining for restart. - */ -export class GatewayDrainingError extends Error { - constructor() { - super("Gateway is draining for restart; new tasks are not accepted"); - this.name = "GatewayDrainingError"; - } -} - // Minimal in-process queue to serialize command executions. // Default lane ("main") preserves the existing behavior. Additional lanes allow // low-risk parallelism (e.g. cron jobs) without interleaving stdin / logs for @@ -122,7 +119,6 @@ const COMMAND_QUEUE_STATE_KEY = Symbol.for("openclaw.commandQueueState"); function getQueueState() { const state = resolveGlobalSingleton(COMMAND_QUEUE_STATE_KEY, () => ({ - gatewayDraining: false, lanes: new Map(), activeTaskWaiters: new Set(), nextTaskId: 1, @@ -464,11 +460,11 @@ function drainLane(lane: string) { * `GatewayDrainingError` instead of being silently killed on shutdown. */ export function markGatewayDraining(): void { - getQueueState().gatewayDraining = true; + markGatewayRestartDraining(); } export function isGatewayDraining(): boolean { - return getQueueState().gatewayDraining; + return isGatewayWorkAdmissionClosed(); } export function setCommandLaneConcurrency(lane: string, maxConcurrent: number) { @@ -488,7 +484,7 @@ export function enqueueCommandInLane( opts?: CommandQueueEnqueueOptions, ): Promise { const queueState = getQueueState(); - if (queueState.gatewayDraining) { + if (isGatewaySubordinateWorkAdmissionClosed()) { return Promise.reject(new GatewayDrainingError()); } const cleaned = normalizeLane(lane); @@ -608,7 +604,7 @@ export function resetCommandLane(lane: string = CommandLane.Main): number { */ export function resetCommandQueueStateForTest(): void { const queueState = getQueueState(); - queueState.gatewayDraining = false; + resetGatewayWorkAdmission(); queueState.lanes.clear(); for (const waiter of Array.from(queueState.activeTaskWaiters)) { resolveActiveTaskWaiter(waiter, { drained: true }); @@ -633,7 +629,7 @@ export function resetCommandQueueStateForTest(): void { */ export function resetAllLanes(): void { const queueState = getQueueState(); - queueState.gatewayDraining = false; + resetGatewayWorkAdmission(); const lanesToDrain: string[] = []; for (const state of queueState.lanes.values()) { state.generation += 1; diff --git a/src/process/gateway-work-admission.test.ts b/src/process/gateway-work-admission.test.ts new file mode 100644 index 000000000000..0e8d19eb0573 --- /dev/null +++ b/src/process/gateway-work-admission.test.ts @@ -0,0 +1,233 @@ +// Covers root work counting and reversible suspension admission transitions. +import { afterEach, beforeEach, expect, it, vi } from "vitest"; +import { + beginGatewayRestartSignalAdmission, + GatewayDrainingError, + getActiveGatewayRootWorkCount, + isGatewaySubordinateWorkAdmissionClosed, + isGatewayWorkAdmissionClosed, + markGatewayRestartDraining, + retainGatewayRootWorkAdmissionContinuation, + resetGatewayWorkAdmission, + runWithGatewayIndependentRootWorkContinuation, + runWithGatewayRootWorkAdmission, + tryBeginGatewayRootWorkAdmission, + tryBeginGatewaySuspendAdmission, +} from "./gateway-work-admission.js"; + +beforeEach(resetGatewayWorkAdmission); +afterEach(resetGatewayWorkAdmission); + +it("counts one nested root chain once and excludes the preparing caller", async () => { + const outer = tryBeginGatewayRootWorkAdmission(); + expect(outer).not.toBeNull(); + expect(outer?.ownsRoot).toBe(true); + await outer?.run(async () => { + expect(getActiveGatewayRootWorkCount()).toBe(1); + expect(getActiveGatewayRootWorkCount({ excludeCurrent: true })).toBe(0); + const nested = tryBeginGatewayRootWorkAdmission(); + expect(nested).not.toBeNull(); + expect(nested?.ownsRoot).toBe(false); + expect(getActiveGatewayRootWorkCount()).toBe(1); + nested?.release(); + }); + outer?.release(); + expect(getActiveGatewayRootWorkCount()).toBe(0); +}); + +it("rolls back or releases a generation-bound suspension without resetting roots", () => { + const invalidated = vi.fn(); + const preparing = tryBeginGatewaySuspendAdmission(invalidated); + expect(preparing).not.toBeNull(); + expect(isGatewayWorkAdmissionClosed()).toBe(true); + expect(tryBeginGatewayRootWorkAdmission()).toBeNull(); + expect(preparing?.rollback()).toBe(true); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + + const prepared = tryBeginGatewaySuspendAdmission(invalidated); + expect(prepared?.commit()).toBe(true); + expect(prepared?.release()).toBe(true); + expect(prepared?.release()).toBe(false); + expect(invalidated).not.toHaveBeenCalled(); + expect(isGatewayWorkAdmissionClosed()).toBe(false); +}); + +it("lets an admitted root cross only the reversible suspension fence", async () => { + const root = tryBeginGatewayRootWorkAdmission(); + expect(root).not.toBeNull(); + await root?.run(async () => { + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(isGatewaySubordinateWorkAdmissionClosed()).toBe(false); + expect(suspension?.rollback()).toBe(true); + + markGatewayRestartDraining(); + expect(isGatewaySubordinateWorkAdmissionClosed()).toBe(true); + }); + root?.release(); +}); + +it("synchronously reserves a tracked continuation across a closed suspension fence", async () => { + const root = tryBeginGatewayRootWorkAdmission(); + expect(root).not.toBeNull(); + let releaseContinuation = () => {}; + let continuation: Promise | undefined; + await root?.run(async () => { + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension).not.toBeNull(); + continuation = runWithGatewayIndependentRootWorkContinuation( + async () => + await new Promise((resolve) => { + releaseContinuation = resolve; + }), + ); + expect(getActiveGatewayRootWorkCount()).toBe(2); + expect(suspension?.rollback()).toBe(true); + }); + + root?.release(); + expect(getActiveGatewayRootWorkCount()).toBe(1); + releaseContinuation(); + await continuation; + expect(getActiveGatewayRootWorkCount()).toBe(0); +}); + +it("retains an admitted request root across its handler return", async () => { + const root = tryBeginGatewayRootWorkAdmission(); + expect(root).not.toBeNull(); + let continueChild = () => {}; + let releaseContinuation = () => {}; + let subordinateAdmissionClosed: boolean | undefined; + let child: Promise | undefined; + const childGate = new Promise((resolve) => { + continueChild = resolve; + }); + + await root?.run(async () => { + const retainedRelease = retainGatewayRootWorkAdmissionContinuation(); + expect(retainedRelease).not.toBeNull(); + releaseContinuation = retainedRelease ?? (() => {}); + child = (async () => { + await childGate; + subordinateAdmissionClosed = isGatewaySubordinateWorkAdmissionClosed(); + })(); + }); + + root?.release(); + expect(getActiveGatewayRootWorkCount()).toBe(1); + continueChild(); + await child; + expect(subordinateAdmissionClosed).toBe(false); + releaseContinuation(); + releaseContinuation(); + expect(getActiveGatewayRootWorkCount()).toBe(0); +}); + +it("runs an admitted continuation when restart drain wins the handoff race", async () => { + const root = tryBeginGatewayRootWorkAdmission(); + expect(root).not.toBeNull(); + const ran = vi.fn(); + await root?.run(async () => { + markGatewayRestartDraining(); + await runWithGatewayIndependentRootWorkContinuation(async () => { + ran(); + }); + }); + root?.release(); + + expect(ran).toHaveBeenCalledOnce(); + expect(getActiveGatewayRootWorkCount()).toBe(0); +}); + +it("does not admit an unrelated continuation through restart drain", async () => { + markGatewayRestartDraining(); + const ran = vi.fn(); + + await expect( + runWithGatewayIndependentRootWorkContinuation(async () => { + ran(); + }), + ).rejects.toThrow("gateway is draining for restart"); + expect(ran).not.toHaveBeenCalled(); +}); + +it("does not let a stale suspension release clear restart drain", () => { + const invalidated = vi.fn(); + const suspension = tryBeginGatewaySuspendAdmission(invalidated); + expect(suspension?.commit()).toBe(true); + + markGatewayRestartDraining(); + + expect(invalidated).toHaveBeenCalledOnce(); + expect(suspension?.release()).toBe(false); + expect(isGatewayWorkAdmissionClosed()).toBe(true); +}); + +it("blocks suspension while restart signal handling is pending", () => { + const pendingSignal = beginGatewayRestartSignalAdmission(); + + expect(isGatewayWorkAdmissionClosed()).toBe(true); + expect(tryBeginGatewayRootWorkAdmission()).toBeNull(); + expect(tryBeginGatewaySuspendAdmission(() => {})).toBeNull(); + expect(pendingSignal.rollback()).toBe(true); + expect(isGatewayWorkAdmissionClosed()).toBe(false); + expect(tryBeginGatewaySuspendAdmission(() => {})?.rollback()).toBe(true); +}); + +it("promotes a pending restart signal to one-way drain", () => { + const pendingSignal = beginGatewayRestartSignalAdmission(); + + markGatewayRestartDraining(); + + expect(pendingSignal.rollback()).toBe(false); + expect(isGatewayWorkAdmissionClosed()).toBe(true); + expect(tryBeginGatewayRootWorkAdmission()).toBeNull(); +}); + +it("defers required internal root work until suspension reopens", async () => { + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + const entered = vi.fn(); + const pending = runWithGatewayRootWorkAdmission(async () => { + entered(); + expect(getActiveGatewayRootWorkCount()).toBe(1); + }); + + await Promise.resolve(); + expect(entered).not.toHaveBeenCalled(); + suspension?.release(); + await pending; + + expect(entered).toHaveBeenCalledOnce(); + expect(getActiveGatewayRootWorkCount()).toBe(0); +}); + +it("retires surviving root records across an in-process reset", async () => { + const root = tryBeginGatewayRootWorkAdmission(); + expect(root).not.toBeNull(); + await root?.run(async () => { + resetGatewayWorkAdmission(); + expect(getActiveGatewayRootWorkCount()).toBe(0); + expect(isGatewaySubordinateWorkAdmissionClosed()).toBe(true); + const nested = tryBeginGatewayRootWorkAdmission(); + expect(nested).not.toBeNull(); + expect(nested?.ownsRoot).toBe(true); + await nested?.run(async () => { + expect(getActiveGatewayRootWorkCount()).toBe(1); + expect(isGatewaySubordinateWorkAdmissionClosed()).toBe(false); + }); + nested?.release(); + expect(isGatewaySubordinateWorkAdmissionClosed()).toBe(true); + }); + root?.release(); + expect(getActiveGatewayRootWorkCount()).toBe(0); +}); + +it("does not wake deferred internal work into a restart drain", async () => { + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + const pending = runWithGatewayRootWorkAdmission(async () => {}); + + markGatewayRestartDraining(); + + await expect(pending).rejects.toBeInstanceOf(GatewayDrainingError); +}); diff --git a/src/process/gateway-work-admission.ts b/src/process/gateway-work-admission.ts new file mode 100644 index 000000000000..4e8400648297 --- /dev/null +++ b/src/process/gateway-work-admission.ts @@ -0,0 +1,361 @@ +// Coordinates process-wide root work admission with reversible host suspension. +import { AsyncLocalStorage } from "node:async_hooks"; +import { resolveGlobalSingleton } from "../shared/global-singleton.js"; + +export type GatewaySuspendAdmissionPhase = "accepting" | "preparing" | "prepared"; + +export class GatewayDrainingError extends Error { + constructor() { + super("Gateway is draining; new tasks are not accepted"); + this.name = "GatewayDrainingError"; + } +} + +type GatewayRootWorkAdmission = { + references: number; + released: boolean; +}; + +type GatewayWorkAdmissionState = { + restartDraining: boolean; + restartSignalPending: boolean; + restartSignalGeneration: number; + suspendPhase: GatewaySuspendAdmissionPhase; + suspendGeneration: number; + suspendInvalidated?: () => void; + activeRootWork: Set; + currentRootWork: AsyncLocalStorage; + suspendOpenWaiters: Set<() => void>; +}; + +const GATEWAY_WORK_ADMISSION_STATE = resolveGlobalSingleton( + Symbol.for("openclaw.gatewayWorkAdmissionState"), + (): GatewayWorkAdmissionState => ({ + restartDraining: false, + restartSignalPending: false, + restartSignalGeneration: 0, + suspendPhase: "accepting", + suspendGeneration: 0, + activeRootWork: new Set(), + currentRootWork: new AsyncLocalStorage(), + suspendOpenWaiters: new Set(), + }), +); + +export type GatewayRootWorkAdmissionLease = { + ownsRoot: boolean; + release: () => void; + run: (run: () => Promise) => Promise; +}; + +export type GatewaySuspendAdmissionLease = { + commit: () => boolean; + rollback: () => boolean; + release: () => boolean; +}; + +export type GatewayRestartSignalAdmissionLease = { + rollback: () => boolean; +}; + +function createGatewayRootWorkAdmission(): GatewayRootWorkAdmissionLease { + const admission: GatewayRootWorkAdmission = { references: 1, released: false }; + GATEWAY_WORK_ADMISSION_STATE.activeRootWork.add(admission); + const release = createGatewayRootWorkRelease(admission); + return { + ownsRoot: true, + release, + run: async (run: () => Promise) => + await GATEWAY_WORK_ADMISSION_STATE.currentRootWork.run(admission, run), + }; +} + +function createGatewayRootWorkRelease(admission: GatewayRootWorkAdmission): () => void { + let leaseReleased = false; + return () => { + if (leaseReleased || admission.released) { + return; + } + leaseReleased = true; + admission.references -= 1; + if (admission.references > 0) { + return; + } + admission.released = true; + GATEWAY_WORK_ADMISSION_STATE.activeRootWork.delete(admission); + }; +} + +function invalidateSuspendAdmission(): void { + const callback = GATEWAY_WORK_ADMISSION_STATE.suspendInvalidated; + GATEWAY_WORK_ADMISSION_STATE.suspendInvalidated = undefined; + GATEWAY_WORK_ADMISSION_STATE.suspendPhase = "accepting"; + GATEWAY_WORK_ADMISSION_STATE.suspendGeneration += 1; + resolveSuspendOpenWaiters(); + callback?.(); +} + +function resolveSuspendOpenWaiters(): void { + const waiters = Array.from(GATEWAY_WORK_ADMISSION_STATE.suspendOpenWaiters); + GATEWAY_WORK_ADMISSION_STATE.suspendOpenWaiters.clear(); + for (const resolve of waiters) { + resolve(); + } +} + +/** True while restart signal/drain or host suspension rejects new process work. */ +export function isGatewayWorkAdmissionClosed(): boolean { + return ( + GATEWAY_WORK_ADMISSION_STATE.restartDraining || + GATEWAY_WORK_ADMISSION_STATE.restartSignalPending || + GATEWAY_WORK_ADMISSION_STATE.suspendPhase !== "accepting" + ); +} + +/** Existing admitted roots may finish spawning subordinate command/session work. + * New async chains still see the global fence, preserving refuse-only suspension. */ +export function isGatewaySubordinateWorkAdmissionClosed(): boolean { + if ( + GATEWAY_WORK_ADMISSION_STATE.restartDraining || + GATEWAY_WORK_ADMISSION_STATE.restartSignalPending + ) { + return true; + } + const current = GATEWAY_WORK_ADMISSION_STATE.currentRootWork.getStore(); + if (current) { + // Reset/release retires inherited ALS descendants. They must explicitly + // re-enter admission instead of spawning untracked subordinate work. + return current.released; + } + return GATEWAY_WORK_ADMISSION_STATE.suspendPhase !== "accepting"; +} + +export function getGatewaySuspendAdmissionPhase(): GatewaySuspendAdmissionPhase { + return GATEWAY_WORK_ADMISSION_STATE.suspendPhase; +} + +export function isGatewayRestartDraining(): boolean { + return ( + GATEWAY_WORK_ADMISSION_STATE.restartDraining || + GATEWAY_WORK_ADMISSION_STATE.restartSignalPending + ); +} + +/** Restart drain is one-way until the in-process restart resets runtime state. */ +export function markGatewayRestartDraining(): void { + GATEWAY_WORK_ADMISSION_STATE.restartSignalPending = false; + GATEWAY_WORK_ADMISSION_STATE.restartSignalGeneration += 1; + GATEWAY_WORK_ADMISSION_STATE.restartDraining = true; + resolveSuspendOpenWaiters(); + if (GATEWAY_WORK_ADMISSION_STATE.suspendPhase !== "accepting") { + // A restart supersedes a reversible suspension. The coordinator callback + // drops its timer/token without reopening the scheduler being shut down. + invalidateSuspendAdmission(); + } +} + +/** Blocks suspension across signal emission until the run loop starts restart drain. */ +export function beginGatewayRestartSignalAdmission(): GatewayRestartSignalAdmissionLease { + if (GATEWAY_WORK_ADMISSION_STATE.restartSignalPending) { + return { rollback: () => false }; + } + GATEWAY_WORK_ADMISSION_STATE.restartSignalPending = true; + const generation = ++GATEWAY_WORK_ADMISSION_STATE.restartSignalGeneration; + return { + rollback: () => { + if ( + !GATEWAY_WORK_ADMISSION_STATE.restartSignalPending || + GATEWAY_WORK_ADMISSION_STATE.restartSignalGeneration !== generation + ) { + return false; + } + GATEWAY_WORK_ADMISSION_STATE.restartSignalPending = false; + GATEWAY_WORK_ADMISSION_STATE.restartSignalGeneration += 1; + resolveSuspendOpenWaiters(); + return true; + }, + }; +} + +/** Root RPC/timer admission. Nested work in the same async chain counts once. */ +export function tryBeginGatewayRootWorkAdmission(): GatewayRootWorkAdmissionLease | null { + const current = GATEWAY_WORK_ADMISSION_STATE.currentRootWork.getStore(); + if (current && !current.released) { + return { + ownsRoot: false, + release: () => {}, + run: async (run: () => Promise) => await run(), + }; + } + // Existing request chains use the ALS path above; new roots stop for either + // restart drain or host suspension. + if ( + GATEWAY_WORK_ADMISSION_STATE.restartDraining || + GATEWAY_WORK_ADMISSION_STATE.restartSignalPending || + GATEWAY_WORK_ADMISSION_STATE.suspendPhase !== "accepting" + ) { + return null; + } + return createGatewayRootWorkAdmission(); +} + +/** Independent detached work counts separately even when launched by an admitted parent. */ +export function tryBeginGatewayIndependentRootWorkAdmission(): GatewayRootWorkAdmissionLease | null { + if ( + GATEWAY_WORK_ADMISSION_STATE.restartDraining || + GATEWAY_WORK_ADMISSION_STATE.restartSignalPending || + GATEWAY_WORK_ADMISSION_STATE.suspendPhase !== "accepting" + ) { + return null; + } + return createGatewayRootWorkAdmission(); +} + +/** Waits through a prepared lease, then joins the root-work set atomically. */ +export async function beginGatewayRootWorkAdmissionWhenOpen(): Promise { + while (true) { + if (GATEWAY_WORK_ADMISSION_STATE.restartDraining) { + throw new GatewayDrainingError(); + } + const admission = tryBeginGatewayRootWorkAdmission(); + if (admission) { + return admission; + } + await new Promise((resolve) => { + GATEWAY_WORK_ADMISSION_STATE.suspendOpenWaiters.add(resolve); + }); + } +} + +export async function runWithGatewayRootWorkAdmission(run: () => Promise): Promise { + const admission = await beginGatewayRootWorkAdmissionWhenOpen(); + try { + return await admission.run(run); + } finally { + admission.release(); + } +} + +export async function runWithGatewayIndependentRootWorkAdmission( + run: () => Promise, +): Promise { + while (true) { + if (GATEWAY_WORK_ADMISSION_STATE.restartDraining) { + throw new Error("gateway is draining for restart"); + } + const admission = tryBeginGatewayIndependentRootWorkAdmission(); + if (admission) { + try { + return await admission.run(run); + } finally { + admission.release(); + } + } + await new Promise((resolve) => { + GATEWAY_WORK_ADMISSION_STATE.suspendOpenWaiters.add(resolve); + }); + } +} + +/** + * Detaches required follow-up from the current admitted transaction. + * A live parent synchronously reserves a tracked root even after restart or + * suspension closes admission; callers without a live parent use the normal + * independent-root fence. + */ +export function runWithGatewayIndependentRootWorkContinuation( + run: () => Promise, +): Promise { + const parent = GATEWAY_WORK_ADMISSION_STATE.currentRootWork.getStore(); + if (!parent || parent.released) { + return runWithGatewayIndependentRootWorkAdmission(run); + } + const admission = createGatewayRootWorkAdmission(); + return admission.run(run).finally(admission.release); +} + +/** Transfers an admitted request root to work that intentionally outlives its handler. */ +export function retainGatewayRootWorkAdmissionContinuation(): (() => void) | null { + const current = GATEWAY_WORK_ADMISSION_STATE.currentRootWork.getStore(); + if (!current || current.released) { + return null; + } + current.references += 1; + return createGatewayRootWorkRelease(current); +} + +/** Active root requests/ticks, optionally excluding the caller running prepare. */ +export function getActiveGatewayRootWorkCount(opts?: { excludeCurrent?: boolean }): number { + let count = GATEWAY_WORK_ADMISSION_STATE.activeRootWork.size; + const current = GATEWAY_WORK_ADMISSION_STATE.currentRootWork.getStore(); + if ( + opts?.excludeCurrent === true && + current && + !current.released && + GATEWAY_WORK_ADMISSION_STATE.activeRootWork.has(current) + ) { + count -= 1; + } + return Math.max(0, count); +} + +/** Atomically closes new suspension admission before synchronous inspection. */ +export function tryBeginGatewaySuspendAdmission( + onInvalidated: () => void, +): GatewaySuspendAdmissionLease | null { + if ( + GATEWAY_WORK_ADMISSION_STATE.restartDraining || + GATEWAY_WORK_ADMISSION_STATE.restartSignalPending || + GATEWAY_WORK_ADMISSION_STATE.suspendPhase !== "accepting" + ) { + return null; + } + GATEWAY_WORK_ADMISSION_STATE.suspendPhase = "preparing"; + const generation = ++GATEWAY_WORK_ADMISSION_STATE.suspendGeneration; + GATEWAY_WORK_ADMISSION_STATE.suspendInvalidated = onInvalidated; + + const transition = ( + expected: GatewaySuspendAdmissionPhase, + next: GatewaySuspendAdmissionPhase, + ): boolean => { + if ( + GATEWAY_WORK_ADMISSION_STATE.suspendGeneration !== generation || + GATEWAY_WORK_ADMISSION_STATE.suspendPhase !== expected + ) { + return false; + } + GATEWAY_WORK_ADMISSION_STATE.suspendPhase = next; + if (next === "accepting") { + GATEWAY_WORK_ADMISSION_STATE.suspendInvalidated = undefined; + resolveSuspendOpenWaiters(); + } + return true; + }; + + return { + commit: () => transition("preparing", "prepared"), + rollback: () => transition("preparing", "accepting"), + release: () => transition("prepared", "accepting"), + }; +} + +/** Clears restart/suspend admission during SIGUSR1 and isolated tests. */ +export function resetGatewayWorkAdmission(): void { + // SIGUSR1 can abandon old async chains before their finally blocks run. + // Retire their ALS records so surviving chains must re-enter admission. + for (const admission of GATEWAY_WORK_ADMISSION_STATE.activeRootWork) { + admission.references = 0; + admission.released = true; + } + GATEWAY_WORK_ADMISSION_STATE.activeRootWork.clear(); + GATEWAY_WORK_ADMISSION_STATE.restartDraining = false; + GATEWAY_WORK_ADMISSION_STATE.restartSignalPending = false; + GATEWAY_WORK_ADMISSION_STATE.restartSignalGeneration += 1; + if (GATEWAY_WORK_ADMISSION_STATE.suspendPhase !== "accepting") { + invalidateSuspendAdmission(); + } else { + GATEWAY_WORK_ADMISSION_STATE.suspendGeneration += 1; + GATEWAY_WORK_ADMISSION_STATE.suspendInvalidated = undefined; + } + resolveSuspendOpenWaiters(); +} diff --git a/src/sessions/session-lifecycle-admission.test.ts b/src/sessions/session-lifecycle-admission.test.ts index de245bfbd560..c2df08c989ad 100644 --- a/src/sessions/session-lifecycle-admission.test.ts +++ b/src/sessions/session-lifecycle-admission.test.ts @@ -1,8 +1,15 @@ // Tests lifecycle/work admission ordering across canonical keys and backing ids. import { importFreshModule } from "openclaw/plugin-sdk/test-fixtures"; import { expect, it } from "vitest"; +import { + resetGatewayWorkAdmission, + tryBeginGatewayRootWorkAdmission, + tryBeginGatewaySuspendAdmission, +} from "../process/gateway-work-admission.js"; import { beginSessionWorkAdmission, + getActiveSessionLifecycleMutationCount, + getActiveSessionWorkAdmissionCount, hasOnlySessionLifecycleMutationKindActive, interruptSessionWorkAdmissions, isSessionWorkAdmissionActive, @@ -17,6 +24,116 @@ function createDeferred() { return { promise, resolve }; } +it("counts one multi-identity admission once", async () => { + const admission = await beginSessionWorkAdmission({ + scope: "store-count", + identities: ["agent:main:child", "session-count"], + assertAllowed: () => {}, + }); + try { + expect(getActiveSessionWorkAdmissionCount()).toBe(1); + } finally { + admission.release(); + } + expect(getActiveSessionWorkAdmissionCount()).toBe(0); +}); + +it("counts one multi-identity lifecycle mutation once across module instances", async () => { + const first = await importFreshModule( + import.meta.url, + "./session-lifecycle-admission.js?scope=session-mutation-count-a", + ); + const second = await importFreshModule( + import.meta.url, + "./session-lifecycle-admission.js?scope=session-mutation-count-b", + ); + const mutationStarted = createDeferred(); + const releaseMutation = createDeferred(); + const mutation = first.runExclusiveSessionLifecycleMutation({ + scope: "store-mutation-count", + identities: ["agent:main:child", "session-mutation-count"], + run: async () => { + mutationStarted.resolve(); + await releaseMutation.promise; + }, + }); + await mutationStarted.promise; + + try { + expect(first.getActiveSessionLifecycleMutationCount()).toBe(1); + expect(second.getActiveSessionLifecycleMutationCount()).toBe(1); + } finally { + releaseMutation.resolve(); + await mutation; + } + expect(second.getActiveSessionLifecycleMutationCount()).toBe(0); +}); + +it("rejects an admission that resumes after suspension closes the async gap", async () => { + resetGatewayWorkAdmission(); + const mutationStarted = createDeferred(); + const releaseMutation = createDeferred(); + const mutation = runExclusiveSessionLifecycleMutation({ + scope: "store-suspend-race", + identities: ["session-suspend-race", "backing-suspend-race"], + run: async () => { + mutationStarted.resolve(); + await releaseMutation.promise; + }, + }); + await mutationStarted.promise; + expect(getActiveSessionLifecycleMutationCount()).toBeGreaterThan(0); + + const admission = beginSessionWorkAdmission({ + scope: "store-suspend-race", + identities: ["session-suspend-race", "backing-suspend-race"], + assertAllowed: () => {}, + }); + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + releaseMutation.resolve(); + await mutation; + expect(getActiveSessionLifecycleMutationCount()).toBe(0); + + await expect(admission).rejects.toMatchObject({ name: "GatewayDrainingError" }); + expect(getActiveSessionWorkAdmissionCount()).toBe(0); + suspension?.release(); + resetGatewayWorkAdmission(); +}); + +it("lets an admitted root enter session work while suspension preparation refuses new roots", async () => { + resetGatewayWorkAdmission(); + const continueRoot = createDeferred(); + const root = tryBeginGatewayRootWorkAdmission(); + expect(root).not.toBeNull(); + const active = root?.run(async () => { + await continueRoot.promise; + const admission = await beginSessionWorkAdmission({ + scope: "store-admitted-root", + identities: ["session-admitted-root"], + assertAllowed: () => {}, + }); + admission.release(); + }); + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + + try { + continueRoot.resolve(); + await expect(active).resolves.toBeUndefined(); + await expect( + beginSessionWorkAdmission({ + scope: "store-new-root", + identities: ["session-new-root"], + assertAllowed: () => {}, + }), + ).rejects.toMatchObject({ name: "GatewayDrainingError" }); + } finally { + suspension?.rollback(); + root?.release(); + resetGatewayWorkAdmission(); + } +}); + it("serializes lifecycle mutation and work admission across identity aliases", async () => { const mutationStarted = createDeferred(); const releaseMutation = createDeferred(); diff --git a/src/sessions/session-lifecycle-admission.ts b/src/sessions/session-lifecycle-admission.ts index 782ea6fe88c1..5ea67b9dcba0 100644 --- a/src/sessions/session-lifecycle-admission.ts +++ b/src/sessions/session-lifecycle-admission.ts @@ -1,5 +1,9 @@ // Serializes lifecycle mutations and work admission for logical session identities. import { AsyncLocalStorage } from "node:async_hooks"; +import { + GatewayDrainingError, + isGatewaySubordinateWorkAdmissionClosed, +} from "../process/gateway-work-admission.js"; import { resolveGlobalSingleton } from "../shared/global-singleton.js"; import { runQueuedStoreWrite, type StoreWriterQueue } from "../shared/store-writer-queue.js"; @@ -14,6 +18,7 @@ type SessionLifecycleAdmissionState = { mutationQueues: Map; activeAdmissions: Map>; activeMutations: Map; + activeMutationRuns?: Set; activeMutationKinds: Map>; idleWaiters: Map void>>; currentAdmissions: AsyncLocalStorage>; @@ -30,6 +35,7 @@ const SESSION_LIFECYCLE_ADMISSION_STATE = resolveGlobalSingleton( mutationQueues: new Map(), activeAdmissions: new Map(), activeMutations: new Map(), + activeMutationRuns: new Set(), activeMutationKinds: new Map(), idleWaiters: new Map(), currentAdmissions: new AsyncLocalStorage(), @@ -44,6 +50,9 @@ const { idleWaiters: SESSION_LIFECYCLE_IDLE_WAITERS, currentAdmissions: CURRENT_SESSION_WORK_ADMISSIONS, } = SESSION_LIFECYCLE_ADMISSION_STATE; +// Older runtime chunks can create the shared state without this newer index. +const ACTIVE_SESSION_LIFECYCLE_MUTATION_RUNS = + (SESSION_LIFECYCLE_ADMISSION_STATE.activeMutationRuns ??= new Set()); export type SessionWorkAdmissionLease = { release: () => void; @@ -208,6 +217,7 @@ export async function runExclusiveSessionLifecycleMutation(params: { const signal = params.signal; signal?.throwIfAborted(); const callerAdmissions = new Set(CURRENT_SESSION_WORK_ADMISSIONS.getStore()); + const mutationRun = {}; let mutationActivated = false; let removeAbortListener = () => {}; const mutation = runWithSessionMutationIdentityLocks( @@ -219,6 +229,7 @@ export async function runExclusiveSessionLifecycleMutation(params: { signal?.throwIfAborted(); mutationActivated = true; removeAbortListener(); + ACTIVE_SESSION_LIFECYCLE_MUTATION_RUNS.add(mutationRun); for (const identity of identities) { ACTIVE_SESSION_LIFECYCLE_MUTATIONS.set( identity, @@ -263,6 +274,7 @@ export async function runExclusiveSessionLifecycleMutation(params: { resolve(); } } + ACTIVE_SESSION_LIFECYCLE_MUTATION_RUNS.delete(mutationRun); }); } }), @@ -324,6 +336,26 @@ export function isSessionWorkAdmissionActive( ); } +/** Unique admitted turns; one lease can be indexed under several identities. */ +export function getActiveSessionWorkAdmissionCount(): number { + const admissions = new Set(); + for (const active of ACTIVE_SESSION_WORK_ADMISSIONS.values()) { + for (const admission of active) { + admissions.add(admission); + } + } + return admissions.size; +} + +/** Unique active lifecycle mutations; one run can be indexed under several identities. */ +export function getActiveSessionLifecycleMutationCount(): number { + if (ACTIVE_SESSION_LIFECYCLE_MUTATION_RUNS.size > 0) { + return ACTIVE_SESSION_LIFECYCLE_MUTATION_RUNS.size; + } + // A mutation from an older loaded chunk may only populate the identity index. + return ACTIVE_SESSION_LIFECYCLE_MUTATIONS.size > 0 ? 1 : 0; +} + export async function beginSessionWorkAdmission(params: { scope: string; identities: Iterable; @@ -331,6 +363,9 @@ export async function beginSessionWorkAdmission(params: { onInterrupt?: () => void; signal?: AbortSignal; }): Promise { + if (isGatewaySubordinateWorkAdmissionClosed()) { + throw new GatewayDrainingError(); + } const identities = normalizeSessionIdentities(params.scope, params.identities); return await runExclusiveSessionLifecycle({ scope: params.scope, @@ -338,6 +373,11 @@ export async function beginSessionWorkAdmission(params: { signal: params.signal, run: async () => { await params.assertAllowed(); + // assertAllowed can yield while a host suspension acquires its fence. + // Recheck immediately before registration to close that admission race. + if (isGatewaySubordinateWorkAdmissionClosed()) { + throw new GatewayDrainingError(); + } let resolveReleased = () => {}; const admission: SessionWorkAdmission = { interrupt: params.onInterrupt, diff --git a/src/tasks/cron-task-cancel.test.ts b/src/tasks/cron-task-cancel.test.ts index 28793c6a8034..04de78ce9462 100644 --- a/src/tasks/cron-task-cancel.test.ts +++ b/src/tasks/cron-task-cancel.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it, vi } from "vitest"; import { CRON_TASK_RUN_SETTLEMENT_TRACKING_MAX_MS, + getSuspensionVisibleCronTaskRunCount, resetActiveCronTaskRunsForTests, retireActiveCronTaskRunTracking, startActiveCronTaskRunSettlementGrace, @@ -9,14 +10,19 @@ import { } from "./cron-task-cancel.js"; describe("cron task cancellation tracking", () => { - it("retires never-settling cron promises at lifecycle cutoff", async () => { + it("retires restart tracking while keeping an unsettled core suspension-visible", async () => { resetActiveCronTaskRunsForTests(); - trackActiveCronTaskRunSettlement(new Promise(() => {})); + let settle = () => {}; + const core = new Promise((resolve) => { + settle = resolve; + }); + trackActiveCronTaskRunSettlement(core); await expect(waitForActiveCronTaskRuns(0)).resolves.toEqual({ drained: false, active: 1, }); + expect(getSuspensionVisibleCronTaskRunCount()).toBe(1); retireActiveCronTaskRunTracking(); @@ -24,6 +30,11 @@ describe("cron task cancellation tracking", () => { drained: true, active: 0, }); + expect(getSuspensionVisibleCronTaskRunCount()).toBe(1); + + settle(); + await core; + await vi.waitFor(() => expect(getSuspensionVisibleCronTaskRunCount()).toBe(0)); }); it("drops never-settling cron promises after a bounded grace period", async () => { @@ -51,9 +62,25 @@ describe("cron task cancellation tracking", () => { drained: true, active: 0, }); + expect(getSuspensionVisibleCronTaskRunCount()).toBe(1); } finally { vi.useRealTimers(); resetActiveCronTaskRunsForTests(); } }); + + it("keeps suspension blocked until a timed-out core actually settles", async () => { + resetActiveCronTaskRunsForTests(); + let settle = () => {}; + const core = new Promise((resolve) => { + settle = resolve; + }); + trackActiveCronTaskRunSettlement(core); + startActiveCronTaskRunSettlementGrace(); + + expect(getSuspensionVisibleCronTaskRunCount()).toBe(1); + settle(); + await core; + await vi.waitFor(() => expect(getSuspensionVisibleCronTaskRunCount()).toBe(0)); + }); }); diff --git a/src/tasks/cron-task-cancel.ts b/src/tasks/cron-task-cancel.ts index 944c00f511ee..2807459751f9 100644 --- a/src/tasks/cron-task-cancel.ts +++ b/src/tasks/cron-task-cancel.ts @@ -11,6 +11,9 @@ type SettlingCronTaskRun = { const activeCronTaskRunsByRunId = new Map(); const settlingCronTaskRuns = new Map, SettlingCronTaskRun>(); +// Restart drain may retire an abort-ignoring core after a bounded grace, but a +// host snapshot must keep refusing readiness until that core actually settles. +const suspensionVisibleCronTaskRuns = new Set>(); const DEFAULT_CRON_TASK_RUN_DRAIN_POLL_MS = 25; export const CRON_TASK_RUN_SETTLEMENT_TRACKING_MAX_MS = 60_000; @@ -65,6 +68,7 @@ export function abortActiveCronTaskRuns(reason = "Gateway restarting."): number export function trackActiveCronTaskRunSettlement(promise: Promise): void { settlingCronTaskRuns.set(promise, {}); + suspensionVisibleCronTaskRuns.add(promise); void promise .catch(() => undefined) .finally(() => { @@ -73,9 +77,16 @@ export function trackActiveCronTaskRunSettlement(promise: Promise): voi clearTimeout(entry.retirementTimer); } settlingCronTaskRuns.delete(promise); + suspensionVisibleCronTaskRuns.delete(promise); }); } +/** Cron cores that can still mutate state even after timeout/cancel returned. */ +export function getSuspensionVisibleCronTaskRunCount(): number { + return suspensionVisibleCronTaskRuns.size; +} + +/** Retires restart-drain bookkeeping without hiding still-running cores from suspension. */ export function retireActiveCronTaskRunTracking(): void { activeCronTaskRunsByRunId.clear(); for (const entry of settlingCronTaskRuns.values()) { @@ -129,4 +140,5 @@ export function cancelActiveCronTaskRun(params: { export function resetActiveCronTaskRunsForTests(): void { retireActiveCronTaskRunTracking(); + suspensionVisibleCronTaskRuns.clear(); } diff --git a/src/tasks/task-registry.maintenance.ts b/src/tasks/task-registry.maintenance.ts index f49601fe1516..854b12a9e313 100644 --- a/src/tasks/task-registry.maintenance.ts +++ b/src/tasks/task-registry.maintenance.ts @@ -30,6 +30,7 @@ import { isPluginStateDatabaseOpen, sweepExpiredPluginStateEntries, } from "../plugin-state/plugin-state-store.js"; +import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js"; import { parseAgentSessionKey } from "../routing/session-key.js"; import { deriveSessionChatTypeFromKey, @@ -1092,7 +1093,9 @@ function startScheduledSweep() { const clearSweepInProgress = () => { sweepInProgress = false; }; - sweepTaskRegistry().then(clearSweepInProgress, clearSweepInProgress); + void runWithGatewayIndependentRootWorkAdmission(async () => { + await sweepTaskRegistry(); + }).then(clearSweepInProgress, clearSweepInProgress); } export async function runTaskRegistryMaintenance(): Promise { diff --git a/src/tasks/task-registry.test.ts b/src/tasks/task-registry.test.ts index b78702ef48e9..682ff54db532 100644 --- a/src/tasks/task-registry.test.ts +++ b/src/tasks/task-registry.test.ts @@ -14,6 +14,11 @@ import { } from "../infra/heartbeat-wake.js"; import type { SessionBindingRecord } from "../infra/outbound/session-binding-service.js"; import { peekSystemEvents, resetSystemEventsForTest } from "../infra/system-events.js"; +import { + getActiveGatewayRootWorkCount, + resetGatewayWorkAdmission, + tryBeginGatewaySuspendAdmission, +} from "../process/gateway-work-admission.js"; import type { ParsedAgentSessionKey } from "../routing/session-key.js"; import { withTempDir } from "../test-helpers/temp-dir.js"; import { withEnvAsync } from "../test-utils/env.js"; @@ -145,6 +150,7 @@ function configureTaskRegistryMaintenanceRuntimeForTest(params: { listTaskRecords?: () => ReturnType[]; acpEntry?: AcpSessionStoreEntry; acpEntries?: AcpSessionStoreEntry[]; + listAcpSessionEntries?: () => Promise; hasActiveAcpTurn?: (sessionKey: string) => boolean; sessionBindings?: SessionBindingRecord[]; closeAcpSession?: (params: { @@ -167,7 +173,7 @@ function configureTaskRegistryMaintenanceRuntimeForTest(params: { storeReadFailed: false, } satisfies AcpSessionStoreEntry; setTaskRegistryMaintenanceRuntimeForTests({ - listAcpSessionEntries: async () => params.acpEntries ?? [], + listAcpSessionEntries: params.listAcpSessionEntries ?? (async () => params.acpEntries ?? []), readAcpSessionEntry: () => params.acpEntry ?? emptyAcpEntry, listSessionBindingsBySession: () => params.sessionBindings ?? [], closeAcpSession: params.closeAcpSession, @@ -458,6 +464,7 @@ function configureInMemoryTaskStoresForLinkValidationTests() { describe("task-registry", () => { beforeEach(() => { + resetGatewayWorkAdmission(); setTaskRegistryDeliveryRuntimeForTests({ sendMessage: hoisted.sendMessageMock, }); @@ -470,6 +477,7 @@ describe("task-registry", () => { }); afterEach(() => { + resetGatewayWorkAdmission(); vi.useRealTimers(); resetSystemEventsForTest(); resetHeartbeatWakeStateForTests(); @@ -1192,7 +1200,11 @@ describe("task-registry", () => { expect(linked?.parentFlowId).toBe(flow.flowId); let remainingUpsertFailures = 2; + const admittedRetryCounts: number[] = []; const upsertFlow = vi.fn(() => { + if (upsertFlow.mock.calls.length > 1) { + admittedRetryCounts.push(getActiveGatewayRootWorkCount()); + } if (remainingUpsertFailures > 0) { remainingUpsertFailures -= 1; throw new Error("SQLITE_FULL: database or disk is full"); @@ -1223,11 +1235,23 @@ describe("task-registry", () => { await vi.advanceTimersByTimeAsync(1_000); await flushAsyncWork(); expect(getTaskFlowById(flow.flowId)?.status).toBe("running"); + expect(upsertFlow).toHaveBeenCalledTimes(2); + expect(admittedRetryCounts).toEqual([1]); + + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); await vi.advanceTimersByTimeAsync(5_000); await flushAsyncWork(); + expect(upsertFlow).toHaveBeenCalledTimes(2); + expect(getActiveGatewayRootWorkCount()).toBe(0); + + expect(suspension?.release()).toBe(true); + await flushAsyncWork(); + expect(upsertFlow).toHaveBeenCalledTimes(3); + expect(admittedRetryCounts).toEqual([1, 1]); const retriedFlow = getTaskFlowById(flow.flowId); expect(retriedFlow?.status).toBe("succeeded"); expect(retriedFlow?.endedAt).toBe(200); @@ -2386,6 +2410,41 @@ describe("task-registry", () => { }); }); + it("keeps detached terminal delivery root-admitted through mirror persistence", async () => { + await withTaskRegistryTempDir(async () => { + resetTaskRegistryMemoryForTest(); + let releaseSend = () => {}; + hoisted.sendMessageMock.mockImplementationOnce( + () => + new Promise((resolve) => { + releaseSend = () => + resolve({ channel: "notifychat", to: "notifychat:123", via: "direct" }); + }), + ); + createTaskRecord({ + runtime: "acp", + ownerKey: "agent:main:main", + scopeKind: "session", + requesterOrigin: { channel: "notifychat", to: "notifychat:123" }, + childSessionKey: "agent:main:acp:child", + runId: "run-held-delivery", + task: "Deliver after completion", + status: "succeeded", + deliveryStatus: "pending", + terminalOutcome: "blocked", + terminalSummary: "Waiting for parent review.", + }); + + await vi.waitFor(() => expect(hoisted.sendMessageMock).toHaveBeenCalledOnce()); + expect(getActiveGatewayRootWorkCount()).toBe(1); + releaseSend(); + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + expectRecordFields(requireTaskByRunId("run-held-delivery"), { + deliveryStatus: "delivered", + }); + }); + }); + it("restores persisted tasks from disk on the next lookup", async () => { await withTaskRegistryTempDir( async () => { @@ -3228,6 +3287,30 @@ describe("task-registry", () => { }); }); + it("keeps scheduled maintenance root-admitted until session cleanup inspection settles", async () => { + await withTaskRegistryTempDir(async () => { + vi.useFakeTimers(); + resetTaskRegistryMemoryForTest(); + let releaseInspection = (_entries: AcpSessionStoreEntry[]) => {}; + const inspection = new Promise((resolve) => { + releaseInspection = resolve; + }); + configureTaskRegistryMaintenanceRuntimeForTest({ + currentTasks: new Map(), + snapshotTasks: [], + listAcpSessionEntries: async () => await inspection, + }); + + startTaskRegistryMaintenance(); + await vi.advanceTimersByTimeAsync(5_000); + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(1)); + + releaseInspection([]); + await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + stopTaskRegistryMaintenance(); + }); + }); + it("does not leak unhandled rejections when the scheduled maintenance sweep fails", async () => { await withTaskRegistryTempDir(async () => { vi.useFakeTimers(); diff --git a/src/tasks/task-registry.ts b/src/tasks/task-registry.ts index a54d10234063..c2fdd26664d6 100644 --- a/src/tasks/task-registry.ts +++ b/src/tasks/task-registry.ts @@ -14,6 +14,10 @@ import { formatErrorMessage } from "../infra/errors.js"; import { requestHeartbeat } from "../infra/heartbeat-wake.js"; import { enqueueSystemEvent } from "../infra/system-events.js"; import { createSubsystemLogger } from "../logging/subsystem.js"; +import { + isGatewayRestartDraining, + runWithGatewayIndependentRootWorkAdmission, +} from "../process/gateway-work-admission.js"; import { normalizeAgentId, parseAgentSessionKey } from "../routing/session-key.js"; import { createLazyPromiseLoader } from "../shared/lazy-runtime.js"; import { normalizeDeliveryContext } from "../utils/delivery-context.shared.js"; @@ -1147,24 +1151,36 @@ function scheduleTaskFlowSyncRetry(task: TaskRecord, operation: string, attempt } const retryTimer = setTimeout(() => { taskFlowSyncRetryTimers.delete(taskId); - const current = tasks.get(taskId); - if (!current) { - return; - } - const flowId = current.parentFlowId?.trim(); - if (!flowId || findLatestTaskForFlowId(flowId)?.taskId !== taskId) { - return; - } - const result = syncFlowFromTaskResult(current); - if (!result.ok) { - log.warn("Failed to retry parent flow sync from task", { + // A terminal task no longer blocks suspension, but its durable parent-flow + // projection still mutates state. Keep every delayed attempt visible and + // prevent it from crossing a prepared host snapshot boundary. + void runWithGatewayIndependentRootWorkAdmission(async () => { + const current = tasks.get(taskId); + if (!current) { + return; + } + const flowId = current.parentFlowId?.trim(); + if (!flowId || findLatestTaskForFlowId(flowId)?.taskId !== taskId) { + return; + } + const result = syncFlowFromTaskResult(current); + if (!result.ok) { + log.warn("Failed to retry parent flow sync from task", { + operation, + taskId, + flowId: current.parentFlowId, + reason: result.reason, + }); + scheduleTaskFlowSyncRetry(current, operation, attempt + 1); + } + }).catch((error: unknown) => { + log.warn("Failed to admit parent flow sync retry from task", { operation, taskId, - flowId: current.parentFlowId, - reason: result.reason, + flowId: task.parentFlowId, + error, }); - scheduleTaskFlowSyncRetry(current, operation, attempt + 1); - } + }); }, delayMs); retryTimer.unref?.(); taskFlowSyncRetryTimers.set(taskId, retryTimer); @@ -1398,6 +1414,36 @@ function queueBlockedTaskFollowup(task: TaskRecord) { } export async function maybeDeliverTaskTerminalUpdate(taskId: string): Promise { + return await runTaskDeliveryWithIndependentAdmission(taskId, async () => + maybeDeliverTaskTerminalUpdateUnderAdmission(taskId), + ); +} + +async function runTaskDeliveryWithIndependentAdmission( + taskId: string, + deliver: () => Promise, +): Promise { + let admitted = false; + try { + return await runWithGatewayIndependentRootWorkAdmission(async () => { + admitted = true; + return await deliver(); + }); + } catch (error) { + // Late lifecycle callbacks must not leak a rejected detached promise after + // restart closes admission. An already-admitted delivery still reports its + // own failures instead of hiding them behind a concurrent restart. + if (!admitted && isGatewayRestartDraining()) { + const current = tasks.get(taskId); + return current ? cloneTaskRecord(current) : null; + } + throw error; + } +} + +async function maybeDeliverTaskTerminalUpdateUnderAdmission( + taskId: string, +): Promise { ensureTaskRegistryReady(); const current = tasks.get(taskId); if (!current || !shouldAutoDeliverTaskTerminalUpdate(current)) { @@ -1548,6 +1594,15 @@ export async function maybeDeliverTaskTerminalUpdate(taskId: string): Promise { + return await runTaskDeliveryWithIndependentAdmission(taskId, async () => + maybeDeliverTaskStateChangeUpdateUnderAdmission(taskId, latestEvent), + ); +} + +async function maybeDeliverTaskStateChangeUpdateUnderAdmission( + taskId: string, + latestEvent?: TaskEventRecord, ): Promise { ensureTaskRegistryReady(); const current = tasks.get(taskId); diff --git a/src/utils/queue-helpers.ts b/src/utils/queue-helpers.ts index ab20274d9a53..ec07cd9632d8 100644 --- a/src/utils/queue-helpers.ts +++ b/src/utils/queue-helpers.ts @@ -169,10 +169,13 @@ export function applyQueueDropPolicy(params: { } /** Wait until the queue has been quiet for its debounce window. */ -export function waitForQueueDebounce(queue: { - debounceMs: number; - lastEnqueuedAt: number; -}): Promise { +export function waitForQueueDebounce( + queue: { + debounceMs: number; + lastEnqueuedAt: number; + }, + abortSignal?: AbortSignal, +): Promise { if (process.env.OPENCLAW_TEST_FAST === "1") { // Tests use this escape hatch so debounce logic does not slow deterministic queue specs. return Promise.resolve(); @@ -181,15 +184,36 @@ export function waitForQueueDebounce(queue: { if (debounceMs <= 0) { return Promise.resolve(); } + if (abortSignal?.aborted) { + return Promise.resolve(); + } return new Promise((resolve) => { - const check = () => { - const since = Date.now() - queue.lastEnqueuedAt; - if (since >= debounceMs) { - resolve(); + let settled = false; + let timer: ReturnType | undefined; + const finish = () => { + if (settled) { return; } - setTimeout(check, debounceMs - since); + settled = true; + if (timer !== undefined) { + clearTimeout(timer); + } + abortSignal?.removeEventListener("abort", finish); + resolve(); }; + const check = () => { + if (abortSignal?.aborted) { + finish(); + return; + } + const since = Date.now() - queue.lastEnqueuedAt; + if (since >= debounceMs) { + finish(); + return; + } + timer = setTimeout(check, debounceMs - since); + }; + abortSignal?.addEventListener("abort", finish, { once: true }); check(); }); }