From e98c7dfbcb9f0fbac2af28fc409b2eb0fbd415a2 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Mon, 13 Jul 2026 08:03:53 -0700 Subject: [PATCH] feat(cloud-workers): session placement, dispatch, and worker turn routing (#106332) * feat(gateway-protocol): add session placement schema Closed state discriminator for session execution placement (local/requested/provisioning/syncing/starting/active/draining/reconciling/reclaimed/failed), sessions.dispatch params, and worker-admission transcript/live cursor extensions. Swift protocol models mirror the schema. * feat(state): add worker session placement table worker_session_placements rows carry placement state, transition generation, worker ownership metadata, ACK cursors, and the turn claim columns used for atomic admission. * feat(cloud-workers): add durable placement state machine store SQLite-backed placement store split by concern: state table (placement-state), discriminated record types + shape invariants (placement-record), row codec + CAS transition values (placement-row-codec), atomic turn-claim admission/release/waiters (placement-turn-claims), and lifecycle CAS transitions (placement-store). * feat(cloud-workers): sync workspaces and attach sessions to worker environments Environment service session attachment + turn credentials, tunnel workspace commands over a dedicated SSH runner, and git/plain workspace sync into $HOME/.openclaw-worker/workspaces with an immutable manifest. Symlink escapes are rejected locally before transfer (macOS openrsync stat-fails them opaquely) and again by the remote manifest guard. * feat(worker): run one-shot embedded turns from launch descriptors Worker runtime executes a single embedded turn from a stdin launch descriptor and reports completed/failed/fenced on stdout for the gateway launcher. Terminal lifecycle live events are deferred past the final transcript flush; transcript projection helpers are shared via transcript-message instead of duplicated in the runtime. * feat(cloud-workers): dispatch placements and route worker turns Dispatch service drives local->requested->provisioning->syncing->starting->active with failure teardown (placement-dispatch-failure) and restart/runtime recovery incl. lost-worker reclaim (placement-dispatch-recovery). Worker turn launcher claims the placement turn atomically, builds a windowed launch descriptor (worker-turn-payload), runs the remote one-shot worker, and reconciles the committed transcript; agent runners route turns through the session placement admission provider. * feat(gateway): expose session placement RPCs and startup reconciliation sessions.dispatch RPC with lifecycle admission barriers, operator-facing placement projection on session listings, placement-aware session reset guard, and startup/interval reconciliation wiring for worker placements. --- .../openclaw/app/gateway/GatewayProtocol.kt | 1 + .../OpenClawProtocol/GatewayModels.swift | 588 ++++++++ docs/cli/worker.md | 6 +- docs/gateway/protocol.md | 3 +- packages/gateway-protocol/src/index.ts | 14 + packages/gateway-protocol/src/schema.ts | 1 + .../src/schema/protocol-schemas.ts | 2 + .../src/schema/session-placement.test.ts | 255 ++++ .../src/schema/session-placement.ts | 198 +++ .../src/schema/worker-admission.test.ts | 25 + .../src/schema/worker-admission.ts | 31 +- src/agents/command/attempt-execution.ts | 249 ++-- src/agents/embedded-agent-runner/run.ts | 39 +- .../session-placement-admission.test.ts | 191 +++ src/agents/session-placement-admission.ts | 82 ++ .../reply/agent-runner-execution.ts | 385 ++--- src/auto-reply/reply/followup-runner.ts | 355 ++--- src/cli/command-catalog.ts | 1 + src/cli/command-path-policy.test.ts | 1 + src/cli/command-startup-policy.test.ts | 1 + src/cron/isolated-agent/run-executor.ts | 111 +- src/gateway/method-scopes.test.ts | 1 + src/gateway/methods/core-descriptors.ts | 6 + src/gateway/server-methods-list.test.ts | 2 +- .../server-methods/sessions.dispatch.test.ts | 403 +++++ src/gateway/server-methods/sessions.ts | 344 ++++- src/gateway/server-methods/shared-types.ts | 10 +- src/gateway/server-request-context.ts | 8 + .../server-startup-post-attach.test.ts | 44 +- src/gateway/server-startup-post-attach.ts | 45 +- .../server-worker-environment-startup.ts | 26 +- .../server-worker-placement-startup.ts | 429 ++++++ src/gateway/server.impl.ts | 98 +- ...rver.sessions.placement-projection.test.ts | 124 ++ ...essions.worker-placement-lifecycle.test.ts | 246 ++++ .../message-handler.worker.test.ts | 3 + src/gateway/session-message-events.test.ts | 1 + src/gateway/session-reset-service.ts | 25 + src/gateway/session-utils.types.ts | 2 + .../worker-environments/admission.test.ts | 29 +- src/gateway/worker-environments/admission.ts | 4 + .../connection-identity.ts | 1 + .../inference-runtime.test.ts | 1 + .../inference-store.test.ts | 1 + .../worker-environments/inference.test.ts | 1 + .../worker-environments/live-events.test.ts | 47 + .../worker-environments/live-events.ts | 15 +- .../placement-dispatch-failure.ts | 262 ++++ .../placement-dispatch-recovery.ts | 231 +++ .../placement-dispatch.test.ts | 694 +++++++++ .../worker-environments/placement-dispatch.ts | 178 +++ .../placement-projector.test.ts | 107 ++ .../placement-projector.ts | 133 ++ .../worker-environments/placement-record.ts | 371 +++++ .../placement-row-codec.ts | 431 ++++++ .../placement-session-runtime.ts | 33 + .../worker-environments/placement-state.ts | 47 + .../placement-store.test.ts | 695 +++++++++ .../worker-environments/placement-store.ts | 414 ++++++ .../placement-turn-claims.ts | 352 +++++ .../placement-worker-gate.test.ts | 153 ++ .../placement-worker-gate.ts | 60 + .../worker-environments/service-contract.ts | 16 + .../worker-environments/service.test.ts | 551 ++++++- src/gateway/worker-environments/service.ts | 358 ++++- .../transcript-commit.test.ts | 1 + .../worker-environments/tunnel-contract.ts | 16 +- .../worker-environments/tunnel-ssh-runner.ts | 129 ++ .../worker-environments/tunnel.test.ts | 394 ++++- src/gateway/worker-environments/tunnel.ts | 192 +-- .../worker-turn-launcher.test.ts | 1305 +++++++++++++++++ .../worker-turn-launcher.ts | 464 ++++++ .../worker-turn-payload.ts | 197 +++ .../workspace-sync-local.test.ts | 57 + .../workspace-sync-local.ts | 235 +++ .../workspace-sync-scripts.ts | 185 +++ .../worker-environments/workspace-sync.ts | 407 +++++ src/infra/agent-events.test.ts | 39 + src/infra/agent-events.ts | 11 +- src/state/openclaw-state-db.generated.d.ts | 25 + src/state/openclaw-state-db.test.ts | 434 ++++++ src/state/openclaw-state-schema.generated.ts | 116 ++ src/state/openclaw-state-schema.sql | 116 ++ src/worker/embedded-agent-live.runtime.ts | 36 +- .../embedded-agent-transcript.runtime.ts | 128 +- src/worker/embedded-agent.runtime.ts | 21 +- src/worker/launch-descriptor.test.ts | 3 +- src/worker/launch-descriptor.ts | 12 +- src/worker/transcript-message.ts | 122 ++ src/worker/worker.fault-injection.test.ts | 7 +- src/worker/worker.runtime.test.ts | 69 +- src/worker/worker.runtime.ts | 19 +- 92 files changed, 13318 insertions(+), 963 deletions(-) create mode 100644 packages/gateway-protocol/src/schema/session-placement.test.ts create mode 100644 packages/gateway-protocol/src/schema/session-placement.ts create mode 100644 src/agents/session-placement-admission.test.ts create mode 100644 src/agents/session-placement-admission.ts create mode 100644 src/gateway/server-methods/sessions.dispatch.test.ts create mode 100644 src/gateway/server-worker-placement-startup.ts create mode 100644 src/gateway/server.sessions.placement-projection.test.ts create mode 100644 src/gateway/server.sessions.worker-placement-lifecycle.test.ts create mode 100644 src/gateway/worker-environments/placement-dispatch-failure.ts create mode 100644 src/gateway/worker-environments/placement-dispatch-recovery.ts create mode 100644 src/gateway/worker-environments/placement-dispatch.test.ts create mode 100644 src/gateway/worker-environments/placement-dispatch.ts create mode 100644 src/gateway/worker-environments/placement-projector.test.ts create mode 100644 src/gateway/worker-environments/placement-projector.ts create mode 100644 src/gateway/worker-environments/placement-record.ts create mode 100644 src/gateway/worker-environments/placement-row-codec.ts create mode 100644 src/gateway/worker-environments/placement-session-runtime.ts create mode 100644 src/gateway/worker-environments/placement-state.ts create mode 100644 src/gateway/worker-environments/placement-store.test.ts create mode 100644 src/gateway/worker-environments/placement-store.ts create mode 100644 src/gateway/worker-environments/placement-turn-claims.ts create mode 100644 src/gateway/worker-environments/placement-worker-gate.test.ts create mode 100644 src/gateway/worker-environments/placement-worker-gate.ts create mode 100644 src/gateway/worker-environments/tunnel-ssh-runner.ts create mode 100644 src/gateway/worker-environments/worker-turn-launcher.test.ts create mode 100644 src/gateway/worker-environments/worker-turn-launcher.ts create mode 100644 src/gateway/worker-environments/worker-turn-payload.ts create mode 100644 src/gateway/worker-environments/workspace-sync-local.test.ts create mode 100644 src/gateway/worker-environments/workspace-sync-local.ts create mode 100644 src/gateway/worker-environments/workspace-sync-scripts.ts create mode 100644 src/gateway/worker-environments/workspace-sync.ts diff --git a/apps/android/app/src/main/java/ai/openclaw/app/gateway/GatewayProtocol.kt b/apps/android/app/src/main/java/ai/openclaw/app/gateway/GatewayProtocol.kt index 571332376cbf..819e0a14f9a3 100644 --- a/apps/android/app/src/main/java/ai/openclaw/app/gateway/GatewayProtocol.kt +++ b/apps/android/app/src/main/java/ai/openclaw/app/gateway/GatewayProtocol.kt @@ -368,6 +368,7 @@ enum class GatewayMethod( ApprovalGet("approval.get"), ApprovalResolve("approval.resolve"), SessionsSearch("sessions.search"), + SessionsDispatch("sessions.dispatch"), } enum class GatewayEvent( diff --git a/apps/shared/OpenClawKit/Sources/OpenClawProtocol/GatewayModels.swift b/apps/shared/OpenClawKit/Sources/OpenClawProtocol/GatewayModels.swift index c51e90f4b0cd..7c514bb48b58 100644 --- a/apps/shared/OpenClawKit/Sources/OpenClawProtocol/GatewayModels.swift +++ b/apps/shared/OpenClawKit/Sources/OpenClawProtocol/GatewayModels.swift @@ -68,6 +68,19 @@ public enum NodePresenceAliveReason: String, Codable, Sendable { case connect = "connect" } +public enum SessionPlacementState: String, Codable, Sendable { + case local = "local" + case requested = "requested" + case provisioning = "provisioning" + case syncing = "syncing" + case starting = "starting" + case active = "active" + case draining = "draining" + case reconciling = "reconciling" + case reclaimed = "reclaimed" + case failed = "failed" +} + public enum SessionFileKind: String, Codable, Sendable { case modified = "modified" case read = "read" @@ -3332,6 +3345,526 @@ public struct SessionOperationEvent: Codable, Sendable { } } +public struct LocalSessionPlacement: Codable, Sendable { + public let state: String + public let generation: Int + public let createdatms: Int + public let updatedatms: Int + public let statechangedatms: Int + + public init( + state: String, + generation: Int, + createdatms: Int, + updatedatms: Int, + statechangedatms: Int) + { + self.state = state + self.generation = generation + self.createdatms = createdatms + self.updatedatms = updatedatms + self.statechangedatms = statechangedatms + } + + private enum CodingKeys: String, CodingKey { + case state + case generation + case createdatms = "createdAtMs" + case updatedatms = "updatedAtMs" + case statechangedatms = "stateChangedAtMs" + } +} + +public struct RequestedSessionPlacement: Codable, Sendable { + public let state: String + public let generation: Int + public let createdatms: Int + public let updatedatms: Int + public let statechangedatms: Int + + public init( + state: String, + generation: Int, + createdatms: Int, + updatedatms: Int, + statechangedatms: Int) + { + self.state = state + self.generation = generation + self.createdatms = createdatms + self.updatedatms = updatedatms + self.statechangedatms = statechangedatms + } + + private enum CodingKeys: String, CodingKey { + case state + case generation + case createdatms = "createdAtMs" + case updatedatms = "updatedAtMs" + case statechangedatms = "stateChangedAtMs" + } +} + +public struct ProvisioningSessionPlacement: Codable, Sendable { + public let state: String + public let generation: Int + public let createdatms: Int + public let updatedatms: Int + public let statechangedatms: Int + public let environmentid: String? + + public init( + state: String, + generation: Int, + createdatms: Int, + updatedatms: Int, + statechangedatms: Int, + environmentid: String? = nil) + { + self.state = state + self.generation = generation + self.createdatms = createdatms + self.updatedatms = updatedatms + self.statechangedatms = statechangedatms + self.environmentid = environmentid + } + + private enum CodingKeys: String, CodingKey { + case state + case generation + case createdatms = "createdAtMs" + case updatedatms = "updatedAtMs" + case statechangedatms = "stateChangedAtMs" + case environmentid = "environmentId" + } +} + +public struct SyncingSessionPlacement: Codable, Sendable { + public let state: String + public let generation: Int + public let createdatms: Int + public let updatedatms: Int + public let statechangedatms: Int + public let environmentid: String + public let workerbundlehash: String + + public init( + state: String, + generation: Int, + createdatms: Int, + updatedatms: Int, + statechangedatms: Int, + environmentid: String, + workerbundlehash: String) + { + self.state = state + self.generation = generation + self.createdatms = createdatms + self.updatedatms = updatedatms + self.statechangedatms = statechangedatms + self.environmentid = environmentid + self.workerbundlehash = workerbundlehash + } + + private enum CodingKeys: String, CodingKey { + case state + case generation + case createdatms = "createdAtMs" + case updatedatms = "updatedAtMs" + case statechangedatms = "stateChangedAtMs" + case environmentid = "environmentId" + case workerbundlehash = "workerBundleHash" + } +} + +public struct StartingSessionPlacement: Codable, Sendable { + public let state: String + public let generation: Int + public let createdatms: Int + public let updatedatms: Int + public let statechangedatms: Int + public let environmentid: String + public let workerbundlehash: String + public let workspacebasemanifestref: String + public let remoteworkspacedir: String + + public init( + state: String, + generation: Int, + createdatms: Int, + updatedatms: Int, + statechangedatms: Int, + environmentid: String, + workerbundlehash: String, + workspacebasemanifestref: String, + remoteworkspacedir: String) + { + self.state = state + self.generation = generation + self.createdatms = createdatms + self.updatedatms = updatedatms + self.statechangedatms = statechangedatms + self.environmentid = environmentid + self.workerbundlehash = workerbundlehash + self.workspacebasemanifestref = workspacebasemanifestref + self.remoteworkspacedir = remoteworkspacedir + } + + private enum CodingKeys: String, CodingKey { + case state + case generation + case createdatms = "createdAtMs" + case updatedatms = "updatedAtMs" + case statechangedatms = "stateChangedAtMs" + case environmentid = "environmentId" + case workerbundlehash = "workerBundleHash" + case workspacebasemanifestref = "workspaceBaseManifestRef" + case remoteworkspacedir = "remoteWorkspaceDir" + } +} + +public struct ActiveWorkerSessionPlacement: Codable, Sendable { + public let state: String + public let generation: Int + public let createdatms: Int + public let updatedatms: Int + public let statechangedatms: Int + public let environmentid: String + public let activeownerepoch: Int + public let workerbundlehash: String + public let workspacebasemanifestref: String + public let remoteworkspacedir: String + public let lasttranscriptackcursor: Int? + public let lastliveeventackcursor: Int? + + public init( + state: String, + generation: Int, + createdatms: Int, + updatedatms: Int, + statechangedatms: Int, + environmentid: String, + activeownerepoch: Int, + workerbundlehash: String, + workspacebasemanifestref: String, + remoteworkspacedir: String, + lasttranscriptackcursor: Int? = nil, + lastliveeventackcursor: Int? = nil) + { + self.state = state + self.generation = generation + self.createdatms = createdatms + self.updatedatms = updatedatms + self.statechangedatms = statechangedatms + self.environmentid = environmentid + self.activeownerepoch = activeownerepoch + self.workerbundlehash = workerbundlehash + self.workspacebasemanifestref = workspacebasemanifestref + self.remoteworkspacedir = remoteworkspacedir + self.lasttranscriptackcursor = lasttranscriptackcursor + self.lastliveeventackcursor = lastliveeventackcursor + } + + private enum CodingKeys: String, CodingKey { + case state + case generation + case createdatms = "createdAtMs" + case updatedatms = "updatedAtMs" + case statechangedatms = "stateChangedAtMs" + case environmentid = "environmentId" + case activeownerepoch = "activeOwnerEpoch" + case workerbundlehash = "workerBundleHash" + case workspacebasemanifestref = "workspaceBaseManifestRef" + case remoteworkspacedir = "remoteWorkspaceDir" + case lasttranscriptackcursor = "lastTranscriptAckCursor" + case lastliveeventackcursor = "lastLiveEventAckCursor" + } +} + +public struct DrainingSessionPlacement: Codable, Sendable { + public let state: String + public let generation: Int + public let createdatms: Int + public let updatedatms: Int + public let statechangedatms: Int + public let environmentid: String + public let activeownerepoch: Int + public let workerbundlehash: String + public let workspacebasemanifestref: String + public let remoteworkspacedir: String + public let lasttranscriptackcursor: Int? + public let lastliveeventackcursor: Int? + + public init( + state: String, + generation: Int, + createdatms: Int, + updatedatms: Int, + statechangedatms: Int, + environmentid: String, + activeownerepoch: Int, + workerbundlehash: String, + workspacebasemanifestref: String, + remoteworkspacedir: String, + lasttranscriptackcursor: Int? = nil, + lastliveeventackcursor: Int? = nil) + { + self.state = state + self.generation = generation + self.createdatms = createdatms + self.updatedatms = updatedatms + self.statechangedatms = statechangedatms + self.environmentid = environmentid + self.activeownerepoch = activeownerepoch + self.workerbundlehash = workerbundlehash + self.workspacebasemanifestref = workspacebasemanifestref + self.remoteworkspacedir = remoteworkspacedir + self.lasttranscriptackcursor = lasttranscriptackcursor + self.lastliveeventackcursor = lastliveeventackcursor + } + + private enum CodingKeys: String, CodingKey { + case state + case generation + case createdatms = "createdAtMs" + case updatedatms = "updatedAtMs" + case statechangedatms = "stateChangedAtMs" + case environmentid = "environmentId" + case activeownerepoch = "activeOwnerEpoch" + case workerbundlehash = "workerBundleHash" + case workspacebasemanifestref = "workspaceBaseManifestRef" + case remoteworkspacedir = "remoteWorkspaceDir" + case lasttranscriptackcursor = "lastTranscriptAckCursor" + case lastliveeventackcursor = "lastLiveEventAckCursor" + } +} + +public struct ReconcilingSessionPlacement: Codable, Sendable { + public let state: String + public let generation: Int + public let createdatms: Int + public let updatedatms: Int + public let statechangedatms: Int + public let environmentid: String + public let activeownerepoch: Int + public let workerbundlehash: String + public let workspacebasemanifestref: String + public let remoteworkspacedir: String + public let lasttranscriptackcursor: Int? + public let lastliveeventackcursor: Int? + + public init( + state: String, + generation: Int, + createdatms: Int, + updatedatms: Int, + statechangedatms: Int, + environmentid: String, + activeownerepoch: Int, + workerbundlehash: String, + workspacebasemanifestref: String, + remoteworkspacedir: String, + lasttranscriptackcursor: Int? = nil, + lastliveeventackcursor: Int? = nil) + { + self.state = state + self.generation = generation + self.createdatms = createdatms + self.updatedatms = updatedatms + self.statechangedatms = statechangedatms + self.environmentid = environmentid + self.activeownerepoch = activeownerepoch + self.workerbundlehash = workerbundlehash + self.workspacebasemanifestref = workspacebasemanifestref + self.remoteworkspacedir = remoteworkspacedir + self.lasttranscriptackcursor = lasttranscriptackcursor + self.lastliveeventackcursor = lastliveeventackcursor + } + + private enum CodingKeys: String, CodingKey { + case state + case generation + case createdatms = "createdAtMs" + case updatedatms = "updatedAtMs" + case statechangedatms = "stateChangedAtMs" + case environmentid = "environmentId" + case activeownerepoch = "activeOwnerEpoch" + case workerbundlehash = "workerBundleHash" + case workspacebasemanifestref = "workspaceBaseManifestRef" + case remoteworkspacedir = "remoteWorkspaceDir" + case lasttranscriptackcursor = "lastTranscriptAckCursor" + case lastliveeventackcursor = "lastLiveEventAckCursor" + } +} + +public struct ReclaimedSessionPlacement: Codable, Sendable { + public let state: String + public let generation: Int + public let createdatms: Int + public let updatedatms: Int + public let statechangedatms: Int + public let environmentid: String? + public let activeownerepoch: Int? + public let workspacebasemanifestref: String? + public let remoteworkspacedir: String? + public let workerbundlehash: String? + public let lasttranscriptackcursor: Int? + public let lastliveeventackcursor: Int? + + public init( + state: String, + generation: Int, + createdatms: Int, + updatedatms: Int, + statechangedatms: Int, + environmentid: String? = nil, + activeownerepoch: Int? = nil, + workspacebasemanifestref: String? = nil, + remoteworkspacedir: String? = nil, + workerbundlehash: String? = nil, + lasttranscriptackcursor: Int? = nil, + lastliveeventackcursor: Int? = nil) + { + self.state = state + self.generation = generation + self.createdatms = createdatms + self.updatedatms = updatedatms + self.statechangedatms = statechangedatms + self.environmentid = environmentid + self.activeownerepoch = activeownerepoch + self.workspacebasemanifestref = workspacebasemanifestref + self.remoteworkspacedir = remoteworkspacedir + self.workerbundlehash = workerbundlehash + self.lasttranscriptackcursor = lasttranscriptackcursor + self.lastliveeventackcursor = lastliveeventackcursor + } + + private enum CodingKeys: String, CodingKey { + case state + case generation + case createdatms = "createdAtMs" + case updatedatms = "updatedAtMs" + case statechangedatms = "stateChangedAtMs" + case environmentid = "environmentId" + case activeownerepoch = "activeOwnerEpoch" + case workspacebasemanifestref = "workspaceBaseManifestRef" + case remoteworkspacedir = "remoteWorkspaceDir" + case workerbundlehash = "workerBundleHash" + case lasttranscriptackcursor = "lastTranscriptAckCursor" + case lastliveeventackcursor = "lastLiveEventAckCursor" + } +} + +public struct FailedSessionPlacement: Codable, Sendable { + public let state: String + public let generation: Int + public let createdatms: Int + public let updatedatms: Int + public let statechangedatms: Int + public let environmentid: String? + public let activeownerepoch: Int? + public let workspacebasemanifestref: String? + public let remoteworkspacedir: String? + public let workerbundlehash: String? + public let lasttranscriptackcursor: Int? + public let lastliveeventackcursor: Int? + public let recoveryerror: String + + public init( + state: String, + generation: Int, + createdatms: Int, + updatedatms: Int, + statechangedatms: Int, + environmentid: String? = nil, + activeownerepoch: Int? = nil, + workspacebasemanifestref: String? = nil, + remoteworkspacedir: String? = nil, + workerbundlehash: String? = nil, + lasttranscriptackcursor: Int? = nil, + lastliveeventackcursor: Int? = nil, + recoveryerror: String) + { + self.state = state + self.generation = generation + self.createdatms = createdatms + self.updatedatms = updatedatms + self.statechangedatms = statechangedatms + self.environmentid = environmentid + self.activeownerepoch = activeownerepoch + self.workspacebasemanifestref = workspacebasemanifestref + self.remoteworkspacedir = remoteworkspacedir + self.workerbundlehash = workerbundlehash + self.lasttranscriptackcursor = lasttranscriptackcursor + self.lastliveeventackcursor = lastliveeventackcursor + self.recoveryerror = recoveryerror + } + + private enum CodingKeys: String, CodingKey { + case state + case generation + case createdatms = "createdAtMs" + case updatedatms = "updatedAtMs" + case statechangedatms = "stateChangedAtMs" + case environmentid = "environmentId" + case activeownerepoch = "activeOwnerEpoch" + case workspacebasemanifestref = "workspaceBaseManifestRef" + case remoteworkspacedir = "remoteWorkspaceDir" + case workerbundlehash = "workerBundleHash" + case lasttranscriptackcursor = "lastTranscriptAckCursor" + case lastliveeventackcursor = "lastLiveEventAckCursor" + case recoveryerror = "recoveryError" + } +} + +public struct SessionsDispatchParams: Codable, Sendable { + public let key: String + public let agentid: String? + public let profileid: String + + public init( + key: String, + agentid: String? = nil, + profileid: String) + { + self.key = key + self.agentid = agentid + self.profileid = profileid + } + + private enum CodingKeys: String, CodingKey { + case key + case agentid = "agentId" + case profileid = "profileId" + } +} + +public struct SessionsDispatchResult: Codable, Sendable { + public let ok: Bool + public let key: String + public let sessionid: String + public let placement: ActiveWorkerSessionPlacement + + public init( + ok: Bool, + key: String, + sessionid: String, + placement: ActiveWorkerSessionPlacement) + { + self.ok = ok + self.key = key + self.sessionid = sessionid + self.placement = placement + } + + private enum CodingKeys: String, CodingKey { + case ok + case key + case sessionid = "sessionId" + case placement + } +} + public struct SessionsCompactionListParams: Codable, Sendable { public let key: String public let agentid: String? @@ -12359,6 +12892,61 @@ public enum GatewaySuspendStatusResult: Codable, Sendable { } } +public enum SessionPlacement: Codable, Sendable { + case local(LocalSessionPlacement) + case requested(RequestedSessionPlacement) + case provisioning(ProvisioningSessionPlacement) + case syncing(SyncingSessionPlacement) + case starting(StartingSessionPlacement) + case active(ActiveWorkerSessionPlacement) + case draining(DrainingSessionPlacement) + case reconciling(ReconcilingSessionPlacement) + case reclaimed(ReclaimedSessionPlacement) + case failed(FailedSessionPlacement) + + private enum CodingKeys: String, CodingKey { + case discriminator = "state" + } + + 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 "local": self = try .local(LocalSessionPlacement(from: decoder)) + case "requested": self = try .requested(RequestedSessionPlacement(from: decoder)) + case "provisioning": self = try .provisioning(ProvisioningSessionPlacement(from: decoder)) + case "syncing": self = try .syncing(SyncingSessionPlacement(from: decoder)) + case "starting": self = try .starting(StartingSessionPlacement(from: decoder)) + case "active": self = try .active(ActiveWorkerSessionPlacement(from: decoder)) + case "draining": self = try .draining(DrainingSessionPlacement(from: decoder)) + case "reconciling": self = try .reconciling(ReconcilingSessionPlacement(from: decoder)) + case "reclaimed": self = try .reclaimed(ReclaimedSessionPlacement(from: decoder)) + case "failed": self = try .failed(FailedSessionPlacement(from: decoder)) + default: + throw DecodingError.dataCorruptedError( + forKey: .discriminator, + in: container, + debugDescription: "Unknown SessionPlacement discriminator value" + ) + } + } + + public func encode(to encoder: Encoder) throws { + switch self { + case .local(let value): try value.encode(to: encoder) + case .requested(let value): try value.encode(to: encoder) + case .provisioning(let value): try value.encode(to: encoder) + case .syncing(let value): try value.encode(to: encoder) + case .starting(let value): try value.encode(to: encoder) + case .active(let value): try value.encode(to: encoder) + case .draining(let value): try value.encode(to: encoder) + case .reconciling(let value): try value.encode(to: encoder) + case .reclaimed(let value): try value.encode(to: encoder) + case .failed(let value): try value.encode(to: encoder) + } + } +} + public enum AuditActivityEventV1: Codable, Sendable { case agentRun(AuditActivityAgentRunV1) case toolAction(AuditActivityToolActionV1) diff --git a/docs/cli/worker.md b/docs/cli/worker.md index f1bbb38f97b4..cd33ffe2f4ad 100644 --- a/docs/cli/worker.md +++ b/docs/cli/worker.md @@ -46,8 +46,10 @@ Worker mode does not start channels, Gateway HTTP surfaces, or plugin auto-start beyond the assigned session toolset. It uses a throwaway state directory and has no standing provider or forge credentials. -Worker-to-worker session dispatch is not exposed in this mode. Agent dispatch -and placement remain gateway-owned milestone-3 surfaces. +Worker-to-worker session dispatch is not exposed in this mode. Placement and +dispatch remain gateway-owned: an operator can dispatch an existing local, +managed-worktree session through the Gateway, while a worker process cannot +dispatch itself or another worker. The prepared assignment carries the transcript context, accepted base leaf, commit sequence, and live-event cursor. On a tunnel reconnect, the process diff --git a/docs/gateway/protocol.md b/docs/gateway/protocol.md index 9827a6e65bfc..f3ecf3513af3 100644 --- a/docs/gateway/protocol.md +++ b/docs/gateway/protocol.md @@ -515,13 +515,14 @@ methods. Treat this as feature discovery, not a full enumeration of - - `sessions.list` returns the current session index, including per-row `agentRuntime` metadata when an agent runtime backend is configured. + - `sessions.list` returns the current session index, including per-row `agentRuntime` metadata when an agent runtime backend is configured. When cloud-worker placement is enabled or durable recovery state exists, session rows also include a closed `placement` state (`local`, `requested`, `provisioning`, `syncing`, `starting`, `active`, `draining`, `reconciling`, `reclaimed`, or `failed`) plus state-specific environment, owner-epoch, workspace, bundle, ACK-cursor, or recovery fields. - `sessions.subscribe` and `sessions.unsubscribe` toggle session change event subscriptions for the current WS client. - `sessions.messages.subscribe` and `sessions.messages.unsubscribe` toggle transcript/message event subscriptions for one session. Pass `includeApprovals: true` to also receive sanitized `session.approval` lifecycle events for approvals whose persisted audience includes that exact session and whose reviewer binding authorizes the subscribing client. The subscribe response then includes a bounded pending `approvalReplay`; it is authoritative when `truncated` is false. The opt-in is per subscribe call, not sticky: re-subscribing to the same session without `includeApprovals: true` removes an existing approval subscription. In addition to normal session-read authority, this opt-in requires `operator.admin`, or `operator.approvals` on a paired device. - `sessions.preview` returns bounded transcript previews for specific session keys. - `sessions.describe` returns one gateway session row for an exact session key. - `sessions.resolve` resolves or canonicalizes a session target. - `sessions.create` creates a new session entry. `worktree: true` provisions a managed worktree; optional `worktreeBaseRef`/`worktreeName` select the base ref and branch name, and `execNode` (`operator.admin`) binds session exec to a node host. The created worktree is echoed in the result and persisted on the session row (`worktree: { id, branch, repoRoot }`). When the entry is created but its nested initial `chat.send` is rejected, the successful result includes `runStarted: false` and `runError`; clients can preserve the prompt and retry against the returned session key. + - `sessions.dispatch` (`operator.admin`) moves an existing local OpenClaw session with a session-owned managed worktree to a configured cloud-worker profile. Pass `{ key, profileId, agentId? }`. The method is absent when no worker profile is configured, closes local turn admission before draining active work, and returns only after placement reaches `active` worker ownership. Dispatch is one-way; worker-to-local pull-back is not part of this RPC. - `sessions.groups.list`, `sessions.groups.put`, `sessions.groups.rename`, and `sessions.groups.delete` manage the gateway-owned custom session group catalog (names + display order). Membership stays on each session's `category` field; rename and delete update member sessions server-side. - `sessions.send` sends a message into an existing session. - `sessions.steer` is the interrupt-and-steer variant for an active session. diff --git a/packages/gateway-protocol/src/index.ts b/packages/gateway-protocol/src/index.ts index 26bbe4055834..587d9fd705cc 100644 --- a/packages/gateway-protocol/src/index.ts +++ b/packages/gateway-protocol/src/index.ts @@ -343,11 +343,15 @@ import { SessionFileEntrySchema, SessionFileKindSchema, SessionFileRelevanceSchema, + SessionPlacementSchema, + SessionPlacementStateSchema, SessionWorktreeInfoSchema, SessionsCreateParamsSchema, SessionsCreateResultSchema, SessionsDeleteParamsSchema, SessionsDescribeParamsSchema, + SessionsDispatchParamsSchema, + SessionsDispatchResultSchema, SessionGroupSchema, SessionsGroupsDeleteParamsSchema, SessionsGroupsListParamsSchema, @@ -706,6 +710,8 @@ export const validateSessionsFilesSetParams = lazyCompile(SessionsFilesSetParams export const validateSessionsDiffParams = lazyCompile(SessionsDiffParamsSchema); export const validateSessionsCreateParams = lazyCompile(SessionsCreateParamsSchema); export const validateSessionsSendParams = lazyCompile(SessionsSendParamsSchema); +export const validateSessionsDispatchParams = lazyCompile(SessionsDispatchParamsSchema); +export const validateSessionsDispatchResult = lazyCompile(SessionsDispatchResultSchema); export const validateSessionsMessagesSubscribeParams = lazyCompile( SessionsMessagesSubscribeParamsSchema, ); @@ -1058,9 +1064,13 @@ export { SessionsCompactionGetParamsSchema, SessionsCompactionBranchParamsSchema, SessionsCompactionRestoreParamsSchema, + SessionPlacementStateSchema, + SessionPlacementSchema, SessionWorktreeInfoSchema, SessionsCreateParamsSchema, SessionsCreateResultSchema, + SessionsDispatchParamsSchema, + SessionsDispatchResultSchema, SessionsSendParamsSchema, SessionsAbortParamsSchema, SessionsPatchParamsSchema, @@ -1600,7 +1610,11 @@ export type { SessionsDescribeParams, SessionsResolveParams, SessionOperationEvent, + SessionPlacementState, + SessionPlacement, SessionWorktreeInfo, + SessionsDispatchParams, + SessionsDispatchResult, SessionsCreateResult, SessionsPatchParams, SessionsResetParams, diff --git a/packages/gateway-protocol/src/schema.ts b/packages/gateway-protocol/src/schema.ts index 1d1abe4382b4..96545144a6c4 100644 --- a/packages/gateway-protocol/src/schema.ts +++ b/packages/gateway-protocol/src/schema.ts @@ -30,6 +30,7 @@ export * from "./schema/nodes.js"; export * from "./schema/protocol-schemas.js"; export * from "./schema/push.js"; export * from "./schema/secrets.js"; +export * from "./schema/session-placement.js"; export * from "./schema/sessions.js"; export * from "./schema/sessions-catalog.js"; export * from "./schema/snapshot.js"; diff --git a/packages/gateway-protocol/src/schema/protocol-schemas.ts b/packages/gateway-protocol/src/schema/protocol-schemas.ts index de8faa789b37..9dcc243c8533 100644 --- a/packages/gateway-protocol/src/schema/protocol-schemas.ts +++ b/packages/gateway-protocol/src/schema/protocol-schemas.ts @@ -362,6 +362,7 @@ import { SecretsResolveParamsSchema, SecretsResolveResultSchema, } from "./secrets.js"; +import { SessionPlacementProtocolSchemas } from "./session-placement.js"; import { SessionCatalogCapabilitiesSchema, SessionCatalogDescriptorSchema, @@ -629,6 +630,7 @@ export const ProtocolSchemas = { SessionsSearchResult: SessionsSearchResultSchema, SessionCompactionCheckpoint: SessionCompactionCheckpointSchema, SessionOperationEvent: SessionOperationEventSchema, + ...SessionPlacementProtocolSchemas, SessionsCompactionListParams: SessionsCompactionListParamsSchema, SessionsCompactionGetParams: SessionsCompactionGetParamsSchema, SessionsCompactionBranchParams: SessionsCompactionBranchParamsSchema, diff --git a/packages/gateway-protocol/src/schema/session-placement.test.ts b/packages/gateway-protocol/src/schema/session-placement.test.ts new file mode 100644 index 000000000000..b33802cd7f5e --- /dev/null +++ b/packages/gateway-protocol/src/schema/session-placement.test.ts @@ -0,0 +1,255 @@ +import { Value } from "typebox/value"; +import { describe, expect, it } from "vitest"; +import { + SessionPlacementSchema, + SessionPlacementStateSchema, + validateSessionsDispatchParams, + validateSessionsDispatchResult, +} from "../index.js"; + +const placementStates = [ + "local", + "requested", + "provisioning", + "syncing", + "starting", + "active", + "draining", + "reconciling", + "reclaimed", + "failed", +] as const; + +const basePlacement = { + generation: 4, + createdAtMs: 100, + updatedAtMs: 200, + stateChangedAtMs: 150, +}; +const workerBundleHash = "a".repeat(64); +const environmentFields = { + environmentId: "environment-1", + workerBundleHash, +}; +const workspaceFields = { + workspaceBaseManifestRef: "manifest-1", + remoteWorkspaceDir: "/workspace/session-1", +}; +const workerOwnedFields = { + ...environmentFields, + ...workspaceFields, + activeOwnerEpoch: 7, +}; + +describe("session dispatch protocol schemas", () => { + it("accepts only the dedicated dispatch selector and configured profile", () => { + expect( + validateSessionsDispatchParams({ + key: "agent:main:dispatch", + agentId: "main", + profileId: "development", + }), + ).toBe(true); + expect(validateSessionsDispatchParams({ key: "agent:main:dispatch" })).toBe(false); + expect( + validateSessionsDispatchParams({ + key: "agent:main:dispatch", + profileId: "development", + task: "run remotely", + }), + ).toBe(false); + }); + + it("keeps placement states closed", () => { + for (const state of placementStates) { + expect(Value.Check(SessionPlacementStateSchema, state)).toBe(true); + } + expect(Value.Check(SessionPlacementStateSchema, "unknown")).toBe(false); + }); + + it("keeps local and requested placement free of worker metadata", () => { + expect(Value.Check(SessionPlacementSchema, { state: "local", ...basePlacement })).toBe(true); + expect(Value.Check(SessionPlacementSchema, { state: "requested", ...basePlacement })).toBe( + true, + ); + expect( + Value.Check(SessionPlacementSchema, { + state: "local", + ...basePlacement, + environmentId: "environment-1", + }), + ).toBe(false); + expect( + Value.Check(SessionPlacementSchema, { + state: "requested", + ...basePlacement, + workerBundleHash, + }), + ).toBe(false); + }); + + it("allows only the optional reserved environment while provisioning", () => { + expect( + Value.Check(SessionPlacementSchema, { + state: "provisioning", + ...basePlacement, + environmentId: "environment-1", + }), + ).toBe(true); + expect(Value.Check(SessionPlacementSchema, { state: "provisioning", ...basePlacement })).toBe( + true, + ); + expect( + Value.Check(SessionPlacementSchema, { + state: "provisioning", + ...basePlacement, + ...environmentFields, + }), + ).toBe(false); + }); + + it("requires the provisioned bundle while syncing", () => { + expect( + Value.Check(SessionPlacementSchema, { + state: "syncing", + ...basePlacement, + ...environmentFields, + }), + ).toBe(true); + expect( + Value.Check(SessionPlacementSchema, { + state: "syncing", + ...basePlacement, + environmentId: "environment-1", + }), + ).toBe(false); + expect( + Value.Check(SessionPlacementSchema, { + state: "syncing", + ...basePlacement, + ...environmentFields, + ...workspaceFields, + }), + ).toBe(false); + }); + + it("requires workspace identity while starting", () => { + expect( + Value.Check(SessionPlacementSchema, { + state: "starting", + ...basePlacement, + ...environmentFields, + ...workspaceFields, + }), + ).toBe(true); + expect( + Value.Check(SessionPlacementSchema, { + state: "starting", + ...basePlacement, + ...environmentFields, + remoteWorkspaceDir: "/workspace/session-1", + }), + ).toBe(false); + expect( + Value.Check(SessionPlacementSchema, { + state: "starting", + ...basePlacement, + ...environmentFields, + ...workspaceFields, + lastTranscriptAckCursor: 0, + }), + ).toBe(false); + }); + + it.each(["active", "draining", "reconciling"] as const)( + "requires complete worker ownership for %s placement", + (state) => { + expect( + Value.Check(SessionPlacementSchema, { + state, + ...basePlacement, + ...workerOwnedFields, + lastTranscriptAckCursor: 2, + lastLiveEventAckCursor: 9, + }), + ).toBe(true); + expect( + Value.Check(SessionPlacementSchema, { + state, + ...basePlacement, + environmentId: "environment-1", + activeOwnerEpoch: 7, + workerBundleHash, + }), + ).toBe(false); + }, + ); + + it("preserves optional provenance only in terminal states", () => { + expect(Value.Check(SessionPlacementSchema, { state: "reclaimed", ...basePlacement })).toBe( + true, + ); + expect( + Value.Check(SessionPlacementSchema, { + state: "reclaimed", + ...basePlacement, + ...workerOwnedFields, + }), + ).toBe(true); + }); + + it("requires recovery evidence for failed placement", () => { + const failed = { + state: "failed" as const, + ...basePlacement, + ...workerOwnedFields, + recoveryError: "worker admission failed", + }; + expect(Value.Check(SessionPlacementSchema, failed)).toBe(true); + expect( + Value.Check(SessionPlacementSchema, { + state: "failed", + ...basePlacement, + }), + ).toBe(false); + }); + + it("accepts only active worker ownership in successful dispatch results", () => { + const active = { + state: "active" as const, + ...basePlacement, + ...workerOwnedFields, + }; + expect( + validateSessionsDispatchResult({ + ok: true, + key: "agent:main:dispatch", + sessionId: "session-1", + placement: active, + }), + ).toBe(true); + expect( + validateSessionsDispatchResult({ + ok: true, + key: "agent:main:dispatch", + sessionId: "session-1", + placement: { + state: "failed", + ...basePlacement, + recoveryError: "worker admission failed", + }, + }), + ).toBe(false); + }); + + it("rejects unknown placement fields", () => { + expect( + Value.Check(SessionPlacementSchema, { + state: "active", + ...basePlacement, + ...workerOwnedFields, + unexpected: true, + }), + ).toBe(false); + }); +}); diff --git a/packages/gateway-protocol/src/schema/session-placement.ts b/packages/gateway-protocol/src/schema/session-placement.ts new file mode 100644 index 000000000000..84468ae33d58 --- /dev/null +++ b/packages/gateway-protocol/src/schema/session-placement.ts @@ -0,0 +1,198 @@ +import type { Static } from "typebox"; +import { Type } from "typebox"; +import { NonEmptyString } from "./primitives.js"; + +/** Durable gateway ownership states for one session execution placement. */ +export const SessionPlacementStateSchema = Type.Union([ + Type.Literal("local"), + Type.Literal("requested"), + Type.Literal("provisioning"), + Type.Literal("syncing"), + Type.Literal("starting"), + Type.Literal("active"), + Type.Literal("draining"), + Type.Literal("reconciling"), + Type.Literal("reclaimed"), + Type.Literal("failed"), +]); + +const SessionPlacementTimingProperties = { + generation: Type.Integer({ minimum: 0, maximum: Number.MAX_SAFE_INTEGER }), + createdAtMs: Type.Integer({ minimum: 0, maximum: Number.MAX_SAFE_INTEGER }), + updatedAtMs: Type.Integer({ minimum: 0, maximum: Number.MAX_SAFE_INTEGER }), + stateChangedAtMs: Type.Integer({ minimum: 0, maximum: Number.MAX_SAFE_INTEGER }), +}; + +const SessionPlacementOwnerEpochSchema = Type.Integer({ + minimum: 1, + maximum: Number.MAX_SAFE_INTEGER, +}); + +const WorkerBundleHashSchema = Type.String({ + minLength: 64, + maxLength: 64, + pattern: "^[a-f0-9]{64}$", +}); + +const SessionPlacementWorkspaceProperties = { + workspaceBaseManifestRef: NonEmptyString, + remoteWorkspaceDir: NonEmptyString, +}; + +const SessionPlacementAckProperties = { + lastTranscriptAckCursor: Type.Optional( + Type.Integer({ minimum: 0, maximum: Number.MAX_SAFE_INTEGER }), + ), + lastLiveEventAckCursor: Type.Optional( + Type.Integer({ minimum: 0, maximum: Number.MAX_SAFE_INTEGER }), + ), +}; + +const TerminalSessionPlacementProperties = { + environmentId: Type.Optional(NonEmptyString), + activeOwnerEpoch: Type.Optional(SessionPlacementOwnerEpochSchema), + workspaceBaseManifestRef: Type.Optional(NonEmptyString), + remoteWorkspaceDir: Type.Optional(NonEmptyString), + workerBundleHash: Type.Optional(WorkerBundleHashSchema), + ...SessionPlacementAckProperties, +}; + +function createUnownedSessionPlacementSchema( + state: State, +) { + return Type.Object( + { state: Type.Literal(state), ...SessionPlacementTimingProperties }, + { additionalProperties: false }, + ); +} + +function createWorkerOwnedSessionPlacementSchema< + const State extends "active" | "draining" | "reconciling", +>(state: State) { + return Type.Object( + { + state: Type.Literal(state), + ...SessionPlacementTimingProperties, + environmentId: NonEmptyString, + activeOwnerEpoch: SessionPlacementOwnerEpochSchema, + workerBundleHash: WorkerBundleHashSchema, + ...SessionPlacementWorkspaceProperties, + ...SessionPlacementAckProperties, + }, + { additionalProperties: false }, + ); +} + +export const LocalSessionPlacementSchema = createUnownedSessionPlacementSchema("local"); +export const RequestedSessionPlacementSchema = createUnownedSessionPlacementSchema("requested"); + +export const ProvisioningSessionPlacementSchema = Type.Object( + { + state: Type.Literal("provisioning"), + ...SessionPlacementTimingProperties, + environmentId: Type.Optional(NonEmptyString), + }, + { additionalProperties: false }, +); + +export const SyncingSessionPlacementSchema = Type.Object( + { + state: Type.Literal("syncing"), + ...SessionPlacementTimingProperties, + environmentId: NonEmptyString, + workerBundleHash: WorkerBundleHashSchema, + }, + { additionalProperties: false }, +); + +export const StartingSessionPlacementSchema = Type.Object( + { + state: Type.Literal("starting"), + ...SessionPlacementTimingProperties, + environmentId: NonEmptyString, + workerBundleHash: WorkerBundleHashSchema, + ...SessionPlacementWorkspaceProperties, + }, + { additionalProperties: false }, +); + +export const ActiveWorkerSessionPlacementSchema = createWorkerOwnedSessionPlacementSchema("active"); +export const DrainingSessionPlacementSchema = createWorkerOwnedSessionPlacementSchema("draining"); +export const ReconcilingSessionPlacementSchema = + createWorkerOwnedSessionPlacementSchema("reconciling"); + +export const ReclaimedSessionPlacementSchema = Type.Object( + { + state: Type.Literal("reclaimed"), + ...SessionPlacementTimingProperties, + ...TerminalSessionPlacementProperties, + }, + { additionalProperties: false }, +); + +export const FailedSessionPlacementSchema = Type.Object( + { + state: Type.Literal("failed"), + ...SessionPlacementTimingProperties, + ...TerminalSessionPlacementProperties, + recoveryError: NonEmptyString, + }, + { additionalProperties: false }, +); + +/** Gateway-visible placement projection; `state` remains the closed discriminator. */ +export const SessionPlacementSchema = Type.Union([ + LocalSessionPlacementSchema, + RequestedSessionPlacementSchema, + ProvisioningSessionPlacementSchema, + SyncingSessionPlacementSchema, + StartingSessionPlacementSchema, + ActiveWorkerSessionPlacementSchema, + DrainingSessionPlacementSchema, + ReconcilingSessionPlacementSchema, + ReclaimedSessionPlacementSchema, + FailedSessionPlacementSchema, +]); + +/** Requests one-way dispatch of an existing local session to a configured worker profile. */ +export const SessionsDispatchParamsSchema = Type.Object( + { + key: NonEmptyString, + agentId: Type.Optional(NonEmptyString), + profileId: NonEmptyString, + }, + { additionalProperties: false }, +); + +/** Result returned once session dispatch reaches durable worker ownership. */ +export const SessionsDispatchResultSchema = Type.Object( + { + ok: Type.Literal(true), + key: NonEmptyString, + sessionId: NonEmptyString, + placement: ActiveWorkerSessionPlacementSchema, + }, + { additionalProperties: false }, +); + +export const SessionPlacementProtocolSchemas = { + SessionPlacementState: SessionPlacementStateSchema, + LocalSessionPlacement: LocalSessionPlacementSchema, + RequestedSessionPlacement: RequestedSessionPlacementSchema, + ProvisioningSessionPlacement: ProvisioningSessionPlacementSchema, + SyncingSessionPlacement: SyncingSessionPlacementSchema, + StartingSessionPlacement: StartingSessionPlacementSchema, + ActiveWorkerSessionPlacement: ActiveWorkerSessionPlacementSchema, + DrainingSessionPlacement: DrainingSessionPlacementSchema, + ReconcilingSessionPlacement: ReconcilingSessionPlacementSchema, + ReclaimedSessionPlacement: ReclaimedSessionPlacementSchema, + FailedSessionPlacement: FailedSessionPlacementSchema, + SessionPlacement: SessionPlacementSchema, + SessionsDispatchParams: SessionsDispatchParamsSchema, + SessionsDispatchResult: SessionsDispatchResultSchema, +} as const; + +export type SessionPlacementState = Static; +export type SessionPlacement = Static; +export type SessionsDispatchParams = Static; +export type SessionsDispatchResult = Static; diff --git a/packages/gateway-protocol/src/schema/worker-admission.test.ts b/packages/gateway-protocol/src/schema/worker-admission.test.ts index de8809364636..c3a02e4e373f 100644 --- a/packages/gateway-protocol/src/schema/worker-admission.test.ts +++ b/packages/gateway-protocol/src/schema/worker-admission.test.ts @@ -46,6 +46,7 @@ const connectParams = { environmentId: "worker-1", credential, sessionId: null, + runId: null, ownerEpoch: 1, rpcSetVersion: WORKER_RPC_SET_VERSION, handshake, @@ -190,6 +191,29 @@ describe("worker protocol schemas", () => { params: connectParams, }), ).toBe(true); + const missingRunId = structuredClone(connectParams); + Reflect.deleteProperty(missingRunId.admission, "runId"); + expect( + validateWorkerConnectRequestFrame({ + type: "req", + id: "connect-missing-run", + method: "connect", + params: missingRunId, + }), + ).toBe(false); + for (const admission of [ + { ...connectParams.admission, sessionId: null, runId: "run-1" }, + { ...connectParams.admission, sessionId: "session-1", runId: null }, + ]) { + expect( + validateWorkerConnectRequestFrame({ + type: "req", + id: "connect-mismatched-session-run", + method: "connect", + params: { ...connectParams, admission }, + }), + ).toBe(false); + } expect( Value.Check(WorkerAdmissionResponseFrameSchema, { type: "res", @@ -449,6 +473,7 @@ describe("worker protocol schemas", () => { it("keeps worker close reasons closed", () => { expect(Value.Check(WorkerProtocolCloseReasonSchema, "credential-replaced")).toBe(true); + expect(Value.Check(WorkerProtocolCloseReasonSchema, "placement-mismatch")).toBe(true); expect(Value.Check(WorkerProtocolCloseReasonSchema, "not-a-worker-reason")).toBe(false); }); }); diff --git a/packages/gateway-protocol/src/schema/worker-admission.ts b/packages/gateway-protocol/src/schema/worker-admission.ts index 62ab4aeca2ae..22c1a5317dde 100644 --- a/packages/gateway-protocol/src/schema/worker-admission.ts +++ b/packages/gateway-protocol/src/schema/worker-admission.ts @@ -58,6 +58,27 @@ export const WorkerAdmissionHandshakeSchema = closedObject({ }), }); +const WorkerConnectAdmissionCommonProperties = { + environmentId: WorkerIdentifierSchema, + credential: WorkerCredentialSchema, + ownerEpoch: Type.Integer({ minimum: 0, maximum: Number.MAX_SAFE_INTEGER }), + rpcSetVersion: Type.Integer({ minimum: 1, maximum: Number.MAX_SAFE_INTEGER }), + handshake: WorkerAdmissionHandshakeSchema, +}; + +const WorkerConnectAdmissionSchema = Type.Union([ + closedObject({ + ...WorkerConnectAdmissionCommonProperties, + sessionId: Type.Null(), + runId: Type.Null(), + }), + closedObject({ + ...WorkerConnectAdmissionCommonProperties, + sessionId: WorkerIdentifierSchema, + runId: WorkerIdentifierSchema, + }), +]); + /** Dedicated first-frame payload accepted only on the worker ingress. */ export const WorkerConnectParamsSchema = closedObject({ minProtocol: Type.Integer({ minimum: 1 }), @@ -69,14 +90,7 @@ export const WorkerConnectParamsSchema = closedObject({ mode: Type.Literal(GATEWAY_CLIENT_MODES.WORKER), }), role: Type.Literal("worker"), - admission: closedObject({ - environmentId: WorkerIdentifierSchema, - credential: WorkerCredentialSchema, - sessionId: Type.Union([WorkerIdentifierSchema, Type.Null()]), - ownerEpoch: Type.Integer({ minimum: 0, maximum: Number.MAX_SAFE_INTEGER }), - rpcSetVersion: Type.Integer({ minimum: 1, maximum: Number.MAX_SAFE_INTEGER }), - handshake: WorkerAdmissionHandshakeSchema, - }), + admission: WorkerConnectAdmissionSchema, }); export const WorkerConnectRequestFrameSchema = closedObject({ @@ -94,6 +108,7 @@ export const WorkerAdmissionFailureReasonSchema = Type.Union([ Type.Literal("bundle-mismatch"), Type.Literal("version-mismatch"), Type.Literal("session-mismatch"), + Type.Literal("placement-mismatch"), Type.Literal("owner-epoch-mismatch"), Type.Literal("rpc-set-mismatch"), Type.Literal("protocol-features-mismatch"), diff --git a/src/agents/command/attempt-execution.ts b/src/agents/command/attempt-execution.ts index 0b63af2dea68..e75ac23bdae0 100644 --- a/src/agents/command/attempt-execution.ts +++ b/src/agents/command/attempt-execution.ts @@ -66,6 +66,7 @@ import type { AgentRunSessionTarget } from "../run-session-target.js"; import { resolveAgentRunAbortLifecycleFields } from "../run-termination.js"; import { buildAgentRuntimeAuthPlan } from "../runtime-plan/auth.js"; import type { AgentMessage } from "../runtime/index.js"; +import { withLocalSessionPlacementTurnAdmission } from "../session-placement-admission.js"; import { buildUsageWithNoCost } from "../stream-message-shared.js"; import { buildClaudeCliFallbackContextPrelude, @@ -715,127 +716,139 @@ export function runAgentAttempt(params: { ...mutableCliSessionStore, } : undefined; - return runCliAgent({ - sessionId: params.sessionId, - sessionKey: params.sessionKey, - sessionEntry: params.sessionEntry, - agentId: params.sessionAgentId, - trigger: "user", - sessionFile: params.sessionFile, - storePath: params.storePath, - workspaceDir: params.workspaceDir, - cwd: params.cwd, - config: params.cfg, - prompt: cliPrompt, - transcriptPrompt: params.transcriptBody, - modelProvider: params.providerOverride, - provider: cliExecutionProvider, - model: params.modelOverride, - thinkLevel: params.resolvedThinkLevel, - timeoutMs: params.timeoutMs, - runTimeoutOverrideMs: params.runTimeoutOverrideMs, - runId: params.runId, - lifecycleGeneration: params.lifecycleGeneration, - lane: params.opts.lane, - extraSystemPrompt: params.opts.extraSystemPrompt, - inputProvenance: params.opts.inputProvenance, - sourceReplyDeliveryMode: params.opts.sourceReplyDeliveryMode, - requireExplicitMessageTarget: - params.opts.requireExplicitMessageTarget ?? isSubagentSessionKey(params.sessionKey), - cliSessionBindingFacts: params.opts.cliSessionBindingFacts, - cliSessionId: nextCliSessionId, - cliSessionBinding: - nextCliSessionId === activeCliSessionBinding?.sessionId - ? activeCliSessionBinding - : undefined, - forkCliSessionOnResume, - ...(forkStoreParams - ? { - claimCliSessionFork: async () => { - const claimed = await consumeCliSessionForkInStore(forkStoreParams); - if (claimed) { - params.sessionEntry = claimed; - } - return Boolean(claimed); - }, - restoreCliSessionFork: async () => { - const restored = await restoreCliSessionForkInStore(forkStoreParams); - if (restored) { - params.sessionEntry = restored; - } - }, - persistCliSessionForkSuccessor: async (successorCliSessionId: string) => { - const persisted = await persistCliSessionForkSuccessorInStore({ - ...forkStoreParams, - successorCliSessionId, - }); - if (!persisted) { - throw new Error("CLI session fork successor could not be persisted"); - } - params.sessionEntry = persisted; - }, - } - : {}), - authProfileId, - bootstrapPromptWarningSignaturesSeen, - bootstrapPromptWarningSignature, - // Image discovery must use the original turn, before retry/history decoration. - imagePrompt: params.body, - // Fallback prompts repeat the current task, so prompt-local images must - // accompany every CLI process. Native dedupe requires a runtime receipt. - images: params.opts.images, - imageOrder: params.opts.imageOrder, - skillsSnapshot: params.skillsSnapshot, - messageChannel: params.messageChannel, - streamParams: params.opts.streamParams, - messageProvider: params.opts.messageProvider ?? params.messageChannel, - currentChannelId: params.runContext.currentChannelId, - chatId: params.runContext.chatId, - channelContext: params.runContext.channelContext, - currentThreadTs: params.runContext.currentThreadTs, - currentInboundAudio: params.runContext.currentInboundAudio, - approvalReviewerDeviceId: params.opts.approvalReviewerDeviceId, - agentAccountId: params.runContext.accountId, - senderId: params.runContext.senderId, - senderIsOwner: params.opts.senderIsOwner, - bashElevated: params.opts.bashElevated, - groupId: params.runContext.groupId, - groupChannel: params.runContext.groupChannel, - groupSpace: params.runContext.groupSpace, - spawnedBy: params.spawnedBy, - toolsAllow: resolveCliRuntimeToolsAllow( - params.opts.toolsAllow, - params.opts.toolsAllowIsDefault, - ), - cleanupBundleMcpOnRunEnd: params.opts.cleanupBundleMcpOnRunEnd, - cleanupCliLiveSessionOnRunEnd: params.opts.cleanupCliLiveSessionOnRunEnd, - oneShotCliRun: params.opts.oneShotCliRun, - userTurnTranscriptRecorder: params.userTurnTranscriptRecorder, - suppressNextUserMessagePersistence: params.suppressPromptPersistenceOnRetry === true, - ...(mutableCliSessionStore && !forkCliSessionOnResume - ? { - onBeforeFreshCliSessionRetry: async (retry) => { - if ( - hasNewGeneratedMediaTaskForSessionKey(params.sessionKey, mediaTaskIdsBefore) || - retry.sessionId !== activeCliSessionBinding?.sessionId - ) { - return false; + return withLocalSessionPlacementTurnAdmission( + { + sessionId: params.sessionId, + sessionKey: params.sessionKey ?? params.sessionId, + agentId: params.sessionAgentId, + runId: params.runId, + }, + () => + runCliAgent({ + sessionId: params.sessionId, + sessionKey: params.sessionKey, + sessionEntry: params.sessionEntry, + agentId: params.sessionAgentId, + trigger: "user", + sessionFile: params.sessionFile, + storePath: params.storePath, + workspaceDir: params.workspaceDir, + cwd: params.cwd, + config: params.cfg, + prompt: cliPrompt, + transcriptPrompt: params.transcriptBody, + modelProvider: params.providerOverride, + provider: cliExecutionProvider, + model: params.modelOverride, + thinkLevel: params.resolvedThinkLevel, + timeoutMs: params.timeoutMs, + runTimeoutOverrideMs: params.runTimeoutOverrideMs, + runId: params.runId, + lifecycleGeneration: params.lifecycleGeneration, + lane: params.opts.lane, + extraSystemPrompt: params.opts.extraSystemPrompt, + inputProvenance: params.opts.inputProvenance, + sourceReplyDeliveryMode: params.opts.sourceReplyDeliveryMode, + requireExplicitMessageTarget: + params.opts.requireExplicitMessageTarget ?? isSubagentSessionKey(params.sessionKey), + cliSessionBindingFacts: params.opts.cliSessionBindingFacts, + cliSessionId: nextCliSessionId, + cliSessionBinding: + nextCliSessionId === activeCliSessionBinding?.sessionId + ? activeCliSessionBinding + : undefined, + forkCliSessionOnResume, + ...(forkStoreParams + ? { + claimCliSessionFork: async () => { + const claimed = await consumeCliSessionForkInStore(forkStoreParams); + if (claimed) { + params.sessionEntry = claimed; + } + return Boolean(claimed); + }, + restoreCliSessionFork: async () => { + const restored = await restoreCliSessionForkInStore(forkStoreParams); + if (restored) { + params.sessionEntry = restored; + } + }, + persistCliSessionForkSuccessor: async (successorCliSessionId: string) => { + const persisted = await persistCliSessionForkSuccessorInStore({ + ...forkStoreParams, + successorCliSessionId, + }); + if (!persisted) { + throw new Error("CLI session fork successor could not be persisted"); + } + params.sessionEntry = persisted; + }, } + : {}), + authProfileId, + bootstrapPromptWarningSignaturesSeen, + bootstrapPromptWarningSignature, + // Image discovery must use the original turn, before retry/history decoration. + imagePrompt: params.body, + // Fallback prompts repeat the current task, so prompt-local images must + // accompany every CLI process. Native dedupe requires a runtime receipt. + images: params.opts.images, + imageOrder: params.opts.imageOrder, + skillsSnapshot: params.skillsSnapshot, + messageChannel: params.messageChannel, + streamParams: params.opts.streamParams, + messageProvider: params.opts.messageProvider ?? params.messageChannel, + currentChannelId: params.runContext.currentChannelId, + chatId: params.runContext.chatId, + channelContext: params.runContext.channelContext, + currentThreadTs: params.runContext.currentThreadTs, + currentInboundAudio: params.runContext.currentInboundAudio, + approvalReviewerDeviceId: params.opts.approvalReviewerDeviceId, + agentAccountId: params.runContext.accountId, + senderId: params.runContext.senderId, + senderIsOwner: params.opts.senderIsOwner, + bashElevated: params.opts.bashElevated, + groupId: params.runContext.groupId, + groupChannel: params.runContext.groupChannel, + groupSpace: params.runContext.groupSpace, + spawnedBy: params.spawnedBy, + toolsAllow: resolveCliRuntimeToolsAllow( + params.opts.toolsAllow, + params.opts.toolsAllowIsDefault, + ), + cleanupBundleMcpOnRunEnd: params.opts.cleanupBundleMcpOnRunEnd, + cleanupCliLiveSessionOnRunEnd: params.opts.cleanupCliLiveSessionOnRunEnd, + oneShotCliRun: params.opts.oneShotCliRun, + userTurnTranscriptRecorder: params.userTurnTranscriptRecorder, + suppressNextUserMessagePersistence: params.suppressPromptPersistenceOnRetry === true, + ...(mutableCliSessionStore && !forkCliSessionOnResume + ? { + onBeforeFreshCliSessionRetry: async (retry) => { + if ( + hasNewGeneratedMediaTaskForSessionKey( + params.sessionKey, + mediaTaskIdsBefore, + ) || + retry.sessionId !== activeCliSessionBinding?.sessionId + ) { + return false; + } - log.warn( - `CLI session failed, clearing before fresh retry: provider=${sanitizeForLog(cliExecutionProvider)} sessionKey=${mutableCliSessionStore.sessionKey} reason=${sanitizeForLog(retry.reason)}`, - ); + log.warn( + `CLI session failed, clearing before fresh retry: provider=${sanitizeForLog(cliExecutionProvider)} sessionKey=${mutableCliSessionStore.sessionKey} reason=${sanitizeForLog(retry.reason)}`, + ); - params.sessionEntry = - (await clearCliSessionInStore({ - provider: cliExecutionProvider, - ...mutableCliSessionStore, - })) ?? params.sessionEntry; - return true; - }, - } - : {}), - }); + params.sessionEntry = + (await clearCliSessionInStore({ + provider: cliExecutionProvider, + ...mutableCliSessionStore, + })) ?? params.sessionEntry; + return true; + }, + } + : {}), + }), + ); }; return resolveReusableCliSessionBinding().then(async (activeCliSessionBinding) => { try { diff --git a/src/agents/embedded-agent-runner/run.ts b/src/agents/embedded-agent-runner/run.ts index 9b11956dc115..18499f4b4e53 100644 --- a/src/agents/embedded-agent-runner/run.ts +++ b/src/agents/embedded-agent-runner/run.ts @@ -168,6 +168,7 @@ import { import type { AgentRuntimePlan } from "../runtime-plan/types.js"; import type { AgentRuntimeAuthPlan } from "../runtime-plan/types.js"; import { ensureRuntimePluginsLoaded } from "../runtime-plugins.js"; +import { withSessionPlacementTurnAdmission } from "../session-placement-admission.js"; import { resolveSessionSuspensionReason, resolveSessionSuspensionTarget, @@ -889,12 +890,15 @@ async function runEmbeddedAgentInternal( }); } }; - const enqueueGlobal = (task: () => Promise, opts?: CommandQueueEnqueueOptions) => { + const enqueueGlobal = ( + task: () => Promise, + opts?: CommandQueueEnqueueOptions, + ) => { const globalOpts: CommandQueueEnqueueOptions = { ...opts, priority: sessionQueuePriority, }; - const taskWithCurrentLifecycle = () => { + const taskWithCurrentLifecycle = async () => { params.onLaneWait?.({ waitMs: 0, queuedAhead: 0, waiting: false }); throwIfAborted(); const currentLifecycleGeneration = getAgentEventLifecycleGeneration(); @@ -914,16 +918,29 @@ async function runEmbeddedAgentInternal( lifecycleGeneration = currentLifecycleGeneration; params = { ...params, lifecycleGeneration }; } - // Queue waits can outlive the durable harness binding that admitted a run. - // Recheck only after lifecycle admission, before any run context or hook can execute. + // Queue waits can outlive durable harness and placement bindings. + // Recheck and claim only after lifecycle admission, before context or hooks execute. assertAgentHarnessRunAdmission(params); - claimAgentRunContext(params.runId, { - ...existingContext, - sessionKey: params.sessionKey ?? existingContext?.sessionKey, - sessionId: params.sessionId ?? existingContext?.sessionId, - lifecycleGeneration, - }); - return withAgentRunLifecycleGeneration(lifecycleGeneration, task); + return await withAgentRunLifecycleGeneration(lifecycleGeneration, () => + withSessionPlacementTurnAdmission( + { + sessionId: params.sessionId, + ...(params.agentId ? { agentId: params.agentId } : {}), + ...(params.sessionKey ? { sessionKey: params.sessionKey } : {}), + runId: params.runId, + }, + params, + () => { + claimAgentRunContext(params.runId, { + ...existingContext, + sessionKey: params.sessionKey ?? existingContext?.sessionKey, + sessionId: params.sessionId ?? existingContext?.sessionId, + lifecycleGeneration, + }); + return task(); + }, + ), + ); }; if (params.enqueue) { return params.enqueue(taskWithCurrentLifecycle, withLaneTimeout(withRunLaneWait(globalOpts))); diff --git a/src/agents/session-placement-admission.test.ts b/src/agents/session-placement-admission.test.ts new file mode 100644 index 000000000000..8c49466309d2 --- /dev/null +++ b/src/agents/session-placement-admission.test.ts @@ -0,0 +1,191 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { + installSessionPlacementAdmissionProvider, + installSessionPlacementResetGuard, + resolveSessionPlacementResetBlock, + type LocalTurnPlacementClaim, + type SessionPlacementAdmissionProvider, + withLocalSessionPlacementTurnAdmission, + withSessionPlacementTurnAdmission, +} from "./session-placement-admission.js"; + +let uninstallProvider: (() => void) | undefined; +let uninstallResetGuard: (() => void) | undefined; +const executeLocalTurn: SessionPlacementAdmissionProvider["executeLocalTurn"] = async ( + _claim, + runLocal, +) => await runLocal(); + +afterEach(() => { + uninstallProvider?.(); + uninstallProvider = undefined; + uninstallResetGuard?.(); + uninstallResetGuard = undefined; +}); + +describe("local turn placement admission", () => { + const turnParams = { + sessionId: "session-1", + sessionFile: "/tmp/session-1.jsonl", + workspaceDir: "/tmp/workspace", + prompt: "test", + timeoutMs: 1_000, + runId: "run-1", + }; + + it("delegates the final turn decision to the installed provider", async () => { + const events: string[] = []; + uninstallProvider = installSessionPlacementAdmissionProvider({ + executeLocalTurn, + executeTurn: async (claim, params, runLocal) => { + events.push("claim"); + expect(claim).toEqual({ + sessionId: "session-1", + sessionKey: "agent:main:main", + runId: "run-1", + }); + expect(params).toBe(turnParams); + const result = await runLocal(); + events.push("release"); + return result; + }, + }); + + const result = await withSessionPlacementTurnAdmission( + { + sessionId: "session-1", + sessionKey: "agent:main:main", + runId: "run-1", + }, + turnParams, + async () => { + events.push("turn"); + return { meta: { durationMs: 1 } }; + }, + ); + + expect(result.meta.durationMs).toBe(1); + expect(events).toEqual(["claim", "turn", "release"]); + }); + + it("does not start a local turn when the provider routes remotely", async () => { + const turn = vi.fn(async () => ({ meta: { durationMs: 1 } })); + const executeTurn = vi.fn(async () => ({ + payloads: [{ text: "remote" }], + meta: { durationMs: 2 }, + })); + uninstallProvider = installSessionPlacementAdmissionProvider({ + executeLocalTurn, + executeTurn, + }); + + const result = await withSessionPlacementTurnAdmission( + { sessionId: "session-2", runId: "run-2" }, + { ...turnParams, sessionId: "session-2", runId: "run-2" }, + turn, + ); + expect(result.payloads).toEqual([{ text: "remote" }]); + expect(executeTurn).toHaveBeenCalledOnce(); + expect(executeTurn.mock.calls[0]?.[0]).toEqual({ sessionId: "session-2", runId: "run-2" }); + expect(turn).not.toHaveBeenCalled(); + }); + + it("does not resurrect a replaced provider during uninstall", async () => { + const firstClaim = vi.fn( + async (_claim, _params, runLocal: () => Promise<{ meta: { durationMs: number } }>) => + await runLocal(), + ); + const uninstallFirst = installSessionPlacementAdmissionProvider({ + executeLocalTurn, + executeTurn: firstClaim, + }); + const secondClaim = vi.fn( + async (_claim, _params, runLocal: () => Promise<{ meta: { durationMs: number } }>) => + await runLocal(), + ); + const uninstallSecond = installSessionPlacementAdmissionProvider({ + executeLocalTurn, + executeTurn: secondClaim, + }); + uninstallProvider = uninstallSecond; + + uninstallFirst(); + await withSessionPlacementTurnAdmission( + { sessionId: "session-4", runId: "run-4" }, + { ...turnParams, sessionId: "session-4", runId: "run-4" }, + async () => ({ meta: { durationMs: 1 } }), + ); + expect(firstClaim).not.toHaveBeenCalled(); + expect(secondClaim).toHaveBeenCalledOnce(); + + uninstallSecond(); + uninstallProvider = undefined; + await withSessionPlacementTurnAdmission( + { sessionId: "session-5", runId: "run-5" }, + { ...turnParams, sessionId: "session-5", runId: "run-5" }, + async () => ({ meta: { durationMs: 1 } }), + ); + expect(firstClaim).not.toHaveBeenCalled(); + expect(secondClaim).toHaveBeenCalledOnce(); + }); + + it("delegates generic local execution through the placement gate", async () => { + const events: string[] = []; + uninstallProvider = installSessionPlacementAdmissionProvider({ + async executeLocalTurn( + claim: LocalTurnPlacementClaim, + runLocal: () => Promise, + ): Promise { + events.push("claim"); + expect(claim).toEqual({ + sessionId: "session-cli", + sessionKey: "agent:main:cli", + agentId: "main", + runId: "run-cli", + }); + const result = await runLocal(); + events.push("release"); + return result; + }, + executeTurn: async (_claim, _params, runLocal) => await runLocal(), + }); + + const result = await withLocalSessionPlacementTurnAdmission( + { + sessionId: "session-cli", + sessionKey: "agent:main:cli", + agentId: "main", + runId: "run-cli", + }, + async () => { + events.push("turn"); + return { kind: "cli", code: 0 } as const; + }, + ); + + expect(result).toEqual({ kind: "cli", code: 0 }); + expect(events).toEqual(["claim", "turn", "release"]); + }); +}); + +describe("session placement reset guard", () => { + it("returns the installed reset block", () => { + uninstallResetGuard = installSessionPlacementResetGuard((sessionId) => + sessionId === "session-worker" ? "cloud worker placement is active" : undefined, + ); + + expect(resolveSessionPlacementResetBlock("session-worker")).toBe( + "cloud worker placement is active", + ); + expect(resolveSessionPlacementResetBlock("session-local")).toBeUndefined(); + }); + + it("does not clear a replacement reset guard during stale uninstall", () => { + const uninstallFirst = installSessionPlacementResetGuard(() => "first"); + uninstallResetGuard = installSessionPlacementResetGuard(() => "second"); + + uninstallFirst(); + + expect(resolveSessionPlacementResetBlock("session-worker")).toBe("second"); + }); +}); diff --git a/src/agents/session-placement-admission.ts b/src/agents/session-placement-admission.ts new file mode 100644 index 000000000000..1c875fe5ffd3 --- /dev/null +++ b/src/agents/session-placement-admission.ts @@ -0,0 +1,82 @@ +import { resolveGlobalSingleton } from "../shared/global-singleton.js"; +import type { RunEmbeddedAgentParams } from "./embedded-agent-runner/run/params.js"; +import type { EmbeddedAgentRunResult } from "./embedded-agent-runner/types.js"; + +export type LocalTurnPlacementClaim = { + sessionId: string; + agentId?: string; + sessionKey?: string; + runId: string; +}; + +export type SessionPlacementTurnParams = RunEmbeddedAgentParams & { sessionFile: string }; + +export type SessionPlacementAdmissionProvider = { + executeLocalTurn: (claim: LocalTurnPlacementClaim, runLocal: () => Promise) => Promise; + executeTurn: ( + claim: LocalTurnPlacementClaim, + params: SessionPlacementTurnParams, + runLocal: () => Promise, + ) => Promise; +}; + +type SessionPlacementResetGuard = (sessionId: string) => string | undefined; + +type SessionPlacementAdmissionState = { + provider?: SessionPlacementAdmissionProvider; + resetGuard?: SessionPlacementResetGuard; +}; + +// Runtime chunks share one provider. The identity guard keeps an older gateway +// shutdown from clearing a newer lifecycle's admission gate. +const state = resolveGlobalSingleton( + Symbol.for("openclaw.sessionPlacementAdmissionState"), + (): SessionPlacementAdmissionState => ({}), +); + +export function installSessionPlacementAdmissionProvider( + provider: SessionPlacementAdmissionProvider, +): () => void { + state.provider = provider; + return () => { + if (state.provider === provider) { + state.provider = undefined; + } + }; +} + +export function installSessionPlacementResetGuard(guard: SessionPlacementResetGuard): () => void { + state.resetGuard = guard; + return () => { + if (state.resetGuard === guard) { + state.resetGuard = undefined; + } + }; +} + +export function resolveSessionPlacementResetBlock(sessionId: string): string | undefined { + return state.resetGuard?.(sessionId); +} + +export async function withSessionPlacementTurnAdmission( + claim: LocalTurnPlacementClaim, + params: SessionPlacementTurnParams, + task: () => Promise, +): Promise { + const provider = state.provider; + if (!provider) { + return await task(); + } + return await provider.executeTurn(claim, params, task); +} + +export async function withLocalSessionPlacementTurnAdmission( + claim: LocalTurnPlacementClaim, + task: () => Promise, +): Promise { + const provider = state.provider; + if (!provider) { + return await task(); + } + return await provider.executeLocalTurn(claim, task); +} diff --git a/src/auto-reply/reply/agent-runner-execution.ts b/src/auto-reply/reply/agent-runner-execution.ts index 9d58636ffaea..1be5723ece3d 100644 --- a/src/auto-reply/reply/agent-runner-execution.ts +++ b/src/auto-reply/reply/agent-runner-execution.ts @@ -70,6 +70,7 @@ import { resolveAgentRunErrorLifecycleFields, } from "../../agents/run-termination.js"; import { buildAgentRuntimeOutcomePlan } from "../../agents/runtime-plan/build.js"; +import { withLocalSessionPlacementTurnAdmission } from "../../agents/session-placement-admission.js"; import { resolveSessionRuntimeOverrideForProvider } from "../../agents/session-runtime-compat.js"; import { resolveCandidateThinkingLevel } from "../../agents/thinking-runtime.js"; import { resolveGroupSessionKey, type SessionEntry } from "../../config/sessions.js"; @@ -2080,196 +2081,210 @@ async function runAgentTurnWithFallbackInternal( }, }); const result = await agentTurnTiming.measure("cli_run", () => - runCliAgentWithLifecycle({ - runId, - lifecycleGeneration, - provider: cliExecutionProvider, - startedAt: cliLifecycleStartedAt, - emitLifecycleTerminal: false, - onAgentRunStart: notifyAgentRunStart, - suppressAssistantBridge: params.followupRun.run.silentExpected, - onActivity: () => params.replyOperation?.recordActivity(), - preserveProgressCallbackStartOrder, - onAssistantText: async (text) => { - if (!preserveProgressCallbackStartOrder) { - const textForTyping = await handlePartialForTyping({ - text, - } as ReplyPayload); - if (textForTyping === undefined || !params.opts?.onPartialReply) { - return; - } - await params.opts.onPartialReply({ text: textForTyping }); - return; - } - const textForTyping = preparePartialForTyping({ text } as ReplyPayload); - if (textForTyping === undefined) { - return; - } - // Assistant and tool CLI bridges drain independently; stage presentation - // before typing I/O so a later tool cannot overtake this text. - await startPresentationWhileTyping( - params.typingSignals.signalTextDelta(textForTyping), - () => params.opts?.onPartialReply?.({ text: textForTyping }), - ); - }, - onReasoningText: createCliReasoningStreamBridge(params.opts?.onReasoningStream), - onReasoningProgress: async (payload) => { - await params.opts?.onReasoningProgress?.(payload); - }, - onToolEvent: async (payload) => { - if (!preserveProgressCallbackStartOrder) { - await cliToolSummaryTracker.noteToolEvent(payload); - if (payload.phase === "result") { - return; - } - const { name, phase, args } = payload; - await Promise.all([ - params.typingSignals.signalToolStart(), - params.opts?.onToolStart?.({ - name, - phase, - args, - detailMode: params.toolProgressDetail, - }), - ]); - return; - } - const summaryPromise = cliToolSummaryTracker.noteToolEvent(payload); - if (payload.phase === "result") { - await summaryPromise; - return; - } - const { name, phase, args } = payload; - // Tool and assistant CLI bridges drain independently. Start channel - // presentation before either bridge can yield and invert source order. - await Promise.all([ - summaryPromise, - startPresentationWhileTyping(params.typingSignals.signalToolStart(), () => - params.opts?.onToolStart?.({ - name, - phase, - args, - detailMode: params.toolProgressDetail, - }), - ), - ]); - }, - onCommentaryText: - params.opts?.commentaryProgressEnabled === true && params.opts.onItemEvent - ? async (payload) => { - await params.opts?.onItemEvent?.({ - itemId: payload.itemId, - kind: "preamble", - progressText: payload.text, - }); - } - : undefined, - onFastModeAutoProgress: async (payload) => { - await params.opts?.onToolResult?.(payload); - }, - transformResult: - params.followupRun.currentInboundEventKind === "room_event" - ? (resultLocal) => - keepCliSessionBindingOnlyWhenReused({ - result: resultLocal, - existingSessionId: cliSessionBinding?.sessionId, - onDroppedReplacement: () => { - droppedCliSessionReplacement = true; - }, - }) - : undefined, - runParams: { + withLocalSessionPlacementTurnAdmission( + { sessionId: params.followupRun.run.sessionId, sessionKey: params.sessionKey, - runtimePolicySessionKey: - params.followupRun.run.runtimePolicySessionKey ?? - params.runtimePolicySessionKey, agentId: params.followupRun.run.agentId, - trigger: params.isHeartbeat ? "heartbeat" : "user", - sessionFile: params.followupRun.run.sessionFile, - workspaceDir: params.followupRun.run.workspaceDir, - cwd: params.followupRun.run.cwd, - config: runtimeConfig, - prompt: params.commandBody, - transcriptPrompt: params.transcriptCommandBody, - suppressNextUserMessagePersistence: suppressQueuedUserPersistenceForCandidate, - userTurnTranscriptRecorder, - onUserMessagePersisted: notifyUserMessagePersisted, - persistAssistantTranscript: - params.followupRun.currentInboundEventKind !== "room_event" && - params.followupRun.run.suppressTranscriptOnlyAssistantPersistence !== true, - storePath: params.storePath, - currentInboundEventKind: params.followupRun.currentInboundEventKind, - currentInboundContext: params.followupRun.currentInboundContext, - inputProvenance: params.followupRun.run.inputProvenance, - modelProvider: provider, - provider: cliExecutionProvider, - execOverrides: params.followupRun.run.execOverrides, - bashElevated: params.followupRun.run.bashElevated, - model, - thinkLevel: candidateThinkLevel, - fastMode: candidateFastMode.fastMode, - fastModeStartedAtMs, - fastModeAutoOnSeconds: candidateFastMode.fastModeAutoOnSeconds, - fastModeAutoProgressState, - isFinalFallbackAttempt: runOptions?.isFinalFallbackAttempt, - timeoutMs: params.followupRun.run.timeoutMs, - runTimeoutOverrideMs: params.followupRun.run.runTimeoutOverrideMs, runId, - lane: runLane, - extraSystemPrompt: params.followupRun.run.extraSystemPrompt, - sourceReplyDeliveryMode: params.followupRun.run.sourceReplyDeliveryMode, - taskSuggestionDeliveryMode: params.followupRun.run.taskSuggestionDeliveryMode, - silentReplyPromptMode: params.followupRun.run.silentReplyPromptMode, - allowEmptyAssistantReplyAsSilent: - params.followupRun.run.allowEmptyAssistantReplyAsSilent, - extraSystemPromptStatic: params.followupRun.run.extraSystemPromptStatic, - cliSessionBindingFacts: params.followupRun.run.cliSessionBindingFacts, - ownerNumbers: params.followupRun.run.ownerNumbers, - cliSessionId: cliSessionBinding?.sessionId, - cliSessionBinding, - authProfileId: authProfile.authProfileId, - bootstrapContextMode: params.opts?.bootstrapContextMode, - bootstrapContextRunKind, - bootstrapPromptWarningSignaturesSeen, - bootstrapPromptWarningSignature: - bootstrapPromptWarningSignaturesSeen[ - bootstrapPromptWarningSignaturesSeen.length - 1 - ], - images: currentTurnImages.images, - imageOrder: currentTurnImages.imageOrder, - skillsSnapshot: params.followupRun.run.skillsSnapshot, - messageChannel: params.followupRun.originatingChannel ?? undefined, - messageProvider: hookMessageProvider, - clientCaps: params.followupRun.run.clientCaps, - currentChannelId: - params.followupRun.originatingTo ?? - params.sessionCtx.OriginatingTo ?? - params.sessionCtx.To, - senderId: params.followupRun.run.senderId, - senderName: params.followupRun.run.senderName, - senderUsername: params.followupRun.run.senderUsername, - senderE164: params.followupRun.run.senderE164, - groupId: params.followupRun.run.groupId, - groupChannel: params.followupRun.run.groupChannel, - groupSpace: params.followupRun.run.groupSpace, - spawnedBy: params.followupRun.run.spawnedBy, - chatId: params.followupRun.originatingChatId, - channelContext: params.followupRun.run.channelContext, - currentThreadTs: - cliCurrentThreadId != null ? String(cliCurrentThreadId) : undefined, - currentMessageId: cliCurrentMessageId, - currentInboundAudio: hasInboundAudio(params.sessionCtx), - agentAccountId: params.followupRun.run.agentAccountId, - senderIsOwner: params.followupRun.run.senderIsOwner, - approvalReviewerDeviceId: params.followupRun.run.approvalReviewerDeviceId, - toolsAllow: params.opts?.toolsAllow, - disableTools: params.opts?.disableTools, - abortSignal: runAbortSignal, - onExecutionPhase: signalExecutionPhaseForTyping, - replyOperation: params.replyOperation, }, - }), + () => + runCliAgentWithLifecycle({ + runId, + lifecycleGeneration, + provider: cliExecutionProvider, + startedAt: cliLifecycleStartedAt, + emitLifecycleTerminal: false, + onAgentRunStart: notifyAgentRunStart, + suppressAssistantBridge: params.followupRun.run.silentExpected, + onActivity: () => params.replyOperation?.recordActivity(), + preserveProgressCallbackStartOrder, + onAssistantText: async (text) => { + if (!preserveProgressCallbackStartOrder) { + const textForTyping = await handlePartialForTyping({ + text, + } as ReplyPayload); + if (textForTyping === undefined || !params.opts?.onPartialReply) { + return; + } + await params.opts.onPartialReply({ text: textForTyping }); + return; + } + const textForTyping = preparePartialForTyping({ text } as ReplyPayload); + if (textForTyping === undefined) { + return; + } + // Assistant and tool CLI bridges drain independently; stage presentation + // before typing I/O so a later tool cannot overtake this text. + await startPresentationWhileTyping( + params.typingSignals.signalTextDelta(textForTyping), + () => params.opts?.onPartialReply?.({ text: textForTyping }), + ); + }, + onReasoningText: createCliReasoningStreamBridge( + params.opts?.onReasoningStream, + ), + onReasoningProgress: async (payload) => { + await params.opts?.onReasoningProgress?.(payload); + }, + onToolEvent: async (payload) => { + if (!preserveProgressCallbackStartOrder) { + await cliToolSummaryTracker.noteToolEvent(payload); + if (payload.phase === "result") { + return; + } + const { name, phase, args } = payload; + await Promise.all([ + params.typingSignals.signalToolStart(), + params.opts?.onToolStart?.({ + name, + phase, + args, + detailMode: params.toolProgressDetail, + }), + ]); + return; + } + const summaryPromise = cliToolSummaryTracker.noteToolEvent(payload); + if (payload.phase === "result") { + await summaryPromise; + return; + } + const { name, phase, args } = payload; + // Tool and assistant CLI bridges drain independently. Start channel + // presentation before either bridge can yield and invert source order. + await Promise.all([ + summaryPromise, + startPresentationWhileTyping(params.typingSignals.signalToolStart(), () => + params.opts?.onToolStart?.({ + name, + phase, + args, + detailMode: params.toolProgressDetail, + }), + ), + ]); + }, + onCommentaryText: + params.opts?.commentaryProgressEnabled === true && params.opts.onItemEvent + ? async (payload) => { + await params.opts?.onItemEvent?.({ + itemId: payload.itemId, + kind: "preamble", + progressText: payload.text, + }); + } + : undefined, + onFastModeAutoProgress: async (payload) => { + await params.opts?.onToolResult?.(payload); + }, + transformResult: + params.followupRun.currentInboundEventKind === "room_event" + ? (resultLocal) => + keepCliSessionBindingOnlyWhenReused({ + result: resultLocal, + existingSessionId: cliSessionBinding?.sessionId, + onDroppedReplacement: () => { + droppedCliSessionReplacement = true; + }, + }) + : undefined, + runParams: { + sessionId: params.followupRun.run.sessionId, + sessionKey: params.sessionKey, + runtimePolicySessionKey: + params.followupRun.run.runtimePolicySessionKey ?? + params.runtimePolicySessionKey, + agentId: params.followupRun.run.agentId, + trigger: params.isHeartbeat ? "heartbeat" : "user", + sessionFile: params.followupRun.run.sessionFile, + workspaceDir: params.followupRun.run.workspaceDir, + cwd: params.followupRun.run.cwd, + config: runtimeConfig, + prompt: params.commandBody, + transcriptPrompt: params.transcriptCommandBody, + suppressNextUserMessagePersistence: + suppressQueuedUserPersistenceForCandidate, + userTurnTranscriptRecorder, + onUserMessagePersisted: notifyUserMessagePersisted, + persistAssistantTranscript: + params.followupRun.currentInboundEventKind !== "room_event" && + params.followupRun.run.suppressTranscriptOnlyAssistantPersistence !== + true, + storePath: params.storePath, + currentInboundEventKind: params.followupRun.currentInboundEventKind, + currentInboundContext: params.followupRun.currentInboundContext, + inputProvenance: params.followupRun.run.inputProvenance, + modelProvider: provider, + provider: cliExecutionProvider, + execOverrides: params.followupRun.run.execOverrides, + bashElevated: params.followupRun.run.bashElevated, + model, + thinkLevel: candidateThinkLevel, + fastMode: candidateFastMode.fastMode, + fastModeStartedAtMs, + fastModeAutoOnSeconds: candidateFastMode.fastModeAutoOnSeconds, + fastModeAutoProgressState, + isFinalFallbackAttempt: runOptions?.isFinalFallbackAttempt, + timeoutMs: params.followupRun.run.timeoutMs, + runTimeoutOverrideMs: params.followupRun.run.runTimeoutOverrideMs, + runId, + lane: runLane, + extraSystemPrompt: params.followupRun.run.extraSystemPrompt, + sourceReplyDeliveryMode: params.followupRun.run.sourceReplyDeliveryMode, + taskSuggestionDeliveryMode: + params.followupRun.run.taskSuggestionDeliveryMode, + silentReplyPromptMode: params.followupRun.run.silentReplyPromptMode, + allowEmptyAssistantReplyAsSilent: + params.followupRun.run.allowEmptyAssistantReplyAsSilent, + extraSystemPromptStatic: params.followupRun.run.extraSystemPromptStatic, + cliSessionBindingFacts: params.followupRun.run.cliSessionBindingFacts, + ownerNumbers: params.followupRun.run.ownerNumbers, + cliSessionId: cliSessionBinding?.sessionId, + cliSessionBinding, + authProfileId: authProfile.authProfileId, + bootstrapContextMode: params.opts?.bootstrapContextMode, + bootstrapContextRunKind, + bootstrapPromptWarningSignaturesSeen, + bootstrapPromptWarningSignature: + bootstrapPromptWarningSignaturesSeen[ + bootstrapPromptWarningSignaturesSeen.length - 1 + ], + images: currentTurnImages.images, + imageOrder: currentTurnImages.imageOrder, + skillsSnapshot: params.followupRun.run.skillsSnapshot, + messageChannel: params.followupRun.originatingChannel ?? undefined, + messageProvider: hookMessageProvider, + clientCaps: params.followupRun.run.clientCaps, + currentChannelId: + params.followupRun.originatingTo ?? + params.sessionCtx.OriginatingTo ?? + params.sessionCtx.To, + senderId: params.followupRun.run.senderId, + senderName: params.followupRun.run.senderName, + senderUsername: params.followupRun.run.senderUsername, + senderE164: params.followupRun.run.senderE164, + groupId: params.followupRun.run.groupId, + groupChannel: params.followupRun.run.groupChannel, + groupSpace: params.followupRun.run.groupSpace, + spawnedBy: params.followupRun.run.spawnedBy, + chatId: params.followupRun.originatingChatId, + channelContext: params.followupRun.run.channelContext, + currentThreadTs: + cliCurrentThreadId != null ? String(cliCurrentThreadId) : undefined, + currentMessageId: cliCurrentMessageId, + currentInboundAudio: hasInboundAudio(params.sessionCtx), + agentAccountId: params.followupRun.run.agentAccountId, + senderIsOwner: params.followupRun.run.senderIsOwner, + approvalReviewerDeviceId: params.followupRun.run.approvalReviewerDeviceId, + toolsAllow: params.opts?.toolsAllow, + disableTools: params.opts?.disableTools, + abortSignal: runAbortSignal, + onExecutionPhase: signalExecutionPhaseForTyping, + replyOperation: params.replyOperation, + }, + }), + ), ); if (droppedCliSessionReplacement) { await clearDroppedCliSessionBinding({ diff --git a/src/auto-reply/reply/followup-runner.ts b/src/auto-reply/reply/followup-runner.ts index 5f5e7dd69331..f48d294fc1ba 100644 --- a/src/auto-reply/reply/followup-runner.ts +++ b/src/auto-reply/reply/followup-runner.ts @@ -40,6 +40,7 @@ import { buildAgentRuntimeDeliveryPlan, buildAgentRuntimeOutcomePlan, } from "../../agents/runtime-plan/build.js"; +import { withLocalSessionPlacementTurnAdmission } from "../../agents/session-placement-admission.js"; import { resolveSessionRuntimeOverrideForProvider } from "../../agents/session-runtime-compat.js"; import { resolveCandidateThinkingLevel } from "../../agents/thinking-runtime.js"; import type { SessionEntry } from "../../config/sessions.js"; @@ -1168,186 +1169,202 @@ export function createFollowupRunner(params: { shouldEmitToolOutput: shouldEmitToolOutputProgress, deliver: deliverFollowupToolSummary, }); - const result = await runCliAgentWithLifecycle({ - runId, - lifecycleGeneration, - provider: cliExecutionProvider, - startedAt: cliLifecycleStartedAt, - emitLifecycleTerminal: false, - onAgentRunStart: () => opts?.onAgentRunStart?.(runId), - suppressAssistantBridge: run.silentExpected, - onActivity: () => replyOperation?.recordActivity(), - preserveProgressCallbackStartOrder, - onReasoningText: createCliReasoningStreamBridge(progressOpts?.onReasoningStream), - onReasoningProgress: async (payload) => { - await progressOpts?.onReasoningProgress?.(payload); + const result = await withLocalSessionPlacementTurnAdmission( + { + sessionId: run.sessionId, + sessionKey: replySessionKey, + agentId: run.agentId, + runId, }, - onToolEvent: async (payload) => { - if (!preserveProgressCallbackStartOrder) { - await cliToolSummaryTracker.noteToolEvent(payload); - if (payload.phase === "result") { - return; - } - await forwardFollowupProgressEvent({ - evt: { - stream: "tool", - data: { name: payload.name, phase: payload.phase, args: payload.args }, - }, - opts: progressOpts, - detailMode: toolProgressDetail, - emitChannelProgress: shouldEmitToolResultProgress(), - }); - return; - } - if (payload.phase === "result") { - await cliToolSummaryTracker.noteToolEvent(payload); - return; - } - // CLI bridges drain independently. Start channel presentation before - // summary bookkeeping can yield and let later progress overtake this tool. - const presentationPromise = forwardFollowupProgressEvent({ - evt: { - stream: "tool", - data: { name: payload.name, phase: payload.phase, args: payload.args }, + () => + runCliAgentWithLifecycle({ + runId, + lifecycleGeneration, + provider: cliExecutionProvider, + startedAt: cliLifecycleStartedAt, + emitLifecycleTerminal: false, + onAgentRunStart: () => opts?.onAgentRunStart?.(runId), + suppressAssistantBridge: run.silentExpected, + onActivity: () => replyOperation?.recordActivity(), + preserveProgressCallbackStartOrder, + onReasoningText: createCliReasoningStreamBridge( + progressOpts?.onReasoningStream, + ), + onReasoningProgress: async (payload) => { + await progressOpts?.onReasoningProgress?.(payload); }, - opts: progressOpts, - detailMode: toolProgressDetail, - emitChannelProgress: shouldEmitToolResultProgress(), - }); - await Promise.all([ - presentationPromise, - cliToolSummaryTracker.noteToolEvent(payload), - ]); - }, - onCommentaryText: - progressOpts?.commentaryProgressEnabled === true && progressOpts.onItemEvent - ? async ({ text, itemId }) => { + onToolEvent: async (payload) => { + if (!preserveProgressCallbackStartOrder) { + await cliToolSummaryTracker.noteToolEvent(payload); + if (payload.phase === "result") { + return; + } await forwardFollowupProgressEvent({ evt: { - stream: "item", - data: { kind: "preamble", progressText: text, itemId }, + stream: "tool", + data: { + name: payload.name, + phase: payload.phase, + args: payload.args, + }, }, opts: progressOpts, detailMode: toolProgressDetail, + emitChannelProgress: shouldEmitToolResultProgress(), }); + return; } - : undefined, - onFastModeAutoProgress: async (payload) => { - await enqueueProgressDelivery(async () => { - // Mirrors direct dispatch progress suppression: ambient - // room events never get automatic fast-mode notices. - if (isRoomEventFollowup()) { - return; - } - await sendRunPayloads( - [payload], - effectiveQueued, - { - provider, - modelId: model, - }, - { kind: "tool", mirror: false, runId }, - ); - }); - }, - transformResult: - queued.currentInboundEventKind === "room_event" - ? (resultLocal) => - keepCliSessionBindingOnlyWhenReused({ - result: resultLocal, - existingSessionId: cliSessionBinding?.sessionId, - onDroppedReplacement: () => { - droppedCliSessionReplacement = true; + if (payload.phase === "result") { + await cliToolSummaryTracker.noteToolEvent(payload); + return; + } + // CLI bridges drain independently. Start channel presentation before + // summary bookkeeping can yield and let later progress overtake this tool. + const presentationPromise = forwardFollowupProgressEvent({ + evt: { + stream: "tool", + data: { name: payload.name, phase: payload.phase, args: payload.args }, + }, + opts: progressOpts, + detailMode: toolProgressDetail, + emitChannelProgress: shouldEmitToolResultProgress(), + }); + await Promise.all([ + presentationPromise, + cliToolSummaryTracker.noteToolEvent(payload), + ]); + }, + onCommentaryText: + progressOpts?.commentaryProgressEnabled === true && progressOpts.onItemEvent + ? async ({ text, itemId }) => { + await forwardFollowupProgressEvent({ + evt: { + stream: "item", + data: { kind: "preamble", progressText: text, itemId }, + }, + opts: progressOpts, + detailMode: toolProgressDetail, + }); + } + : undefined, + onFastModeAutoProgress: async (payload) => { + await enqueueProgressDelivery(async () => { + // Mirrors direct dispatch progress suppression: ambient + // room events never get automatic fast-mode notices. + if (isRoomEventFollowup()) { + return; + } + await sendRunPayloads( + [payload], + effectiveQueued, + { + provider, + modelId: model, }, - }) - : undefined, - runParams: { - replyOperation, - sessionId: run.sessionId, - sessionKey: replySessionKey, - runtimePolicySessionKey: run.runtimePolicySessionKey, - agentId: run.agentId, - trigger: opts?.isHeartbeat === true ? "heartbeat" : "user", - sessionFile: run.sessionFile, - workspaceDir: run.workspaceDir, - cwd: run.cwd, - config: runtimeConfig, - prompt: queued.prompt, - transcriptPrompt: queued.transcriptPrompt, - suppressNextUserMessagePersistence: suppressQueuedUserPersistenceForCandidate, - userTurnTranscriptRecorder, - onUserMessagePersisted: notifyUserMessagePersisted, - persistAssistantTranscript: - queued.currentInboundEventKind !== "room_event" && - run.suppressTranscriptOnlyAssistantPersistence !== true, - storePath, - currentInboundEventKind: queued.currentInboundEventKind, - currentInboundAudio: queued.currentInboundAudio, - currentInboundContext, - inputProvenance: run.inputProvenance, - modelProvider: provider, - provider: cliExecutionProvider, - execOverrides: run.execOverrides, - bashElevated: run.bashElevated, - model, - ...resolveRunAuthProfile(candidateRun, cliExecutionProvider, { - config: runtimeConfig, + { kind: "tool", mirror: false, runId }, + ); + }); + }, + transformResult: + queued.currentInboundEventKind === "room_event" + ? (resultLocal) => + keepCliSessionBindingOnlyWhenReused({ + result: resultLocal, + existingSessionId: cliSessionBinding?.sessionId, + onDroppedReplacement: () => { + droppedCliSessionReplacement = true; + }, + }) + : undefined, + runParams: { + replyOperation, + sessionId: run.sessionId, + sessionKey: replySessionKey, + runtimePolicySessionKey: run.runtimePolicySessionKey, + agentId: run.agentId, + trigger: opts?.isHeartbeat === true ? "heartbeat" : "user", + sessionFile: run.sessionFile, + workspaceDir: run.workspaceDir, + cwd: run.cwd, + config: runtimeConfig, + prompt: queued.prompt, + transcriptPrompt: queued.transcriptPrompt, + suppressNextUserMessagePersistence: + suppressQueuedUserPersistenceForCandidate, + userTurnTranscriptRecorder, + onUserMessagePersisted: notifyUserMessagePersisted, + persistAssistantTranscript: + queued.currentInboundEventKind !== "room_event" && + run.suppressTranscriptOnlyAssistantPersistence !== true, + storePath, + currentInboundEventKind: queued.currentInboundEventKind, + currentInboundAudio: queued.currentInboundAudio, + currentInboundContext, + inputProvenance: run.inputProvenance, + modelProvider: provider, + provider: cliExecutionProvider, + execOverrides: run.execOverrides, + bashElevated: run.bashElevated, + model, + ...resolveRunAuthProfile(candidateRun, cliExecutionProvider, { + config: runtimeConfig, + }), + thinkLevel: candidateThinkLevel, + fastMode: candidateFastMode.fastMode, + fastModeStartedAtMs, + fastModeAutoOnSeconds: candidateFastMode.fastModeAutoOnSeconds, + fastModeAutoProgressState, + isFinalFallbackAttempt: runOptions?.isFinalFallbackAttempt, + timeoutMs: run.timeoutMs, + runTimeoutOverrideMs: run.runTimeoutOverrideMs, + runId, + extraSystemPrompt: run.extraSystemPrompt, + sourceReplyDeliveryMode: run.sourceReplyDeliveryMode, + taskSuggestionDeliveryMode: run.taskSuggestionDeliveryMode, + silentReplyPromptMode: run.silentReplyPromptMode, + allowEmptyAssistantReplyAsSilent: run.allowEmptyAssistantReplyAsSilent, + extraSystemPromptStatic: run.extraSystemPromptStatic, + cliSessionBindingFacts: run.cliSessionBindingFacts, + ownerNumbers: run.ownerNumbers, + cliSessionId: cliSessionBinding?.sessionId, + cliSessionBinding, + bootstrapPromptWarningSignaturesSeen, + bootstrapPromptWarningSignature: + bootstrapPromptWarningSignaturesSeen[ + bootstrapPromptWarningSignaturesSeen.length - 1 + ], + images: queuedImages, + imageOrder: queuedImageOrder, + skillsSnapshot: run.skillsSnapshot, + messageChannel: queued.originatingChannel ?? undefined, + messageProvider: resolveOriginMessageProvider({ + originatingChannel: queued.originatingChannel, + provider: run.messageProvider, + }), + clientCaps: run.clientCaps, + currentChannelId: queued.originatingTo, + senderId: run.senderId, + senderName: run.senderName, + senderUsername: run.senderUsername, + senderE164: run.senderE164, + groupId: run.groupId, + groupChannel: run.groupChannel, + groupSpace: run.groupSpace, + spawnedBy: run.spawnedBy, + chatId: queued.originatingChatId, + channelContext: run.channelContext, + currentThreadTs: + queued.originatingThreadId != null + ? String(queued.originatingThreadId) + : undefined, + currentMessageId: followupCurrentMessageId, + agentAccountId: run.agentAccountId, + senderIsOwner: run.senderIsOwner, + disableTools: opts?.disableTools, + abortSignal: runAbortSignal, + }, }), - thinkLevel: candidateThinkLevel, - fastMode: candidateFastMode.fastMode, - fastModeStartedAtMs, - fastModeAutoOnSeconds: candidateFastMode.fastModeAutoOnSeconds, - fastModeAutoProgressState, - isFinalFallbackAttempt: runOptions?.isFinalFallbackAttempt, - timeoutMs: run.timeoutMs, - runTimeoutOverrideMs: run.runTimeoutOverrideMs, - runId, - extraSystemPrompt: run.extraSystemPrompt, - sourceReplyDeliveryMode: run.sourceReplyDeliveryMode, - taskSuggestionDeliveryMode: run.taskSuggestionDeliveryMode, - silentReplyPromptMode: run.silentReplyPromptMode, - allowEmptyAssistantReplyAsSilent: run.allowEmptyAssistantReplyAsSilent, - extraSystemPromptStatic: run.extraSystemPromptStatic, - cliSessionBindingFacts: run.cliSessionBindingFacts, - ownerNumbers: run.ownerNumbers, - cliSessionId: cliSessionBinding?.sessionId, - cliSessionBinding, - bootstrapPromptWarningSignaturesSeen, - bootstrapPromptWarningSignature: - bootstrapPromptWarningSignaturesSeen[ - bootstrapPromptWarningSignaturesSeen.length - 1 - ], - images: queuedImages, - imageOrder: queuedImageOrder, - skillsSnapshot: run.skillsSnapshot, - messageChannel: queued.originatingChannel ?? undefined, - messageProvider: resolveOriginMessageProvider({ - originatingChannel: queued.originatingChannel, - provider: run.messageProvider, - }), - clientCaps: run.clientCaps, - currentChannelId: queued.originatingTo, - senderId: run.senderId, - senderName: run.senderName, - senderUsername: run.senderUsername, - senderE164: run.senderE164, - groupId: run.groupId, - groupChannel: run.groupChannel, - groupSpace: run.groupSpace, - spawnedBy: run.spawnedBy, - chatId: queued.originatingChatId, - channelContext: run.channelContext, - currentThreadTs: - queued.originatingThreadId != null - ? String(queued.originatingThreadId) - : undefined, - currentMessageId: followupCurrentMessageId, - agentAccountId: run.agentAccountId, - senderIsOwner: run.senderIsOwner, - disableTools: opts?.disableTools, - abortSignal: runAbortSignal, - }, - }); + ); if (droppedCliSessionReplacement) { await clearDroppedCliSessionBinding({ provider: cliExecutionProvider, diff --git a/src/cli/command-catalog.ts b/src/cli/command-catalog.ts index 5455cc7c0117..9238dfb0b693 100644 --- a/src/cli/command-catalog.ts +++ b/src/cli/command-catalog.ts @@ -338,6 +338,7 @@ export const cliCommandCatalog: readonly CliCommandCatalogEntry[] = [ bypassConfigGuard: true, hideBanner: true, loadPlugins: "never", + ownsProtocolStdout: true, networkProxy: "bypass", }, }, diff --git a/src/cli/command-path-policy.test.ts b/src/cli/command-path-policy.test.ts index 0d35d7b38db6..f72e4b474e8a 100644 --- a/src/cli/command-path-policy.test.ts +++ b/src/cli/command-path-policy.test.ts @@ -191,6 +191,7 @@ describe("command-path-policy", () => { bypassConfigGuard: true, loadPlugins: "never", hideBanner: true, + ownsProtocolStdout: true, networkProxy: "bypass", }); expectResolvedPolicy(["configure"], { diff --git a/src/cli/command-startup-policy.test.ts b/src/cli/command-startup-policy.test.ts index 293a32350268..35b14aeefc8d 100644 --- a/src/cli/command-startup-policy.test.ts +++ b/src/cli/command-startup-policy.test.ts @@ -254,6 +254,7 @@ describe("command-startup-policy", () => { expect(shouldBypassConfigGuardForCommandPath(["worker"])).toBe(true); expect(policy.hideBanner).toBe(true); expect(policy.loadPlugins).toBe(false); + expect(policy.suppressDoctorStdout).toBe(true); }); it("suppresses startup stdout for the bare acp protocol", () => { diff --git a/src/cron/isolated-agent/run-executor.ts b/src/cron/isolated-agent/run-executor.ts index 3b79aa7aa42a..8aedf4c4fb40 100644 --- a/src/cron/isolated-agent/run-executor.ts +++ b/src/cron/isolated-agent/run-executor.ts @@ -8,6 +8,7 @@ import { runAgentHarnessBeforeMessageWriteHook } from "../../agents/harness/hook import type { ModelCatalogEntry } from "../../agents/model-catalog.types.js"; import { resolveCliRuntimeExecutionProvider } from "../../agents/model-runtime-aliases.js"; import { wrapUntrustedPromptDataBlock } from "../../agents/sanitize-for-prompt.js"; +import { withLocalSessionPlacementTurnAdmission } from "../../agents/session-placement-admission.js"; import { resolveSessionRuntimeOverrideForProvider } from "../../agents/session-runtime-compat.js"; import type { ThinkLevel, VerboseLevel } from "../../auto-reply/thinking.js"; import type { CliSessionBinding } from "../../config/sessions.js"; @@ -415,55 +416,69 @@ export function createCronPromptExecutor(params: { cliSessionBinding && hasCliSessionReuseMetadata(cliSessionBinding) ? cliSessionBinding : undefined; - const result = await runCliAgent({ - sessionId: params.cronSession.sessionEntry.sessionId, - sessionKey: params.runSessionKey, - sessionEntry: params.cronSession.sessionEntry, - agentId: params.agentId, - trigger: "cron", - jobId: params.job.id, - cleanupCliLiveSessionOnRunEnd: params.usesDetachedRunSession === true, - sessionFile, - workspaceDir: params.workspaceDir, - config: params.cfgWithAgentDefaults, - prompt: modelPrompt, - transcriptPrompt: deliveryTargetRuntimeContext ? promptText : undefined, - modelProvider: providerOverride, - provider: executionProvider, - model: modelOverride, - thinkLevel: candidateThinkLevel, - timeoutMs: params.timeoutMs, - runId: params.cronSession.sessionEntry.sessionId, - lane: resolveCronAgentLane(params.lane), - allowEmptyAssistantReplyAsSilent, - cliSessionId: cliSessionBinding?.sessionId, - cliSessionBinding: guardedCliSessionBinding, - skillsSnapshot: params.skillsSnapshot, - messageChannel, - sourceReplyDeliveryMode, - requireExplicitMessageTarget: sourceDelivery.messageTool.requireExplicitTarget, - cliSessionBindingFacts: { - sourceReplyDeliveryMode, - requireExplicitMessageTarget: sourceDelivery.messageTool.requireExplicitTarget, + // Cron intentionally reuses its durable session id as the run id; turn + // claims stay unique via per-claim ids and the worker gate handles this + // via credential rotation (see worker-environments/service.ts fences). + const runId = params.cronSession.sessionEntry.sessionId; + const result = await withLocalSessionPlacementTurnAdmission( + { + sessionId: params.cronSession.sessionEntry.sessionId, + sessionKey: params.runSessionKey, + agentId: params.agentId, + runId, }, - toolsAllow: resolveCliRuntimeToolsAllow( - params.agentPayload?.toolsAllow, - params.agentPayload?.toolsAllowIsDefault, - ), - abortSignal: params.abortSignal, - onExecutionStarted: params.onExecutionStarted, - onExecutionPhase: params.onExecutionPhase, - bootstrapContextMode, - bootstrapContextRunKind: "cron", - bootstrapPromptWarningSignaturesSeen, - bootstrapPromptWarningSignature, - fastModeStartedAtMs, - fastModeAutoProgressState, - isFinalFallbackAttempt: runOptions?.isFinalFallbackAttempt, - userTurnTranscriptRecorder, - suppressNextUserMessagePersistence: - userTurnTranscriptRecorder.hasPersisted() || userTurnTranscriptRecorder.isBlocked(), - }); + () => + runCliAgent({ + sessionId: params.cronSession.sessionEntry.sessionId, + sessionKey: params.runSessionKey, + sessionEntry: params.cronSession.sessionEntry, + agentId: params.agentId, + trigger: "cron", + jobId: params.job.id, + cleanupCliLiveSessionOnRunEnd: params.usesDetachedRunSession === true, + sessionFile, + workspaceDir: params.workspaceDir, + config: params.cfgWithAgentDefaults, + prompt: modelPrompt, + transcriptPrompt: deliveryTargetRuntimeContext ? promptText : undefined, + modelProvider: providerOverride, + provider: executionProvider, + model: modelOverride, + thinkLevel: candidateThinkLevel, + timeoutMs: params.timeoutMs, + runId, + lane: resolveCronAgentLane(params.lane), + allowEmptyAssistantReplyAsSilent, + cliSessionId: cliSessionBinding?.sessionId, + cliSessionBinding: guardedCliSessionBinding, + skillsSnapshot: params.skillsSnapshot, + messageChannel, + sourceReplyDeliveryMode, + requireExplicitMessageTarget: sourceDelivery.messageTool.requireExplicitTarget, + cliSessionBindingFacts: { + sourceReplyDeliveryMode, + requireExplicitMessageTarget: sourceDelivery.messageTool.requireExplicitTarget, + }, + toolsAllow: resolveCliRuntimeToolsAllow( + params.agentPayload?.toolsAllow, + params.agentPayload?.toolsAllowIsDefault, + ), + abortSignal: params.abortSignal, + onExecutionStarted: params.onExecutionStarted, + onExecutionPhase: params.onExecutionPhase, + bootstrapContextMode, + bootstrapContextRunKind: "cron", + bootstrapPromptWarningSignaturesSeen, + bootstrapPromptWarningSignature, + fastModeStartedAtMs, + fastModeAutoProgressState, + isFinalFallbackAttempt: runOptions?.isFinalFallbackAttempt, + userTurnTranscriptRecorder, + suppressNextUserMessagePersistence: + userTurnTranscriptRecorder.hasPersisted() || + userTurnTranscriptRecorder.isBlocked(), + }), + ); bootstrapPromptWarningSignaturesSeen = resolveBootstrapWarningSignaturesSeen( result.meta?.systemPromptReport, ); diff --git a/src/gateway/method-scopes.test.ts b/src/gateway/method-scopes.test.ts index faedd78c2c30..aacfaa03bf66 100644 --- a/src/gateway/method-scopes.test.ts +++ b/src/gateway/method-scopes.test.ts @@ -51,6 +51,7 @@ describe("method scope resolution", () => { ["taskSuggestions.dismiss", ["operator.write"]], ["config.schema.lookup", ["operator.read"]], ["sessions.create", ["operator.write"]], + ["sessions.dispatch", ["operator.admin"]], ["sessions.send", ["operator.write"]], ["sessions.abort", ["operator.write"]], ["tasks.cancel", ["operator.write"]], diff --git a/src/gateway/methods/core-descriptors.ts b/src/gateway/methods/core-descriptors.ts index 463f506132f3..0c784351d849 100644 --- a/src/gateway/methods/core-descriptors.ts +++ b/src/gateway/methods/core-descriptors.ts @@ -345,6 +345,12 @@ const CORE_GATEWAY_METHOD_SPECS: readonly CoreGatewayMethodSpec[] = [ { name: "approval.get", scope: "operator.approvals" }, { name: "approval.resolve", scope: "operator.approvals" }, { name: "sessions.search", scope: "operator.read" }, + { + name: "sessions.dispatch", + scope: "operator.admin", + startup: true, + controlPlaneWrite: true, + }, ] as const; const CORE_GATEWAY_METHOD_SPEC_BY_NAME: ReadonlyMap = new Map( diff --git a/src/gateway/server-methods-list.test.ts b/src/gateway/server-methods-list.test.ts index 4e309216bb78..ba79540b98e9 100644 --- a/src/gateway/server-methods-list.test.ts +++ b/src/gateway/server-methods-list.test.ts @@ -86,11 +86,11 @@ describe("listGatewayMethods", () => { ]); expect(methods).toContain("tts.speak"); expect(coreMethods.slice(-5)).toEqual([ - "sessions.catalog.continue", "sessions.catalog.archive", "approval.get", "approval.resolve", "sessions.search", + "sessions.dispatch", ]); expect(methods.indexOf("approval.get")).toBeGreaterThan(methods.indexOf("tts.speak")); expect(methods.indexOf("approval.resolve")).toBe(methods.indexOf("approval.get") + 1); diff --git a/src/gateway/server-methods/sessions.dispatch.test.ts b/src/gateway/server-methods/sessions.dispatch.test.ts new file mode 100644 index 000000000000..2d0263a83552 --- /dev/null +++ b/src/gateway/server-methods/sessions.dispatch.test.ts @@ -0,0 +1,403 @@ +import { expectDefined } from "@openclaw/normalization-core"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { ErrorCodes } from "../../../packages/gateway-protocol/src/index.js"; +import type { WorkerSessionPlacementRecord } from "../worker-environments/placement-store.js"; +import type { GatewayRequestContext, RespondFn } from "./types.js"; + +const mocks = vi.hoisted(() => ({ + findLiveByOwner: vi.fn(), + resolveTarget: vi.fn(), +})); + +vi.mock("../../agents/worktrees/service.js", () => ({ + managedWorktrees: { + findLiveByOwner: mocks.findLiveByOwner, + }, +})); + +vi.mock("../session-utils.js", async () => { + const actual = await vi.importActual("../session-utils.js"); + return { + ...actual, + resolveGatewaySessionStoreTargetWithStore: mocks.resolveTarget, + }; +}); + +import { sessionsHandlers } from "./sessions.js"; + +const sessionKey = "agent:main:cloud-test"; +const sessionId = "session-cloud-test"; + +function reclaimedPlacementRecord(): WorkerSessionPlacementRecord { + return { + sessionId, + agentId: "main", + sessionKey, + state: "reclaimed", + environmentId: "environment-previous", + generation: 4, + activeOwnerEpoch: 1, + workspaceBaseManifestRef: "manifest-previous", + remoteWorkspaceDir: "/worker/session-cloud-test", + workerBundleHash: "c".repeat(64), + lastTranscriptAckCursor: 3, + lastLiveEventAckCursor: 2, + recoveryError: null, + turnClaim: null, + createdAtMs: 1, + updatedAtMs: 2, + stateChangedAtMs: 2, + }; +} + +function targetWithEntry(entry?: { + sessionId: string; + worktree?: { id: string; branch: string; repoRoot: string }; + agentHarnessId?: string; + agentRuntimeOverride?: string; + archivedAt?: number; + modelSelectionLocked?: boolean; + providerOverride?: string; + modelOverride?: string; +}) { + // Pin an anthropic model by default: the effective-runtime fallback consults + // the process-global harness registry, so the default openai model resolves + // to "codex" whenever a sibling test in the shard registered that harness. + const pinnedEntry = entry + ? { providerOverride: "anthropic", modelOverride: "claude-test", ...entry } + : undefined; + return { + agentId: "main", + storePath: "/tmp/openclaw-agent.sqlite", + canonicalKey: sessionKey, + storeKeys: [sessionKey], + store: pinnedEntry ? { [sessionKey]: pinnedEntry } : {}, + }; +} + +function makeContext(overrides: Partial = {}): GatewayRequestContext { + return { + getRuntimeConfig: () => ({ + cloudWorkers: { + profiles: { + test: { provider: "fake", region: "test", size: "small" }, + }, + }, + }), + ...overrides, + } as unknown as GatewayRequestContext; +} + +async function invoke(context: GatewayRequestContext) { + const respond = vi.fn() as unknown as RespondFn; + await expectDefined( + sessionsHandlers["sessions.dispatch"], + 'sessionsHandlers["sessions.dispatch"] test invariant', + )({ + req: { id: "dispatch-request" } as never, + params: { key: sessionKey, profileId: "test" }, + respond, + context, + client: null, + isWebchatConnect: () => false, + }); + return respond; +} + +describe("sessions.dispatch", () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.resolveTarget.mockReturnValue(targetWithEntry()); + }); + + it("stays unavailable without a configured placement dispatcher", async () => { + const respond = await invoke(makeContext()); + + expect(respond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ code: ErrorCodes.INVALID_REQUEST }), + ); + }); + + it("rejects a missing session before dispatch", async () => { + const dispatch = vi.fn(); + const respond = await invoke( + makeContext({ + workerPlacementDispatchService: { dispatch }, + workerSessionPlacementService: { getMany: () => new Map() }, + }), + ); + + expect(dispatch).not.toHaveBeenCalled(); + expect(respond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ code: ErrorCodes.INVALID_REQUEST }), + ); + }); + + it("rejects sessions without their bound managed worktree", async () => { + mocks.resolveTarget.mockReturnValue(targetWithEntry({ sessionId })); + const dispatch = vi.fn(); + const respond = await invoke( + makeContext({ + workerPlacementDispatchService: { dispatch }, + workerSessionPlacementService: { getMany: () => new Map() }, + }), + ); + + expect(dispatch).not.toHaveBeenCalled(); + expect(respond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ + code: ErrorCodes.INVALID_REQUEST, + message: expect.stringContaining("session-owned managed worktree"), + }), + ); + }); + + it("rejects dispatch from a nonlocal placement", async () => { + mocks.resolveTarget.mockReturnValue(targetWithEntry({ sessionId })); + const dispatch = vi.fn(); + const respond = await invoke( + makeContext({ + workerPlacementDispatchService: { dispatch }, + workerSessionPlacementService: { + getMany: () => new Map([[sessionId, { state: "requested" } as never]]), + }, + }), + ); + + expect(dispatch).not.toHaveBeenCalled(); + expect(respond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ + code: ErrorCodes.INVALID_REQUEST, + message: expect.stringContaining("placement requested"), + }), + ); + }); + + it("rejects sessions owned by an unsupported runtime", async () => { + mocks.resolveTarget.mockReturnValue( + targetWithEntry({ + sessionId, + agentRuntimeOverride: "codex", + worktree: { id: "worktree-1", branch: "openclaw/cloud-test", repoRoot: "/repo" }, + }), + ); + const dispatch = vi.fn(); + const respond = await invoke( + makeContext({ + workerPlacementDispatchService: { dispatch }, + workerSessionPlacementService: { getMany: () => new Map() }, + }), + ); + + expect(dispatch).not.toHaveBeenCalled(); + expect(respond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ + code: ErrorCodes.INVALID_REQUEST, + message: expect.stringContaining("OpenClaw runtime"), + }), + ); + }); + + it("rejects an archived session before dispatch", async () => { + mocks.resolveTarget.mockReturnValue( + targetWithEntry({ + sessionId, + archivedAt: 2, + worktree: { id: "worktree-1", branch: "openclaw/cloud-test", repoRoot: "/repo" }, + }), + ); + const dispatch = vi.fn(); + const respond = await invoke( + makeContext({ + workerPlacementDispatchService: { dispatch }, + workerSessionPlacementService: { getMany: () => new Map() }, + }), + ); + + expect(dispatch).not.toHaveBeenCalled(); + expect(respond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ + code: ErrorCodes.INVALID_REQUEST, + message: expect.stringContaining("archived"), + }), + ); + }); + + it("allows an explicitly reclaimed session to dispatch again", async () => { + mocks.resolveTarget.mockReturnValue( + targetWithEntry({ + sessionId, + worktree: { id: "worktree-1", branch: "openclaw/cloud-test", repoRoot: "/repo" }, + }), + ); + mocks.findLiveByOwner.mockReturnValue({ + id: "worktree-1", + ownerKind: "session", + ownerId: sessionKey, + }); + const dispatchedPlacement: WorkerSessionPlacementRecord = { + sessionId, + agentId: "main", + sessionKey, + state: "active", + environmentId: "environment-2", + generation: 5, + activeOwnerEpoch: 2, + workspaceBaseManifestRef: "manifest-2", + remoteWorkspaceDir: "/worker/session-cloud-test", + workerBundleHash: "d".repeat(64), + lastTranscriptAckCursor: null, + lastLiveEventAckCursor: null, + recoveryError: null, + turnClaim: null, + createdAtMs: 1, + updatedAtMs: 3, + stateChangedAtMs: 3, + }; + const dispatch = vi.fn().mockResolvedValue(dispatchedPlacement); + const respond = await invoke( + makeContext({ + workerPlacementDispatchService: { dispatch }, + workerSessionPlacementService: { + getMany: () => new Map([[sessionId, reclaimedPlacementRecord()]]), + }, + }), + ); + + expect(dispatch).toHaveBeenCalledWith({ + sessionId, + sessionKey, + agentId: "main", + profileId: "test", + }); + expect(respond).toHaveBeenCalledWith( + true, + expect.objectContaining({ + placement: expect.objectContaining({ + state: "active", + environmentId: "environment-2", + generation: 5, + }), + }), + undefined, + ); + }); + + it.each([ + ["CLI", "claude-cli"], + ["plugin", "test-harness"], + ])("rejects sessions assigned to a configured %s runtime", async (_kind, runtimeId) => { + const modelRef = "anthropic/claude-test"; + mocks.resolveTarget.mockReturnValue( + targetWithEntry({ + sessionId, + providerOverride: "anthropic", + modelOverride: "claude-test", + worktree: { id: "worktree-1", branch: "openclaw/cloud-test", repoRoot: "/repo" }, + }), + ); + const dispatch = vi.fn(); + const respond = await invoke( + makeContext({ + getRuntimeConfig: () => ({ + cloudWorkers: { + profiles: { + test: { provider: "fake", region: "test", size: "small" }, + }, + }, + agents: { + defaults: { + models: { + [modelRef]: { agentRuntime: { id: runtimeId } }, + }, + }, + }, + }), + workerPlacementDispatchService: { dispatch }, + workerSessionPlacementService: { getMany: () => new Map() }, + }), + ); + + expect(dispatch).not.toHaveBeenCalled(); + expect(respond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ + code: ErrorCodes.INVALID_REQUEST, + message: expect.stringContaining(runtimeId), + }), + ); + }); + + it("dispatches an existing managed-worktree session and projects placement", async () => { + mocks.resolveTarget.mockReturnValue( + targetWithEntry({ + sessionId, + worktree: { id: "worktree-1", branch: "openclaw/cloud-test", repoRoot: "/repo" }, + }), + ); + mocks.findLiveByOwner.mockReturnValue({ + id: "worktree-1", + ownerKind: "session", + ownerId: sessionKey, + }); + const dispatch = vi.fn().mockResolvedValue({ + sessionId, + agentId: "main", + sessionKey, + state: "active", + environmentId: "environment-1", + generation: 5, + activeOwnerEpoch: 2, + workspaceBaseManifestRef: "manifest-1", + remoteWorkspaceDir: "/worker/session-cloud-test", + workerBundleHash: "b".repeat(64), + lastTranscriptAckCursor: null, + lastLiveEventAckCursor: null, + recoveryError: null, + turnClaim: null, + createdAtMs: 1, + updatedAtMs: 2, + stateChangedAtMs: 2, + }); + const respond = await invoke( + makeContext({ + workerPlacementDispatchService: { dispatch }, + workerSessionPlacementService: { getMany: () => new Map() }, + }), + ); + + expect(dispatch).toHaveBeenCalledWith({ + sessionId, + sessionKey, + agentId: "main", + profileId: "test", + }); + expect(respond).toHaveBeenCalledWith( + true, + expect.objectContaining({ + ok: true, + key: sessionKey, + sessionId, + placement: expect.objectContaining({ + state: "active", + environmentId: "environment-1", + activeOwnerEpoch: 2, + }), + }), + undefined, + ); + }); +}); diff --git a/src/gateway/server-methods/sessions.ts b/src/gateway/server-methods/sessions.ts index a73d38cdaf79..d20e7bdfbc29 100644 --- a/src/gateway/server-methods/sessions.ts +++ b/src/gateway/server-methods/sessions.ts @@ -12,7 +12,9 @@ import { GATEWAY_CLIENT_IDS } from "../../../packages/gateway-protocol/src/clien import { ErrorCodes, errorShape, + type SessionPlacement, type SessionOperationEvent, + type SessionsPatchParams, validateSessionsAbortParams, validateSessionsCleanupParams, validateSessionsCompactParams, @@ -23,6 +25,7 @@ import { validateSessionsCreateParams, validateSessionsDeleteParams, validateSessionsDescribeParams, + validateSessionsDispatchParams, validateSessionsGroupsDeleteParams, validateSessionsGroupsListParams, validateSessionsGroupsPutParams, @@ -149,6 +152,11 @@ import { import { projectSessionsPatchEntry } from "../sessions-patch.js"; import { resolveSessionKeyFromResolveParams } from "../sessions-resolve.js"; import { asWorkerInferenceControl } from "../worker-environments/inference-control.js"; +import { projectWorkerSessionPlacement } from "../worker-environments/placement-projector.js"; +import { + isWorkerPlacementSessionRuntimeSupported, + resolveWorkerPlacementSessionRuntime, +} from "../worker-environments/placement-session-runtime.js"; import { resolveWorkerSessionTarget } from "../worker-environments/session-target.js"; import { setGatewayDedupeEntry } from "./agent-job.js"; import { chatHandlers } from "./chat.js"; @@ -174,6 +182,81 @@ const compactionCheckpointStore = createFileBackedCompactionCheckpointStore(); const MODEL_SELECTION_LOCKED_CHECKPOINT_MESSAGE = "Checkpoint branch and restore are unavailable while model selection is locked."; +class SessionWorkerPlacementMutationError extends Error { + constructor( + readonly placementState: SessionPlacement["state"], + action: "delete" | "reset" | "restore", + key: string, + ) { + super(`Session ${key} cannot ${action} while cloud worker placement is ${placementState}.`); + } +} + +function resolveSessionWorkerPlacementMutationError(params: { + action: "delete" | "reset" | "restore"; + context: GatewayRequestContext; + key: string; + sessionId: string | undefined; +}): SessionWorkerPlacementMutationError | undefined { + if (!params.sessionId) { + return undefined; + } + const placement = params.context.workerSessionPlacementService + ?.getMany([params.sessionId]) + .get(params.sessionId); + if ( + !placement || + placement.state === "local" || + (params.action === "delete" && placement.state === "reclaimed") + ) { + return undefined; + } + return new SessionWorkerPlacementMutationError(placement.state, params.action, params.key); +} + +function respondSessionWorkerPlacementMutationError( + error: SessionWorkerPlacementMutationError, + respond: RespondFn, +): void { + respond(false, undefined, errorShape(ErrorCodes.INVALID_REQUEST, error.message)); +} + +function resolveSessionWorkerPlacementPatchError(params: { + agentId: string; + cfg: OpenClawConfig; + context: GatewayRequestContext; + entry: SessionEntry | undefined; + key: string; + patch: SessionsPatchParams; + sessionKey: string; + validateModelRuntime: boolean; +}): string | undefined { + const placement = params.entry?.sessionId + ? params.context.workerSessionPlacementService + ?.getMany([params.entry.sessionId]) + .get(params.entry.sessionId) + : undefined; + if (!placement || placement.state === "local") { + return undefined; + } + if (params.patch.archived !== undefined) { + return `Session ${params.key} cannot change archive state while cloud worker placement is ${placement.state}.`; + } + if (!params.validateModelRuntime || params.patch.model === undefined || !params.entry) { + return undefined; + } + const runtime = resolveWorkerPlacementSessionRuntime({ + cfg: params.cfg, + entry: params.entry, + agentId: params.agentId, + sessionKey: params.sessionKey, + }); + if (isWorkerPlacementSessionRuntimeSupported(runtime)) { + return undefined; + } + return `Session ${params.key} cannot select the ${runtime} runtime while cloud worker placement is ${placement.state}.`; +} + function filterSessionStoreToConfiguredAgents( cfg: OpenClawConfig, store: Record, @@ -363,7 +446,7 @@ function emitSessionOperation( } function rejectWebchatSessionMutation(params: { - action: "patch" | "delete" | "compact" | "restore"; + action: "patch" | "delete" | "compact" | "restore" | "dispatch"; client: GatewayClient | null; isWebchatConnect: (params: GatewayClient["connect"] | null | undefined) => boolean; respond: RespondFn; @@ -385,6 +468,14 @@ function rejectWebchatSessionMutation(params: { return true; } +function isWorkerDispatchInputError(error: unknown): boolean { + if (typeof error !== "object" || error === null || !("code" in error)) { + return false; + } + const code = error.code; + return code === "invalid_profile" || code === "profile_not_found" || code === "invalid_state"; +} + function isAgentMainSessionKey(cfg: OpenClawConfig, sessionKey: string): boolean { const parsed = parseAgentSessionKey(sessionKey); if (!parsed) { @@ -967,10 +1058,16 @@ export const sessionsHandlers: GatewayRequestHandlers = { }, }, ); + const placementsBySessionId = context.workerSessionPlacementService?.getMany( + result.sessions.flatMap((session) => (session.sessionId ? [session.sessionId] : [])), + ); const sessions = measureDiagnosticsTimelineSpanSync( "gateway.sessions.list.active_run_flags", () => { return result.sessions.map((session) => { + const placementRecord = session.sessionId + ? placementsBySessionId?.get(session.sessionId) + : undefined; const activeRunState = resolveVisibleActiveSessionRunState({ context, requestedKey: session.key, @@ -981,6 +1078,9 @@ export const sessionsHandlers: GatewayRequestHandlers = { }); return Object.assign({}, session, { hasActiveRun: activeRunState.active, + ...(placementRecord + ? { placement: projectWorkerSessionPlacement(placementRecord) } + : {}), ...(activeRunState.runIds.length > 0 ? { activeRunIds: activeRunState.runIds } : {}), @@ -1282,7 +1382,16 @@ export const sessionsHandlers: GatewayRequestHandlers = { includeLastMessage: p.includeLastMessage, transcriptUsageMaxBytes: 64 * 1024, }); - respond(true, { session: row }, undefined); + const placement = row.sessionId + ? context.workerSessionPlacementService?.getMany([row.sessionId]).get(row.sessionId) + : undefined; + respond( + true, + { + session: placement ? { ...row, placement: projectWorkerSessionPlacement(placement) } : row, + }, + undefined, + ); }, "sessions.resolve": async ({ params, respond, context }) => { if (!assertValidParams(params, validateSessionsResolveParams, "sessions.resolve", respond)) { @@ -1924,6 +2033,16 @@ export const sessionsHandlers: GatewayRequestHandlers = { ); return; } + const initialPlacementError = resolveSessionWorkerPlacementMutationError({ + action: "restore", + context, + key, + sessionId: entry.sessionId, + }); + if (initialPlacementError) { + respondSessionWorkerPlacementMutationError(initialPlacementError, respond); + return; + } const lifecycleIdentities = [ key, canonicalKey, @@ -1935,6 +2054,7 @@ export const sessionsHandlers: GatewayRequestHandlers = { let admittedWorkReleased = true; let restoreTargetStillCurrent = true; let restoreBlockedByModelLock = false; + let restorePlacementError: SessionWorkerPlacementMutationError | undefined; // Restore replaces the active transcript identity. Hold the same lifecycle fence as // compaction so neither operation can publish state from the other's obsolete session. await runExclusiveSessionLifecycleMutation({ @@ -1960,6 +2080,15 @@ export const sessionsHandlers: GatewayRequestHandlers = { if (restoreBlockedByModelLock) { return; } + restorePlacementError = resolveSessionWorkerPlacementMutationError({ + action: "restore", + context, + key, + sessionId: current.entry?.sessionId, + }); + if (restorePlacementError) { + return; + } clearSessionQueues([ key, current.canonicalKey, @@ -1993,6 +2122,10 @@ export const sessionsHandlers: GatewayRequestHandlers = { ); return; } + if (restorePlacementError) { + respondSessionWorkerPlacementMutationError(restorePlacementError, respond); + return; + } if (!admittedWorkReleased) { respond( false, @@ -2113,6 +2246,145 @@ export const sessionsHandlers: GatewayRequestHandlers = { }, }); }, + "sessions.dispatch": async ({ params, respond, context, client, isWebchatConnect }) => { + if (!assertValidParams(params, validateSessionsDispatchParams, "sessions.dispatch", respond)) { + return; + } + const key = requireSessionKey(params.key, respond); + if (!key) { + return; + } + if (rejectWebchatSessionMutation({ action: "dispatch", client, isWebchatConnect, respond })) { + return; + } + const dispatchService = context.workerPlacementDispatchService; + const placementReader = context.workerSessionPlacementService; + if (!dispatchService || !placementReader) { + respond( + false, + undefined, + errorShape(ErrorCodes.INVALID_REQUEST, "cloud worker dispatch is not configured"), + ); + return; + } + const cfg = context.getRuntimeConfig(); + const requestedAgent = resolveRequestedGlobalAgentId(cfg, key, params.agentId); + if (!requestedAgent.ok) { + respond(false, undefined, requestedAgent.error); + return; + } + if (!Object.hasOwn(cfg.cloudWorkers?.profiles ?? {}, params.profileId)) { + respond( + false, + undefined, + errorShape( + ErrorCodes.INVALID_REQUEST, + `cloud worker profile is not configured: ${params.profileId}`, + ), + ); + return; + } + const target = loadAccessorSessionEntryForGatewayTarget({ + key, + cfg, + agentId: requestedAgent.agentId, + }); + const entry = target.entry; + const sessionId = normalizeOptionalString(entry?.sessionId); + if (!entry || !sessionId) { + respond( + false, + undefined, + errorShape(ErrorCodes.INVALID_REQUEST, `session not found: ${key}`), + ); + return; + } + if (entry.archivedAt !== undefined) { + respond( + false, + undefined, + errorShape(ErrorCodes.INVALID_REQUEST, "cannot dispatch an archived session"), + ); + return; + } + const sessionRuntime = resolveWorkerPlacementSessionRuntime({ + cfg, + entry, + agentId: target.target.agentId, + sessionKey: target.canonicalKey, + }); + if (!isWorkerPlacementSessionRuntimeSupported(sessionRuntime)) { + respond( + false, + undefined, + errorShape( + ErrorCodes.INVALID_REQUEST, + `cloud worker dispatch requires the OpenClaw runtime, not ${sessionRuntime}`, + ), + ); + return; + } + const existingPlacement = placementReader.getMany([sessionId]).get(sessionId); + if ( + existingPlacement && + existingPlacement.state !== "local" && + existingPlacement.state !== "reclaimed" + ) { + respond( + false, + undefined, + errorShape( + ErrorCodes.INVALID_REQUEST, + `session cannot dispatch from placement ${existingPlacement.state}`, + ), + ); + return; + } + const worktree = managedWorktrees.findLiveByOwner("session", target.canonicalKey); + if ( + !target.entry?.worktree?.id || + !worktree || + worktree.id !== target.entry.worktree.id || + worktree.ownerId !== target.canonicalKey + ) { + respond( + false, + undefined, + errorShape( + ErrorCodes.INVALID_REQUEST, + "sessions.dispatch requires a session-owned managed worktree", + ), + ); + return; + } + try { + const placement = await dispatchService.dispatch({ + sessionId, + sessionKey: target.canonicalKey, + agentId: target.target.agentId, + profileId: params.profileId, + }); + respond( + true, + { + ok: true, + key: target.canonicalKey, + sessionId, + placement: projectWorkerSessionPlacement(placement), + }, + undefined, + ); + } catch (error) { + respond( + false, + undefined, + errorShape( + isWorkerDispatchInputError(error) ? ErrorCodes.INVALID_REQUEST : ErrorCodes.UNAVAILABLE, + formatErrorMessage(error), + ), + ); + } + }, "sessions.send": async ({ req, params, respond, context, client, isWebchatConnect }) => { await handleSessionSend({ method: "sessions.send", @@ -2354,6 +2626,20 @@ export const sessionsHandlers: GatewayRequestHandlers = { respond(false, undefined, errorShape(ErrorCodes.INVALID_REQUEST, missingHarnessSessionError)); return; } + const initialPlacementPatchError = resolveSessionWorkerPlacementPatchError({ + agentId: target.agentId, + cfg, + context, + entry: lifecycleEntry, + key, + patch: p, + sessionKey: canonicalKey, + validateModelRuntime: false, + }); + if (initialPlacementPatchError) { + respond(false, undefined, errorShape(ErrorCodes.INVALID_REQUEST, initialPlacementPatchError)); + return; + } const lifecycleIdentities = [canonicalKey, key, lifecycleEntry?.sessionId]; if (p.archived === true && isSessionLifecycleMutationActive(storePath, lifecycleIdentities)) { respond( @@ -2436,8 +2722,8 @@ export const sessionsHandlers: GatewayRequestHandlers = { }); return { primaryKey, candidateKeys: migratedTarget.storeKeys }; }, - project: async ({ primaryKey, existingEntry, entries }) => - await projectSessionsPatchEntry({ + project: async ({ primaryKey, existingEntry, entries }) => { + const projected = await projectSessionsPatchEntry({ cfg, entries, existingEntry, @@ -2445,7 +2731,27 @@ export const sessionsHandlers: GatewayRequestHandlers = { agentId: requestedAgentId, patch: p, loadGatewayModelCatalog: loadPatchModelCatalog, - }), + }); + if (!projected.ok) { + return projected; + } + const placementPatchError = resolveSessionWorkerPlacementPatchError({ + agentId: target.agentId, + cfg, + context, + entry: projected.entry, + key, + patch: p, + sessionKey: canonicalKey, + validateModelRuntime: true, + }); + return placementPatchError + ? { + ok: false, + error: errorShape(ErrorCodes.INVALID_REQUEST, placementPatchError), + } + : projected; + }, }); }; const applied = await runExclusiveSessionLifecycleMutation({ @@ -2774,6 +3080,16 @@ export const sessionsHandlers: GatewayRequestHandlers = { if (rejectExpectedSessionMismatch(initialDeleteEntry)) { return; } + const initialPlacementError = resolveSessionWorkerPlacementMutationError({ + action: "delete", + context, + key, + sessionId: normalizeOptionalString(initialDeleteEntry?.sessionId), + }); + if (initialPlacementError) { + respondSessionWorkerPlacementMutationError(initialPlacementError, respond); + return; + } if ( rejectPluginRuntimeDeleteMismatch({ client, @@ -2821,6 +3137,7 @@ export const sessionsHandlers: GatewayRequestHandlers = { let admittedWorkReleased = true; let expectedSessionStillCurrent = true; let deleteBlockedByModelLock = false; + let deleteBlockedByWorkerPlacement = false; const deletion = await runExclusiveSessionLifecycleMutation({ scope: storePath, identities: deleteLifecycleIdentities, @@ -2834,6 +3151,17 @@ export const sessionsHandlers: GatewayRequestHandlers = { if (!expectedSessionStillCurrent) { return; } + const placementError = resolveSessionWorkerPlacementMutationError({ + action: "delete", + context, + key, + sessionId: normalizeOptionalString(preparedEntry?.sessionId), + }); + if (placementError) { + deleteBlockedByWorkerPlacement = true; + respondSessionWorkerPlacementMutationError(placementError, respond); + return; + } admittedWorkReleased = await interruptSessionWorkAdmissions({ scope: storePath, identities: deleteLifecycleIdentities, @@ -2841,7 +3169,11 @@ export const sessionsHandlers: GatewayRequestHandlers = { }); }, run: async () => { - if (deleteBlockedByModelLock || !expectedSessionStillCurrent) { + if ( + deleteBlockedByModelLock || + deleteBlockedByWorkerPlacement || + !expectedSessionStillCurrent + ) { return undefined; } if (!admittedWorkReleased) { diff --git a/src/gateway/server-methods/shared-types.ts b/src/gateway/server-methods/shared-types.ts index f4eb970b583d..80136f87a8e3 100644 --- a/src/gateway/server-methods/shared-types.ts +++ b/src/gateway/server-methods/shared-types.ts @@ -40,7 +40,11 @@ import type { DedupeEntry } from "../server-shared.js"; import type { GatewayEventLoopHealth } from "../server/event-loop-health.js"; import type { TerminalLaunchResolution } from "../terminal/launch.js"; import type { TerminalSessionManager } from "../terminal/session-manager.js"; -import type { WorkerEnvironmentServiceContract } from "../worker-environments/service-contract.js"; +import type { WorkerSessionPlacementReader } from "../worker-environments/placement-projector.js"; +import type { + WorkerEnvironmentServiceContract, + WorkerPlacementDispatchContract, +} from "../worker-environments/service-contract.js"; /** * Shared gateway request types used by every server-method module. @@ -142,6 +146,10 @@ export type GatewayRequestContext = { nodeRegistry: NodeRegistry; /** Durable cloud-worker lifecycle; absent from lightweight in-process contexts. */ workerEnvironmentService?: WorkerEnvironmentServiceContract; + /** Durable per-session worker placement; absent when cloud workers are disabled. */ + workerSessionPlacementService?: WorkerSessionPlacementReader; + /** One-way local-to-worker dispatch; absent when cloud workers are disabled. */ + workerPlacementDispatchService?: WorkerPlacementDispatchContract; // Operator terminal session store. Absent in local/in-process contexts where // no PTY surface is served. terminalSessions?: TerminalSessionManager; diff --git a/src/gateway/server-request-context.ts b/src/gateway/server-request-context.ts index 57f88235852c..ddfa5858fec2 100644 --- a/src/gateway/server-request-context.ts +++ b/src/gateway/server-request-context.ts @@ -48,6 +48,8 @@ type GatewayRequestContextParams = { enforceSharedGatewayAuthGenerationForConfigWrite: (nextConfig: OpenClawConfig) => void; nodeRegistry: GatewayRequestContext["nodeRegistry"]; workerEnvironmentService?: GatewayRequestContext["workerEnvironmentService"]; + workerSessionPlacementService?: GatewayRequestContext["workerSessionPlacementService"]; + workerPlacementDispatchService?: GatewayRequestContext["workerPlacementDispatchService"]; terminalSessions?: GatewayRequestContext["terminalSessions"]; agentRunSeq: GatewayRequestContext["agentRunSeq"]; chatAbortControllers: GatewayRequestContext["chatAbortControllers"]; @@ -211,6 +213,12 @@ export function createGatewayRequestContext( ...(params.workerEnvironmentService ? { workerEnvironmentService: params.workerEnvironmentService } : {}), + ...(params.workerSessionPlacementService + ? { workerSessionPlacementService: params.workerSessionPlacementService } + : {}), + ...(params.workerPlacementDispatchService + ? { workerPlacementDispatchService: params.workerPlacementDispatchService } + : {}), terminalSessions: params.terminalSessions, agentRunSeq: params.agentRunSeq, chatAbortControllers: params.chatAbortControllers, diff --git a/src/gateway/server-startup-post-attach.test.ts b/src/gateway/server-startup-post-attach.test.ts index 951cfdaed1b9..97aca6b8ab34 100644 --- a/src/gateway/server-startup-post-attach.test.ts +++ b/src/gateway/server-startup-post-attach.test.ts @@ -2173,16 +2173,26 @@ describe("startGatewayPostAttachRuntime", () => { expect(startGatewaySidecarsValue).toHaveBeenCalledTimes(1); }); - it("starts the worker environment sidecar before releasing startup-gated methods", async () => { + it("reconciles worker placement before starting channels and sidecars", async () => { let finishReconcile: (() => void) | undefined; const reconcileReady = new Promise((resolve) => { finishReconcile = resolve; }); + const startupOrder: string[] = []; const workerSidecar = { stop: vi.fn() }; const startWorkerEnvironmentRuntime = vi.fn(async () => { + startupOrder.push("worker-reconcile"); await reconcileReady; + startupOrder.push("worker-ready"); return workerSidecar; }); + const startGatewaySidecarsValue = vi.fn(async () => { + startupOrder.push("gateway-sidecars"); + return { + pluginServices: null, + postReadySidecars: [], + }; + }); const onGatewayLifetimeSidecars = vi.fn(); const unavailableGatewayMethods = new Set(STARTUP_UNAVAILABLE_GATEWAY_METHODS); @@ -2195,25 +2205,47 @@ describe("startGatewayPostAttachRuntime", () => { onGatewayLifetimeSidecars, }, createPostAttachRuntimeDeps({ - startGatewaySidecars: vi.fn(async () => ({ - pluginServices: null, - postReadySidecars: [], - })), + startGatewaySidecars: startGatewaySidecarsValue, }), ); await vi.waitFor(() => { expect(startWorkerEnvironmentRuntime).toHaveBeenCalledTimes(1); }); + expect(startGatewaySidecarsValue).not.toHaveBeenCalled(); + expect(startupOrder).toEqual(["worker-reconcile"]); expect([...unavailableGatewayMethods]).toEqual([...STARTUP_UNAVAILABLE_GATEWAY_METHODS]); finishReconcile?.(); await vi.waitFor(() => { - expect([...unavailableGatewayMethods]).toEqual([]); + expect(startGatewaySidecarsValue).toHaveBeenCalledTimes(1); }); + expect(startupOrder).toEqual(["worker-reconcile", "worker-ready", "gateway-sidecars"]); + expect([...unavailableGatewayMethods]).toEqual([]); expect(onGatewayLifetimeSidecars).toHaveBeenCalledWith(expect.arrayContaining([workerSidecar])); }); + it("stops worker placement runtime when channel and sidecar startup fails", async () => { + const workerSidecar = { stop: vi.fn(async () => {}) }; + const startupError = new Error("sidecar startup failed"); + + await expect( + startGatewayPostAttachRuntime( + { + ...createPostAttachParams(), + startWorkerEnvironmentRuntime: vi.fn(() => workerSidecar), + }, + createPostAttachRuntimeDeps({ + startGatewaySidecars: vi.fn(async () => { + throw startupError; + }), + }), + ), + ).rejects.toBe(startupError); + + expect(workerSidecar.stop).toHaveBeenCalledTimes(1); + }); + it("does not start the worker environment sidecar after close begins", async () => { const startWorkerEnvironmentRuntime = vi.fn(() => ({ stop: vi.fn() })); const startGatewaySidecarsValue = vi.fn(async () => ({ diff --git a/src/gateway/server-startup-post-attach.ts b/src/gateway/server-startup-post-attach.ts index 5fd26583a2ff..0d19617de437 100644 --- a/src/gateway/server-startup-post-attach.ts +++ b/src/gateway/server-startup-post-attach.ts @@ -1273,28 +1273,35 @@ export async function startGatewayPostAttachRuntime( ? Promise.resolve({ pluginServices: null, pluginRegistry, postReadySidecars: [] }) : waitForSidecarStartTurn().then(async () => { await loadStartupPluginsIfNeeded(); - params.log.info("starting channels and sidecars..."); - const loaderStatsBefore = getPluginModuleLoaderStats(); - const result = await measureStartup(params.startupTrace, "sidecars.total", () => - runtimeDeps.startGatewaySidecars({ - cfg: params.gatewayPluginConfigAtStart, - pluginRegistry, - defaultWorkspaceDir: params.defaultWorkspaceDir, - deps: params.deps, - startChannels: params.startChannels, - log: params.log, - logHooks: params.logHooks, - logChannels: params.logChannels, - startupTrace: params.startupTrace, - onChannelsStarted: params.onChannelsStarted, - onPluginServices: reportPluginServices, - shouldStartPluginServices: () => params.isClosing?.() !== true, - startupOutcomes, - }), - ); const workerEnvironmentSidecar = params.isClosing?.() ? null : ((await params.startWorkerEnvironmentRuntime?.()) ?? null); + params.log.info("starting channels and sidecars..."); + const loaderStatsBefore = getPluginModuleLoaderStats(); + const result = await (async () => { + try { + return await measureStartup(params.startupTrace, "sidecars.total", () => + runtimeDeps.startGatewaySidecars({ + cfg: params.gatewayPluginConfigAtStart, + pluginRegistry, + defaultWorkspaceDir: params.defaultWorkspaceDir, + deps: params.deps, + startChannels: params.startChannels, + log: params.log, + logHooks: params.logHooks, + logChannels: params.logChannels, + startupTrace: params.startupTrace, + onChannelsStarted: params.onChannelsStarted, + onPluginServices: reportPluginServices, + shouldStartPluginServices: () => params.isClosing?.() !== true, + startupOutcomes, + }), + ); + } catch (error) { + await workerEnvironmentSidecar?.stop(); + throw error; + } + })(); const loaderStatsAfter = getPluginModuleLoaderStats(); params.startupTrace?.detail("sidecars.plugin-loader", [ ["callsCount", loaderStatsAfter.calls - loaderStatsBefore.calls], diff --git a/src/gateway/server-worker-environment-startup.ts b/src/gateway/server-worker-environment-startup.ts index bba7cee28073..7a515a5e089a 100644 --- a/src/gateway/server-worker-environment-startup.ts +++ b/src/gateway/server-worker-environment-startup.ts @@ -8,6 +8,7 @@ import { import { createLazyRuntimeModule } from "../shared/lazy-runtime.js"; import type { WorkerBundleProducer, WorkerNpmArtifact } from "./worker-environments/bundle.js"; import type { WorkerLiveEventReceiver } from "./worker-environments/live-events.js"; +import type { WorkerSessionPlacementStore } from "./worker-environments/placement-store.js"; import type { WorkerEnvironmentService } from "./worker-environments/service.js"; type WorkerEnvironmentStore = ReturnType< @@ -24,6 +25,8 @@ export type GatewayWorkerEnvironmentStartupState = { listDurableProviderIds: () => string[]; records: WorkerEnvironmentRecord[]; store: WorkerEnvironmentStore; + placementStore: WorkerSessionPlacementStore; + hasNonlocalPlacementRecords: boolean; }; export type GatewayWorkerEnvironmentRuntime = { @@ -39,8 +42,13 @@ const loadWorkerInferenceRuntimeModule = createLazyRuntimeModule( ); export async function loadGatewayWorkerEnvironmentStartupState(): Promise { - const { createWorkerEnvironmentStore } = await import("./worker-environments/store.js"); + const [{ createWorkerEnvironmentStore }, { createWorkerSessionPlacementStore }] = + await Promise.all([ + import("./worker-environments/store.js"), + import("./worker-environments/placement-store.js"), + ]); const store = createWorkerEnvironmentStore(); + const placementStore = createWorkerSessionPlacementStore(); const records = store.list(); const durableProviderIds = uniqueStrings( records.flatMap((record) => @@ -51,7 +59,15 @@ export async function loadGatewayWorkerEnvironmentStartupState(): Promise uniqueStrings(store.listForReconcile().map((record) => record.providerId)); - return { durableProviderIds, listDurableProviderIds, records, store }; + return { + durableProviderIds, + listDurableProviderIds, + records, + store, + placementStore, + // Non-local placements must revive the worker service even without configured profiles. + hasNonlocalPlacementRecords: placementStore.listForReconcile().length > 0, + }; } export async function createGatewayWorkerEnvironmentRuntime(params: { @@ -63,16 +79,21 @@ export async function createGatewayWorkerEnvironmentRuntime(params: { const [ { createWorkerEnvironmentService }, { createWorkerLiveEventReceiver }, + { createWorkerSessionPlacementGate }, { createWorkerTranscriptCommitter }, { createWorkerTunnelManager }, { resolveWorkerProvider }, ] = await Promise.all([ import("./worker-environments/service.js"), import("./worker-environments/live-events.js"), + import("./worker-environments/placement-worker-gate.js"), import("./worker-environments/transcript-commit.js"), import("./worker-environments/tunnel.js"), import("../plugins/worker-provider-registry.js"), ]); + // A crashed gateway can leak local turn claims; drop them before workers re-admit turns. + params.startup.placementStore.clearLocalTurnClaimsAfterRestart(); + const placementGate = createWorkerSessionPlacementGate(params.startup.placementStore); let workerBundleProducer: WorkerBundleProducer | undefined; let workerNpmArtifact: Promise | undefined; const prepareInstallation = async (install: "bundle" | "npm") => { @@ -128,6 +149,7 @@ export async function createGatewayWorkerEnvironmentRuntime(params: { const workerInferenceRuntime = await loadWorkerInferenceRuntimeModule(); return await workerInferenceRuntime.executeWorkerInference(inferenceParams); }, + placementStore: placementGate, liveEvents: workerLiveEvents, resolveSshIdentity: async ({ provider, leaseId, profile, keyRef }) => { const workerRuntime = await loadWorkerEnvironmentRuntimeModule(); diff --git a/src/gateway/server-worker-placement-startup.ts b/src/gateway/server-worker-placement-startup.ts new file mode 100644 index 000000000000..a3fa97ed9b8c --- /dev/null +++ b/src/gateway/server-worker-placement-startup.ts @@ -0,0 +1,429 @@ +import { + installSessionPlacementAdmissionProvider, + installSessionPlacementResetGuard, +} from "../agents/session-placement-admission.js"; +import { clearSessionQueues } from "../auto-reply/reply/queue/cleanup.js"; +import { getRuntimeConfig } from "../config/config.js"; +import { runExclusiveSessionStoreWrite } from "../config/sessions/store-writer.js"; +import { formatErrorMessage } from "../infra/errors.js"; +import { + interruptSessionWorkAdmissions, + runExclusiveSessionLifecycleMutation, + SESSION_WORK_ADMISSION_DRAIN_TIMEOUT_MS, +} from "../sessions/session-lifecycle-admission.js"; +import { createLazyRuntimeModule } from "../shared/lazy-runtime.js"; +import { + createWorkerPlacementDispatchService, + type WorkerPlacementDispatchService, +} from "./worker-environments/placement-dispatch.js"; +import type { WorkerSessionPlacementStore } from "./worker-environments/placement-store.js"; +import type { WorkerEnvironmentService } from "./worker-environments/service.js"; +import { createWorkerSessionTurnPlacementProvider } from "./worker-environments/worker-turn-launcher.js"; + +const WORKER_PLACEMENT_RECONCILE_INTERVAL_MS = 60_000; + +const loadWorkerPlacementSessionRuntimeModule = createLazyRuntimeModule(async () => { + const [placementSessionRuntime, { managedWorktrees }, sessionUtils] = await Promise.all([ + import("./worker-environments/placement-session-runtime.js"), + import("../agents/worktrees/service.js"), + import("./session-utils.js"), + ]); + return { + isWorkerPlacementSessionRuntimeSupported: + placementSessionRuntime.isWorkerPlacementSessionRuntimeSupported, + managedWorktrees, + resolveWorkerPlacementSessionRuntime: + placementSessionRuntime.resolveWorkerPlacementSessionRuntime, + resolveFreshestSessionEntryFromStoreKeys: sessionUtils.resolveFreshestSessionEntryFromStoreKeys, + resolveGatewaySessionStoreTargetWithStore: + sessionUtils.resolveGatewaySessionStoreTargetWithStore, + }; +}); + +class WorkerDispatchTargetChangedError extends Error { + readonly code = "invalid_state"; +} + +/** Serializes reconciliation sweeps against in-flight dispatches so a sweep never + * observes a placement mid-transition. Dispatches wait out any pending sweep. */ +function coordinateWorkerPlacementDispatch( + service: WorkerPlacementDispatchService, +): WorkerPlacementDispatchService { + let activeDispatchCount = 0; + let reconciliation: Promise | undefined; + const dispatchIdleWaiters = new Set<() => void>(); + const waitForDispatchIdle = (): Promise => { + if (activeDispatchCount === 0) { + return Promise.resolve(); + } + return new Promise((resolve) => { + dispatchIdleWaiters.add(resolve); + }); + }; + const runReconciliation = (operation: () => Promise): Promise => { + if (reconciliation) { + return reconciliation; + } + const current = (async () => { + await waitForDispatchIdle(); + await operation(); + })(); + reconciliation = current; + const clearCurrent = () => { + if (reconciliation === current) { + reconciliation = undefined; + } + }; + void current.then(clearCurrent, clearCurrent); + return current; + }; + return { + dispatch: async (request) => { + for (;;) { + const pendingReconciliation = reconciliation; + if (!pendingReconciliation) { + break; + } + await pendingReconciliation.catch(() => undefined); + } + activeDispatchCount += 1; + try { + return await service.dispatch(request); + } finally { + activeDispatchCount -= 1; + if (activeDispatchCount === 0) { + const waiters = [...dispatchIdleWaiters]; + dispatchIdleWaiters.clear(); + for (const resolve of waiters) { + resolve(); + } + } + } + }, + reconcile: () => runReconciliation(service.reconcile), + reconcileActive: () => runReconciliation(service.reconcileActive), + }; +} + +type WorkerPlacementSidecar = { stop: () => Promise }; + +export type GatewayWorkerPlacementRuntimeParams = { + placements: WorkerSessionPlacementStore; + environments: WorkerEnvironmentService; + admitNewPlacements: boolean; + revokeSessionAuthority: (request: { sessionId: string; sessionKeys: readonly string[] }) => void; + warn: (message: string) => void; +}; + +export type GatewayWorkerPlacementRuntime = ReturnType; + +export function createGatewayWorkerPlacementRuntime(params: GatewayWorkerPlacementRuntimeParams) { + const dispatchService = coordinateWorkerPlacementDispatch( + createWorkerPlacementDispatchService({ + placements: params.placements, + environments: params.environments, + runLocalBarrier: async ({ sessionId, sessionKey, agentId, startDispatch }) => { + const { + isWorkerPlacementSessionRuntimeSupported, + managedWorktrees, + resolveFreshestSessionEntryFromStoreKeys, + resolveGatewaySessionStoreTargetWithStore, + resolveWorkerPlacementSessionRuntime, + } = await loadWorkerPlacementSessionRuntimeModule(); + const target = resolveGatewaySessionStoreTargetWithStore({ + cfg: getRuntimeConfig(), + key: sessionKey, + agentId, + clone: false, + }); + const lifecycleIdentities = [ + sessionKey, + target.canonicalKey, + ...target.storeKeys, + sessionId, + ]; + let placement: ReturnType | undefined; + await runExclusiveSessionLifecycleMutation({ + scope: target.storePath, + identities: lifecycleIdentities, + prepare: async () => { + const currentConfig = getRuntimeConfig(); + const currentTarget = resolveGatewaySessionStoreTargetWithStore({ + cfg: currentConfig, + key: sessionKey, + agentId, + clone: false, + }); + const currentEntry = resolveFreshestSessionEntryFromStoreKeys( + currentTarget.store, + currentTarget.storeKeys, + ); + const worktree = managedWorktrees.findLiveByOwner( + "session", + currentTarget.canonicalKey, + ); + if ( + currentTarget.storePath !== target.storePath || + currentTarget.canonicalKey !== target.canonicalKey || + currentTarget.agentId !== target.agentId || + currentEntry?.sessionId !== sessionId || + !currentEntry.worktree?.id || + !worktree || + worktree.id !== currentEntry.worktree.id || + worktree.ownerId !== currentTarget.canonicalKey + ) { + throw new WorkerDispatchTargetChangedError( + `Session ${sessionKey} changed before cloud worker dispatch. Retry.`, + ); + } + if (currentEntry.archivedAt !== undefined) { + throw new WorkerDispatchTargetChangedError( + `Session ${sessionKey} was archived before cloud worker dispatch. Retry.`, + ); + } + const currentRuntime = resolveWorkerPlacementSessionRuntime({ + cfg: currentConfig, + entry: currentEntry, + agentId: currentTarget.agentId, + sessionKey: currentTarget.canonicalKey, + }); + if (!isWorkerPlacementSessionRuntimeSupported(currentRuntime)) { + throw new WorkerDispatchTargetChangedError( + `Session ${sessionKey} runtime changed to ${currentRuntime} before cloud worker dispatch. Retry.`, + ); + } + placement = startDispatch(); + clearSessionQueues(lifecycleIdentities); + params.revokeSessionAuthority({ + sessionId, + sessionKeys: lifecycleIdentities, + }); + const released = await interruptSessionWorkAdmissions({ + scope: target.storePath, + identities: lifecycleIdentities, + timeoutMs: SESSION_WORK_ADMISSION_DRAIN_TIMEOUT_MS, + }); + if (!released) { + throw new Error(`Session ${sessionKey} is still active; dispatch stopped`); + } + await params.placements.waitForTurnClaimRelease(sessionId, { + timeoutMs: SESSION_WORK_ADMISSION_DRAIN_TIMEOUT_MS, + }); + await runExclusiveSessionStoreWrite(target.storePath, async () => {}, { + reentrant: true, + }); + }, + run: async () => { + if (!placement) { + throw new Error(`Session ${sessionKey} dispatch barrier did not start`); + } + }, + }); + if (!placement) { + throw new Error(`Session ${sessionKey} dispatch barrier did not complete`); + } + return placement; + }, + runActivationBarrier: async ({ sessionId, sessionKey, agentId, activate }) => { + const { + isWorkerPlacementSessionRuntimeSupported, + managedWorktrees, + resolveFreshestSessionEntryFromStoreKeys, + resolveGatewaySessionStoreTargetWithStore, + resolveWorkerPlacementSessionRuntime, + } = await loadWorkerPlacementSessionRuntimeModule(); + const target = resolveGatewaySessionStoreTargetWithStore({ + cfg: getRuntimeConfig(), + key: sessionKey, + agentId, + clone: false, + }); + const lifecycleIdentities = [ + sessionKey, + target.canonicalKey, + ...target.storeKeys, + sessionId, + ]; + let activePlacement: ReturnType | undefined; + await runExclusiveSessionLifecycleMutation({ + scope: target.storePath, + identities: lifecycleIdentities, + run: async () => { + const currentConfig = getRuntimeConfig(); + const currentTarget = resolveGatewaySessionStoreTargetWithStore({ + cfg: currentConfig, + key: sessionKey, + agentId, + clone: false, + }); + const currentEntry = resolveFreshestSessionEntryFromStoreKeys( + currentTarget.store, + currentTarget.storeKeys, + ); + const worktree = managedWorktrees.findLiveByOwner( + "session", + currentTarget.canonicalKey, + ); + if ( + currentTarget.storePath !== target.storePath || + currentTarget.canonicalKey !== target.canonicalKey || + currentTarget.agentId !== target.agentId || + currentEntry?.sessionId !== sessionId || + !currentEntry.worktree?.id || + !worktree || + worktree.id !== currentEntry.worktree.id || + worktree.ownerId !== currentTarget.canonicalKey + ) { + throw new WorkerDispatchTargetChangedError( + `Session ${sessionKey} changed before cloud worker activation. Retry.`, + ); + } + if (currentEntry.archivedAt !== undefined) { + throw new WorkerDispatchTargetChangedError( + `Session ${sessionKey} was archived before cloud worker activation. Retry.`, + ); + } + const currentRuntime = resolveWorkerPlacementSessionRuntime({ + cfg: currentConfig, + entry: currentEntry, + agentId: currentTarget.agentId, + sessionKey: currentTarget.canonicalKey, + }); + if (!isWorkerPlacementSessionRuntimeSupported(currentRuntime)) { + throw new WorkerDispatchTargetChangedError( + `Session ${sessionKey} runtime changed to ${currentRuntime} before cloud worker activation. Retry.`, + ); + } + activePlacement = activate(); + }, + }); + if (!activePlacement) { + throw new Error(`Session ${sessionKey} activation barrier did not complete`); + } + return activePlacement; + }, + resolveWorkspacePath: async ({ sessionId, sessionKey, agentId }) => { + const { + managedWorktrees, + resolveFreshestSessionEntryFromStoreKeys, + resolveGatewaySessionStoreTargetWithStore, + } = await loadWorkerPlacementSessionRuntimeModule(); + const target = resolveGatewaySessionStoreTargetWithStore({ + cfg: getRuntimeConfig(), + key: sessionKey, + agentId, + clone: false, + }); + const sessionEntry = resolveFreshestSessionEntryFromStoreKeys( + target.store, + target.storeKeys, + ); + const worktree = managedWorktrees.findLiveByOwner("session", target.canonicalKey); + if ( + sessionEntry?.sessionId !== sessionId || + !sessionEntry.worktree?.id || + !worktree || + worktree.id !== sessionEntry.worktree.id || + worktree.ownerId !== target.canonicalKey + ) { + throw new Error( + `Session ${sessionKey} dispatch requires a session-owned managed worktree`, + ); + } + return worktree.path; + }, + }), + ); + const admissionProvider = createWorkerSessionTurnPlacementProvider({ + environments: params.environments, + placements: params.placements, + admitNewPlacements: params.admitNewPlacements, + }); + const startRuntime = async (hooks: { + isClosePreludeStarted: () => boolean; + registerSidecar: (sidecar: WorkerPlacementSidecar) => void; + }): Promise => { + const uninstallPlacementAdmission = installSessionPlacementAdmissionProvider(admissionProvider); + const uninstallPlacementResetGuard = installSessionPlacementResetGuard((sessionId) => { + const placement = params.placements.get(sessionId); + if (!placement || placement.state === "local") { + return undefined; + } + return `cloud worker placement is ${placement.state}`; + }); + let placementReconcileInterval: ReturnType | undefined; + let placementReconcileInFlight: Promise | undefined; + let stopped = false; + const reconcileActivePlacements = (): Promise => { + if (stopped) { + return Promise.resolve(); + } + if (placementReconcileInFlight) { + return placementReconcileInFlight; + } + const current = dispatchService.reconcileActive(); + placementReconcileInFlight = current; + const clearCurrent = () => { + if (placementReconcileInFlight === current) { + placementReconcileInFlight = undefined; + } + }; + void current.then(clearCurrent, (error: unknown) => { + params.warn(`Worker placement reconcile sweep failed: ${formatErrorMessage(error)}`); + clearCurrent(); + }); + return current; + }; + const sidecar: WorkerPlacementSidecar = { + stop: async () => { + if (stopped) { + return; + } + stopped = true; + clearInterval(placementReconcileInterval); + placementReconcileInterval = undefined; + uninstallPlacementAdmission(); + uninstallPlacementResetGuard(); + const environmentStop = params.environments.stop(); + const stopResults = await Promise.allSettled([ + ...(placementReconcileInFlight ? [placementReconcileInFlight] : []), + environmentStop, + ]); + const environmentStopResult = stopResults.at(-1); + if (environmentStopResult?.status === "rejected") { + throw environmentStopResult.reason; + } + }, + }; + // Close must see the drain handle before reconciliation can yield. + hooks.registerSidecar(sidecar); + // Track startup reconciliation in the shared in-flight slot so a concurrent + // close prelude drains it before uninstalling guards and stopping environments. + const startupReconcile = dispatchService.reconcile(); + placementReconcileInFlight = startupReconcile; + try { + try { + await startupReconcile; + } finally { + if (placementReconcileInFlight === startupReconcile) { + placementReconcileInFlight = undefined; + } + } + if (hooks.isClosePreludeStarted()) { + await sidecar.stop(); + return null; + } + params.environments.start(); + placementReconcileInterval = setInterval( + () => void reconcileActivePlacements(), + WORKER_PLACEMENT_RECONCILE_INTERVAL_MS, + ); + placementReconcileInterval.unref?.(); + return sidecar; + } catch (error) { + await sidecar.stop(); + throw error; + } + }; + return { dispatchService, admissionProvider, placements: params.placements, startRuntime }; +} diff --git a/src/gateway/server.impl.ts b/src/gateway/server.impl.ts index d63c7aaf761b..570b3f0e6070 100644 --- a/src/gateway/server.impl.ts +++ b/src/gateway/server.impl.ts @@ -85,6 +85,7 @@ import { createAuthRateLimiter, type AuthRateLimiter } from "./auth-rate-limit.j import { resolveGatewayAuth } from "./auth.js"; import type { RestartRecoveryCandidate } from "./chat-abort.js"; import type { ExecApprovalManager } from "./exec-approval-manager.js"; +import { revokeAttachGrantsForSession } from "./mcp-grant-store.js"; import { ADMIN_SCOPE } from "./method-scopes.js"; import { STARTUP_UNAVAILABLE_GATEWAY_METHODS, @@ -164,6 +165,9 @@ const loadGatewayModelCatalogModule = createLazyRuntimeModule( const loadWorkerEnvironmentStartupModule = createLazyRuntimeModule( () => import("./server-worker-environment-startup.js"), ); +const loadWorkerPlacementStartupModule = createLazyRuntimeModule( + () => import("./server-worker-placement-startup.js"), +); export async function resetModelCatalogCacheForTest(): Promise { const { resetModelCatalogCacheForTest: resetModelCatalogCacheForTestLocal } = @@ -177,6 +181,21 @@ const MAX_MEDIA_TTL_HOURS = 24 * 7; const POST_READY_MAINTENANCE_DELAY_MS = 250; const RETAINED_PLUGIN_CLEANUP_DELAY_MS = 30_000; +function approvalRequestTargetsSession( + request: unknown, + sessionKeys: ReadonlySet, + sessionId: string, +): boolean { + if (typeof request !== "object" || request === null) { + return false; + } + const record = request as { sessionKey?: unknown; sessionId?: unknown }; + return ( + (typeof record.sessionId === "string" && record.sessionId === sessionId) || + (typeof record.sessionKey === "string" && sessionKeys.has(record.sessionKey)) + ); +} + type GatewayStartupChannelPlugin = { id: ChannelId; gatewayMethods?: readonly string[]; @@ -901,9 +920,12 @@ export async function startGatewayServer( } let { pluginRegistry, baseGatewayMethods } = pluginBootstrap; // Unconfigured clean installs get no service; durable rows still need list/status projection. + const hasConfiguredWorkerProfiles = + Object.keys(gatewayPluginConfigAtStart.cloudWorkers?.profiles ?? {}).length > 0; const shouldStartWorkerEnvironmentService = - Object.keys(gatewayPluginConfigAtStart.cloudWorkers?.profiles ?? {}).length > 0 || - Boolean(workerEnvironmentStartup?.records.length); + hasConfiguredWorkerProfiles || + Boolean(workerEnvironmentStartup?.records.length) || + Boolean(workerEnvironmentStartup?.hasNonlocalPlacementRecords); let resolveWorkerGatewayEndpoint: () => | { host: "127.0.0.1" | "::1"; port: number } | undefined = () => undefined; @@ -920,6 +942,30 @@ export async function startGatewayServer( }) : {}; const { workerEnvironmentService, workerLiveEvents } = workerEnvironmentRuntime; + // Assigned once approval managers exist; placement dispatch must not run before then. + let revokeWorkerDispatchSessionAuthority = (_params: { + sessionId: string; + sessionKeys: readonly string[]; + }): void => { + throw new Error("Worker dispatch authority revocation is not ready"); + }; + const workerPlacementRuntime = + workerEnvironmentService && workerEnvironmentStartup + ? await startupTrace.measure("worker-environments.placement-runtime", async () => { + const placementModule = await loadWorkerPlacementStartupModule(); + return placementModule.createGatewayWorkerPlacementRuntime({ + placements: workerEnvironmentStartup.placementStore, + environments: workerEnvironmentService, + admitNewPlacements: hasConfiguredWorkerProfiles, + revokeSessionAuthority: (request) => revokeWorkerDispatchSessionAuthority(request), + warn: (message) => log.warn(message), + }); + }) + : undefined; + // Without configured profiles, existing placements still reconcile but new dispatches stay off. + const workerPlacementDispatchAvailable = hasConfiguredWorkerProfiles + ? workerPlacementRuntime?.dispatchService + : undefined; const channelLogs = Object.fromEntries( listGatewayStartupChannelPlugins().map((plugin) => [plugin.id, logChannels.child(plugin.id)]), ) as Record>; @@ -937,7 +983,9 @@ export async function startGatewayServer( return methods; }; const listActiveGatewayMethods = (nextBaseGatewayMethods: string[]) => - uniqueStrings([...nextBaseGatewayMethods, ...listStartupChannelGatewayMethods()]); + uniqueStrings([...nextBaseGatewayMethods, ...listStartupChannelGatewayMethods()]).filter( + (method) => workerPlacementDispatchAvailable || method !== "sessions.dispatch", + ); const runtimeConfig = await startupTrace.measure("runtime.config", async () => { const { resolveGatewayRuntimeConfig } = await import("./server-runtime-config.js"); return resolveGatewayRuntimeConfig({ @@ -1549,6 +1597,22 @@ export async function startGatewayServer( }); approvalManagersForReplay.exec = execApprovalManager; approvalManagersForReplay.plugin = pluginApprovalManager; + revokeWorkerDispatchSessionAuthority = ({ sessionId, sessionKeys }) => { + const keys = new Set(sessionKeys); + for (const sessionKey of keys) { + revokeAttachGrantsForSession(sessionKey); + } + for (const record of execApprovalManager.listPendingRecords()) { + if (approvalRequestTargetsSession(record.request, keys, sessionId)) { + execApprovalManager.expire(record.id, "worker-dispatch"); + } + } + for (const record of pluginApprovalManager.listPendingRecords()) { + if (approvalRequestTargetsSession(record.request, keys, sessionId)) { + pluginApprovalManager.expire(record.id, "worker-dispatch"); + } + } + }; const attachedGatewayExtraHandlers: GatewayRequestHandlers = { ...pluginRegistry.gatewayHandlers, ...extraHandlers, @@ -1568,8 +1632,10 @@ export async function startGatewayServer( } const coreDescriptors = createCoreGatewayMethodDescriptors(coreDescriptorHandlers).filter( (descriptor) => - workerEnvironmentService || - (descriptor.name !== "environments.create" && descriptor.name !== "environments.destroy"), + (workerEnvironmentService || + (descriptor.name !== "environments.create" && + descriptor.name !== "environments.destroy")) && + (workerPlacementDispatchAvailable || descriptor.name !== "sessions.dispatch"), ); return createGatewayMethodRegistry([ ...coreDescriptors, @@ -1816,6 +1882,12 @@ export async function startGatewayServer( }, nodeRegistry, ...(workerEnvironmentService ? { workerEnvironmentService } : {}), + ...(workerPlacementRuntime + ? { workerSessionPlacementService: workerPlacementRuntime.placements } + : {}), + ...(workerPlacementDispatchAvailable + ? { workerPlacementDispatchService: workerPlacementDispatchAvailable } + : {}), terminalSessions, agentRunSeq, chatAbortControllers, @@ -2054,17 +2126,19 @@ export async function startGatewayServer( runtimeState.gatewayLifetimeSidecars = []; } }, - ...(workerEnvironmentService + ...(workerPlacementRuntime ? { - startWorkerEnvironmentRuntime: () => { + startWorkerEnvironmentRuntime: async () => { if (closePreludeStarted) { return null; } - const sidecar = { stop: () => workerEnvironmentService.stop() }; - // Close must see the drain handle before reconciliation can yield. - runtimeState.gatewayLifetimeSidecars.push(sidecar); - workerEnvironmentService.start(); - return sidecar; + return await workerPlacementRuntime.startRuntime({ + isClosePreludeStarted: () => closePreludeStarted, + // Close must see the drain handle before reconciliation can yield. + registerSidecar: (sidecar) => { + runtimeState.gatewayLifetimeSidecars.push(sidecar); + }, + }); }, } : {}), diff --git a/src/gateway/server.sessions.placement-projection.test.ts b/src/gateway/server.sessions.placement-projection.test.ts new file mode 100644 index 000000000000..9582d01cd81b --- /dev/null +++ b/src/gateway/server.sessions.placement-projection.test.ts @@ -0,0 +1,124 @@ +import { expect, test, vi } from "vitest"; +import type { GatewaySessionRow } from "./session-utils.types.js"; +import { writeSessionStore } from "./test-helpers.js"; +import { + directSessionReq, + setupGatewaySessionsTestHarness, +} from "./test/server-sessions.test-helpers.js"; +import type { WorkerSessionPlacementReader } from "./worker-environments/placement-projector.js"; +import type { WorkerSessionPlacementRecord } from "./worker-environments/placement-store.js"; + +const { createSessionStoreDir } = setupGatewaySessionsTestHarness(); + +function activePlacementRecord(): WorkerSessionPlacementRecord { + return { + sessionId: "sess-main", + agentId: "main", + sessionKey: "agent:main:main", + state: "active", + environmentId: "env-placement", + generation: 7, + activeOwnerEpoch: 12, + workspaceBaseManifestRef: "manifest-base", + remoteWorkspaceDir: "/workspace/main", + workerBundleHash: ["a", "b"].join("").repeat(32), + lastTranscriptAckCursor: 23, + lastLiveEventAckCursor: 9, + recoveryError: null, + turnClaim: null, + createdAtMs: 100, + updatedAtMs: 300, + stateChangedAtMs: 200, + }; +} + +async function seedSessionRows(): Promise { + await createSessionStoreDir(); + await writeSessionStore({ + entries: { + main: { sessionId: "sess-main", updatedAt: 200 }, + "agent:main:other": { sessionId: "sess-other", updatedAt: 100 }, + }, + }); +} + +test("sessions.list omits placement when the worker placement service is disabled", async () => { + await seedSessionRows(); + + const result = await directSessionReq<{ sessions: GatewaySessionRow[] }>("sessions.list", {}); + + expect(result.ok).toBe(true); + expect(result.payload?.sessions).toHaveLength(2); + expect(result.payload?.sessions.every((session) => session.placement === undefined)).toBe(true); +}); + +test("sessions.list batch-projects durable worker placement", async () => { + await seedSessionRows(); + const placement = activePlacementRecord(); + const getMany = vi.fn((sessionIds) => { + expect(sessionIds).toEqual(expect.arrayContaining(["sess-main", "sess-other"])); + return new Map([[placement.sessionId, placement]]); + }); + + const result = await directSessionReq<{ sessions: GatewaySessionRow[] }>( + "sessions.list", + {}, + { + context: { workerSessionPlacementService: { getMany } }, + }, + ); + + expect(result.ok).toBe(true); + expect(getMany).toHaveBeenCalledTimes(1); + const main = result.payload?.sessions.find((session) => session.sessionId === "sess-main"); + const other = result.payload?.sessions.find((session) => session.sessionId === "sess-other"); + expect(main?.placement).toEqual({ + state: "active", + environmentId: "env-placement", + generation: 7, + activeOwnerEpoch: 12, + workspaceBaseManifestRef: "manifest-base", + remoteWorkspaceDir: "/workspace/main", + workerBundleHash: ["a", "b"].join("").repeat(32), + lastTranscriptAckCursor: 23, + lastLiveEventAckCursor: 9, + createdAtMs: 100, + updatedAtMs: 300, + stateChangedAtMs: 200, + }); + expect(other?.placement).toBeUndefined(); +}); + +test("sessions.describe projects durable worker placement", async () => { + await seedSessionRows(); + const placement = activePlacementRecord(); + const getMany = vi.fn((sessionIds) => { + expect(sessionIds).toEqual(["sess-main"]); + return new Map([[placement.sessionId, placement]]); + }); + + const result = await directSessionReq<{ session: GatewaySessionRow | null }>( + "sessions.describe", + { key: "main" }, + { + context: { workerSessionPlacementService: { getMany } }, + }, + ); + + expect(result.ok).toBe(true); + expect(getMany).toHaveBeenCalledTimes(1); + expect(result.payload?.session?.placement).toEqual({ + state: "active", + environmentId: "env-placement", + generation: 7, + activeOwnerEpoch: 12, + workspaceBaseManifestRef: "manifest-base", + remoteWorkspaceDir: "/workspace/main", + workerBundleHash: ["a", "b"].join("").repeat(32), + lastTranscriptAckCursor: 23, + lastLiveEventAckCursor: 9, + createdAtMs: 100, + updatedAtMs: 300, + stateChangedAtMs: 200, + }); +}); diff --git a/src/gateway/server.sessions.worker-placement-lifecycle.test.ts b/src/gateway/server.sessions.worker-placement-lifecycle.test.ts new file mode 100644 index 000000000000..8bb3d1ff02d1 --- /dev/null +++ b/src/gateway/server.sessions.worker-placement-lifecycle.test.ts @@ -0,0 +1,246 @@ +import { afterEach, expect, test } from "vitest"; +import { installSessionPlacementResetGuard } from "../agents/session-placement-admission.js"; +import { closeOpenClawStateDatabaseForTest } from "../state/openclaw-state-db.js"; +import { loadSessionEntry } from "./session-utils.js"; +import { embeddedRunMock, writeSessionStore } from "./test-helpers.js"; +import { + directSessionReq, + sessionStoreEntry, + setupGatewaySessionsTestHarness, +} from "./test/server-sessions.test-helpers.js"; +import type { WorkerSessionPlacementReader } from "./worker-environments/placement-projector.js"; +import type { WorkerSessionPlacementRecord } from "./worker-environments/placement-store.js"; + +const { createSessionStoreDir, seedActiveMainSession } = setupGatewaySessionsTestHarness(); +let uninstallResetGuard: (() => void) | undefined; + +afterEach(() => { + uninstallResetGuard?.(); + uninstallResetGuard = undefined; + closeOpenClawStateDatabaseForTest(); +}); + +function placementRecord( + sessionId: string, + state: "active" | "local", +): WorkerSessionPlacementRecord { + const identity = { + sessionId, + agentId: "main", + sessionKey: "agent:main:worker-session", + turnClaim: null, + createdAtMs: 1, + updatedAtMs: 2, + stateChangedAtMs: 2, + }; + if (state === "active") { + return { + ...identity, + state, + generation: 2, + environmentId: "worker-environment", + activeOwnerEpoch: 1, + workspaceBaseManifestRef: "manifest-ref", + remoteWorkspaceDir: "/workspace", + workerBundleHash: "bundle-hash", + lastTranscriptAckCursor: null, + lastLiveEventAckCursor: null, + recoveryError: null, + }; + } + return { + ...identity, + state, + generation: 0, + environmentId: null, + activeOwnerEpoch: null, + workspaceBaseManifestRef: null, + remoteWorkspaceDir: null, + workerBundleHash: null, + lastTranscriptAckCursor: null, + lastLiveEventAckCursor: null, + recoveryError: null, + }; +} + +function terminalPlacementRecord( + sessionId: string, + state: "failed" | "reclaimed", +): WorkerSessionPlacementRecord { + const terminalMetadata = { + environmentId: "worker-environment", + activeOwnerEpoch: 1, + workspaceBaseManifestRef: "manifest-ref", + remoteWorkspaceDir: "/workspace", + workerBundleHash: "bundle-hash", + lastTranscriptAckCursor: null, + lastLiveEventAckCursor: null, + }; + const identity = { + sessionId, + agentId: "main", + sessionKey: "agent:main:worker-session", + generation: 2, + turnClaim: null, + createdAtMs: 1, + updatedAtMs: 2, + stateChangedAtMs: 2, + }; + if (state === "failed") { + return { + ...identity, + ...terminalMetadata, + state, + recoveryError: "worker recovery stopped", + }; + } + return { + ...identity, + ...terminalMetadata, + state, + recoveryError: null, + }; +} + +function sequencedPlacementReader( + records: readonly WorkerSessionPlacementRecord[], +): WorkerSessionPlacementReader { + let readIndex = 0; + return { + getMany(sessionIds) { + const record = records[Math.min(readIndex, records.length - 1)]; + readIndex += 1; + const result = new Map(); + if (record && sessionIds.includes(record.sessionId)) { + result.set(record.sessionId, record); + } + return result; + }, + }; +} + +test("sessions.reset rechecks worker placement inside the lifecycle fence", async () => { + await seedActiveMainSession(); + let resetGuardReadCount = 0; + uninstallResetGuard = installSessionPlacementResetGuard((sessionId) => { + expect(sessionId).toBe("sess-main"); + resetGuardReadCount += 1; + return resetGuardReadCount === 1 ? undefined : "cloud worker placement is active"; + }); + + const reset = await directSessionReq("sessions.reset", { key: "main" }); + + expect(reset.ok).toBe(false); + expect(reset.error?.message).toContain("cloud worker placement is active"); + expect(resetGuardReadCount).toBe(2); + expect(loadSessionEntry("main").entry?.sessionId).toBe("sess-main"); + expect(embeddedRunMock.abortCalls).toEqual([]); +}); + +test("sessions.delete rechecks worker placement before destructive cleanup", async () => { + await createSessionStoreDir(); + const sessionKey = "discord:group:worker-session"; + const sessionId = "sess-worker-delete"; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + const placementReader = sequencedPlacementReader([ + placementRecord(sessionId, "local"), + placementRecord(sessionId, "active"), + ]); + + const deleted = await directSessionReq( + "sessions.delete", + { key: sessionKey }, + { + context: { workerSessionPlacementService: placementReader }, + }, + ); + + expect(deleted.ok).toBe(false); + expect(deleted.error?.message).toContain("cloud worker placement is active"); + expect(loadSessionEntry(sessionKey).entry?.sessionId).toBe(sessionId); + expect(embeddedRunMock.abortCalls).toEqual([]); +}); + +test("sessions.delete rejects failed placement with unresolved worker ownership", async () => { + await createSessionStoreDir(); + const sessionKey = "discord:group:failed-worker-session"; + const sessionId = "sess-failed-worker-delete"; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + const placementReader = sequencedPlacementReader([terminalPlacementRecord(sessionId, "failed")]); + + const deleted = await directSessionReq( + "sessions.delete", + { key: sessionKey }, + { + context: { workerSessionPlacementService: placementReader }, + }, + ); + + expect(deleted.ok).toBe(false); + expect(deleted.error?.message).toContain("cloud worker placement is failed"); + expect(loadSessionEntry(sessionKey).entry?.sessionId).toBe(sessionId); + expect(embeddedRunMock.abortCalls).toEqual([]); +}); + +test("sessions.delete allows reclaimed placement with no live worker owner", async () => { + await createSessionStoreDir(); + const sessionKey = "discord:group:reclaimed-worker-session"; + const sessionId = "sess-reclaimed-worker-delete"; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + const placementReader = sequencedPlacementReader([ + terminalPlacementRecord(sessionId, "reclaimed"), + ]); + + const deleted = await directSessionReq( + "sessions.delete", + { key: sessionKey }, + { + context: { workerSessionPlacementService: placementReader }, + }, + ); + + expect(deleted.ok).toBe(true); + expect(deleted.payload).toMatchObject({ ok: true, deleted: true }); + expect(loadSessionEntry(sessionKey).entry).toBeUndefined(); +}); + +test("sessions.compaction.restore rechecks worker placement inside the lifecycle fence", async () => { + await createSessionStoreDir(); + const sessionKey = "discord:group:worker-restore"; + const sessionId = "sess-worker-restore"; + const checkpointId = "checkpoint-worker-restore"; + await writeSessionStore({ + entries: { + [sessionKey]: sessionStoreEntry(sessionId, { + compactionCheckpoints: [ + { + checkpointId, + sessionKey, + sessionId, + createdAt: 1, + reason: "manual", + preCompaction: { sessionId }, + postCompaction: { sessionId }, + }, + ], + }), + }, + }); + const placementReader = sequencedPlacementReader([ + placementRecord(sessionId, "local"), + placementRecord(sessionId, "active"), + ]); + + const restored = await directSessionReq( + "sessions.compaction.restore", + { key: sessionKey, checkpointId }, + { + context: { workerSessionPlacementService: placementReader }, + }, + ); + + expect(restored.ok).toBe(false); + expect(restored.error?.message).toContain("cloud worker placement is active"); + expect(loadSessionEntry(sessionKey).entry?.sessionId).toBe(sessionId); + expect(embeddedRunMock.abortCalls).toEqual([]); +}); diff --git a/src/gateway/server/ws-connection/message-handler.worker.test.ts b/src/gateway/server/ws-connection/message-handler.worker.test.ts index 00c635a2fb5a..5a8bcd6a4660 100644 --- a/src/gateway/server/ws-connection/message-handler.worker.test.ts +++ b/src/gateway/server/ws-connection/message-handler.worker.test.ts @@ -53,6 +53,7 @@ const WORKER_CONNECT: WorkerConnectParams = { environmentId: "worker-1", credential: CREDENTIAL, sessionId: null, + runId: null, ownerEpoch: 1, rpcSetVersion: 1, handshake: HANDSHAKE, @@ -63,6 +64,7 @@ const IDENTITY: WorkerConnectionIdentity = { credentialHash: "h".repeat(43), bundleHash: HANDSHAKE.bundleHash, sessionId: null, + runId: null, ownerEpoch: 1, rpcSetVersion: 1, protocolFeatures: [...HANDSHAKE.protocolFeatures], @@ -90,6 +92,7 @@ const LIVE_EVENT = { const ATTACHED_IDENTITY: WorkerConnectionIdentity = { ...IDENTITY, sessionId: "session-1", + runId: "run-1", }; const INFERENCE_IDS = { runEpoch: 1, diff --git a/src/gateway/session-message-events.test.ts b/src/gateway/session-message-events.test.ts index 0fa3bdc83392..5e7fcb0fccf0 100644 --- a/src/gateway/session-message-events.test.ts +++ b/src/gateway/session-message-events.test.ts @@ -1246,6 +1246,7 @@ describe("session.message websocket events", () => { credentialHash: ["fanout", "credential", "hash"].join("-"), bundleHash: "f".repeat(64), sessionId, + runId: "run-fanout", ownerEpoch: 4, rpcSetVersion: 1, protocolFeatures: ["worker-live-event-v1", "worker-transcript-commit-v1"], diff --git a/src/gateway/session-reset-service.ts b/src/gateway/session-reset-service.ts index e174f4d59ae2..6985ff0266bd 100644 --- a/src/gateway/session-reset-service.ts +++ b/src/gateway/session-reset-service.ts @@ -20,6 +20,7 @@ import { clearAllCliSessions } from "../agents/cli-session.js"; import { abortEmbeddedAgentRun, waitForEmbeddedAgentRunEnd } from "../agents/embedded-agent.js"; import { resetRegisteredAgentHarnessSessions } from "../agents/harness/registry.js"; import { resolveSessionModelRef } from "../agents/session-model-ref.js"; +import { resolveSessionPlacementResetBlock } from "../agents/session-placement-admission.js"; import { stopSubagentsForRequester } from "../auto-reply/reply/abort.js"; import { buildSessionEndHookPayload, @@ -927,6 +928,18 @@ export async function performGatewaySessionReset(params: { error: errorShape(ErrorCodes.INVALID_REQUEST, MODEL_SELECTION_LOCKED_RESET_MESSAGE), }; } + const initialPlacementBlock = initialResetEntry?.sessionId + ? resolveSessionPlacementResetBlock(initialResetEntry.sessionId) + : undefined; + if (initialPlacementBlock) { + return { + ok: false, + error: errorShape( + ErrorCodes.INVALID_REQUEST, + `Session ${params.key} cannot reset while ${initialPlacementBlock}.`, + ), + }; + } const resetLifecycleIdentities = [ resetTarget.target.canonicalKey, params.key, @@ -979,6 +992,18 @@ export async function performGatewaySessionReset(params: { params.key, requestedAgentId ? { agentId: requestedAgentId } : undefined, ); + const placementBlock = entry?.sessionId + ? resolveSessionPlacementResetBlock(entry.sessionId) + : undefined; + if (placementBlock) { + return { + ok: false, + error: errorShape( + ErrorCodes.INVALID_REQUEST, + `Session ${params.key} cannot reset while ${placementBlock}.`, + ), + }; + } const archivedSessionError = resolveSessionWorkStartError(canonicalKey, entry); if (archivedSessionError) { return { diff --git a/src/gateway/session-utils.types.ts b/src/gateway/session-utils.types.ts index b9d38bc39268..8a079e2e5e7a 100644 --- a/src/gateway/session-utils.types.ts +++ b/src/gateway/session-utils.types.ts @@ -1,6 +1,7 @@ // Shared Gateway session projection types. // Keeps server methods and Control UI payloads aligned. import type { FastMode } from "@openclaw/normalization-core/string-coerce"; +import type { SessionPlacement } from "../../packages/gateway-protocol/src/index.js"; import type { ChatType } from "../channels/chat-type.js"; import type { SessionCompactionCheckpoint, @@ -77,6 +78,7 @@ export type GatewaySessionRow = { lastReadAt?: number; lastActivityAt?: number; sessionId?: string; + placement?: SessionPlacement; systemSent?: boolean; abortedLastRun?: boolean; thinkingLevel?: string; diff --git a/src/gateway/worker-environments/admission.test.ts b/src/gateway/worker-environments/admission.test.ts index 7bad2dd5b11c..ac0ca3174538 100644 --- a/src/gateway/worker-environments/admission.test.ts +++ b/src/gateway/worker-environments/admission.test.ts @@ -52,15 +52,23 @@ describe("worker admission", () => { } as WorkerEnvironmentStore; }); - const admission = (overrides: Partial = {}) => ({ - environmentId: "worker-1", - credential: CREDENTIAL, - sessionId: null, - ownerEpoch: 1, - rpcSetVersion: 1, - handshake: RECEIPT, - ...overrides, - }); + type AdmissionCommon = Omit; + type AdmissionOverrides = + | (Partial & { sessionId?: null; runId?: null }) + | (Partial & { sessionId: string; runId: string }); + const admission = (overrides: AdmissionOverrides = {}): WorkerConnectParams["admission"] => { + const common = { + environmentId: overrides.environmentId ?? "worker-1", + credential: overrides.credential ?? CREDENTIAL, + ownerEpoch: overrides.ownerEpoch ?? 1, + rpcSetVersion: overrides.rpcSetVersion ?? 1, + handshake: overrides.handshake ?? RECEIPT, + }; + if (typeof overrides.sessionId === "string" && typeof overrides.runId === "string") { + return { ...common, sessionId: overrides.sessionId, runId: overrides.runId }; + } + return { ...common, sessionId: null, runId: null }; + }; const admit = (workerAdmission = admission(), expectedBuild: typeof RECEIPT = RECEIPT) => admitWorkerConnection({ store, admission: workerAdmission, expectedBuild, nowMs }); @@ -72,6 +80,7 @@ describe("worker admission", () => { identity: { environmentId: "worker-1", sessionId: null, + runId: null, ownerEpoch: 1, rpcSetVersion: 1, protocolFeatures: ["worker-heartbeat-v1"], @@ -86,7 +95,7 @@ describe("worker admission", () => { ["environment-mismatch", () => admission({ environmentId: " worker-1 " })], ["bundle-mismatch", () => admission({ handshake: { ...RECEIPT, bundleHash: "b".repeat(64) } })], ["version-mismatch", () => admission({ handshake: { ...RECEIPT, openclawVersion: "other" } })], - ["session-mismatch", () => admission({ sessionId: "session-other" })], + ["session-mismatch", () => admission({ sessionId: "session-other", runId: "run-other" })], ["owner-epoch-mismatch", () => admission({ ownerEpoch: 2 })], ["rpc-set-mismatch", () => admission({ rpcSetVersion: 2 })], [ diff --git a/src/gateway/worker-environments/admission.ts b/src/gateway/worker-environments/admission.ts index 39378bff50df..f92106e24f89 100644 --- a/src/gateway/worker-environments/admission.ts +++ b/src/gateway/worker-environments/admission.ts @@ -93,6 +93,9 @@ export function admitWorkerConnection(params: { if (admission.sessionId !== credential.sessionId) { return { ok: false, reason: "session-mismatch" }; } + if ((admission.sessionId === null) !== (admission.runId === null)) { + return { ok: false, reason: "session-mismatch" }; + } if ( admission.ownerEpoch !== credential.ownerEpoch || admission.ownerEpoch !== environment.ownerEpoch @@ -121,6 +124,7 @@ export function admitWorkerConnection(params: { credentialHash: credential.credentialHash, bundleHash: credential.bundleHash, sessionId: credential.sessionId, + runId: admission.runId, ownerEpoch: credential.ownerEpoch, rpcSetVersion: credential.rpcSetVersion, protocolFeatures: [...environment.bootstrapReceipt.protocolFeatures], diff --git a/src/gateway/worker-environments/connection-identity.ts b/src/gateway/worker-environments/connection-identity.ts index b695f01a2a8c..e23f940ec136 100644 --- a/src/gateway/worker-environments/connection-identity.ts +++ b/src/gateway/worker-environments/connection-identity.ts @@ -4,6 +4,7 @@ export type WorkerConnectionIdentity = { credentialHash: string; bundleHash: string; sessionId: string | null; + runId: string | null; ownerEpoch: number; rpcSetVersion: number; protocolFeatures: string[]; diff --git a/src/gateway/worker-environments/inference-runtime.test.ts b/src/gateway/worker-environments/inference-runtime.test.ts index 5c7ca7040461..2827226d9962 100644 --- a/src/gateway/worker-environments/inference-runtime.test.ts +++ b/src/gateway/worker-environments/inference-runtime.test.ts @@ -73,6 +73,7 @@ const identity: WorkerConnectionIdentity = { credentialHash: ["credential", "hash", "runtime", "test"].join("-"), bundleHash: "bundle-hash-runtime-test", sessionId: SESSION_ID, + runId: "run-runtime-test", ownerEpoch: 3, rpcSetVersion: 1, protocolFeatures: ["worker-inference-v1"], diff --git a/src/gateway/worker-environments/inference-store.test.ts b/src/gateway/worker-environments/inference-store.test.ts index 1127d7a8b5fa..9b7af5e37fe1 100644 --- a/src/gateway/worker-environments/inference-store.test.ts +++ b/src/gateway/worker-environments/inference-store.test.ts @@ -41,6 +41,7 @@ const IDENTITY: WorkerConnectionIdentity = { credentialHash: ["fixture", "digest"].join("-"), bundleHash: ["fixture", "bundle", "digest"].join("-"), sessionId: REQUEST.sessionId, + runId: REQUEST.runId, ownerEpoch: REQUEST.runEpoch, rpcSetVersion: 1, protocolFeatures: ["worker-inference-v1"], diff --git a/src/gateway/worker-environments/inference.test.ts b/src/gateway/worker-environments/inference.test.ts index 76c2069229fe..7a0e3e453b02 100644 --- a/src/gateway/worker-environments/inference.test.ts +++ b/src/gateway/worker-environments/inference.test.ts @@ -27,6 +27,7 @@ const IDENTITY: WorkerConnectionIdentity = { credentialHash: "d", bundleHash: "b", sessionId: REQUEST.sessionId, + runId: REQUEST.runId, ownerEpoch: REQUEST.runEpoch, rpcSetVersion: 1, protocolFeatures: ["worker-inference-v1"], diff --git a/src/gateway/worker-environments/live-events.test.ts b/src/gateway/worker-environments/live-events.test.ts index ef0828daed03..96f64cbbe3d5 100644 --- a/src/gateway/worker-environments/live-events.test.ts +++ b/src/gateway/worker-environments/live-events.test.ts @@ -12,6 +12,7 @@ import { claimAgentRunContext, clearAgentRunContext, emitAgentEvent, + getAgentEventLifecycleGeneration, getAgentRunContext, onAgentRuntimeEvent, sweepStaleRunContexts, @@ -33,6 +34,7 @@ const ID: Identity = { credentialHash: ["credential", "hash", "live"].join("-"), bundleHash: "b".repeat(64), sessionId: SID, + runId: RUN, ownerEpoch: EPOCH, rpcSetVersion: 1, protocolFeatures: ["worker-live-event-v1"], @@ -455,6 +457,51 @@ describe("worker live events", () => { fail(msg(1, "pending", 0, "run-pending"), "invalid-event"); }); + it("adopts a compatible pre-registered gateway run context", () => { + const lifecycleGeneration = getAgentEventLifecycleGeneration(); + claimAgentRunContext(RUN, { + ...LOCAL, + isControlUiVisible: false, + lifecycleGeneration, + }); + + ack(msg(1, "worker")); + + expect(getAgentRunContext(RUN)).toMatchObject({ + ...LOCAL, + isControlUiVisible: false, + lifecycleGeneration, + projectSessionActive: true, + }); + expect(deltas()).toEqual(["worker"]); + }); + + it("rejects pre-registered gateway run contexts with mismatched identity", () => { + const lifecycleGeneration = getAgentEventLifecycleGeneration(); + const mismatches: Array<{ + context: Parameters[1]; + name: string; + }> = [ + { name: "session-id", context: { ...LOCAL, sessionId: `${SID}-other` } }, + { name: "session-key", context: { ...LOCAL, sessionKey: `${KEY}-other` } }, + { name: "agent-id", context: { ...LOCAL, agentId: "other" } }, + { name: "lifecycle", context: { ...LOCAL, lifecycleGeneration: "other-lifecycle" } }, + { name: "visibility", context: { ...LOCAL, isControlUiVisible: true } }, + ]; + + for (const mismatch of mismatches) { + const runId = `run-mismatch-${mismatch.name}`; + claimAgentRunContext(runId, { + isControlUiVisible: false, + lifecycleGeneration, + ...mismatch.context, + }); + fail(msg(1, "blocked", 0, runId), "invalid-event"); + clearAgentRunContext(runId); + } + expect(events).toEqual([]); + }); + it("keeps run ids exclusive", () => { const local = "run-local-first"; claimAgentRunContext(local, LOCAL); diff --git a/src/gateway/worker-environments/live-events.ts b/src/gateway/worker-environments/live-events.ts index 2baecac86bab..e0e6a1472b6b 100644 --- a/src/gateway/worker-environments/live-events.ts +++ b/src/gateway/worker-environments/live-events.ts @@ -569,11 +569,19 @@ export function createWorkerLiveEventReceiver(options: WorkerLiveEventReceiverOp } const lifecycleGeneration = getAgentEventLifecycleGeneration(); const existingContext = getAgentRunContext(runId); - if (existingContext?.lifecycleGeneration === lifecycleGeneration) { - return invalidEvent(); - } // Turn placement owns wider visibility; otherwise scope to this session. const controlUiVisible = false; + const adoptExistingUnowned = existingContext !== undefined; + if ( + existingContext && + (existingContext.sessionId !== window.sessionId || + existingContext.sessionKey !== window.target.sessionKey || + existingContext.agentId !== window.target.agentId || + existingContext.lifecycleGeneration !== lifecycleGeneration || + existingContext.isControlUiVisible !== controlUiVisible) + ) { + return invalidEvent(); + } const claimId = claimAgentRunContext( runId, { @@ -585,6 +593,7 @@ export function createWorkerLiveEventReceiver(options: WorkerLiveEventReceiverOp sessionKey: window.target.sessionKey, }, { + adoptExistingUnowned, exclusive: true, onClearRequested: (clearedClaimId) => { if (window.activeRuns.get(runId)?.claimId === clearedClaimId) { diff --git a/src/gateway/worker-environments/placement-dispatch-failure.ts b/src/gateway/worker-environments/placement-dispatch-failure.ts new file mode 100644 index 000000000000..9a8e3bc13e01 --- /dev/null +++ b/src/gateway/worker-environments/placement-dispatch-failure.ts @@ -0,0 +1,262 @@ +import { truncateUtf16Safe } from "@openclaw/normalization-core/utf16-slice"; +import { formatErrorMessage } from "../../infra/errors.js"; +import { redactSensitiveText } from "../../logging/redact.js"; +import type { + createWorkerSessionPlacementStore, + WorkerSessionPlacementRecord, +} from "./placement-store.js"; +import type { WorkerEnvironmentService } from "./service.js"; + +export type WorkerDispatchPlacement = WorkerSessionPlacementRecord; +export type WorkerActiveDispatchPlacement = Extract< + WorkerSessionPlacementRecord, + { state: "active" } +>; +export type WorkerFailedDispatchPlacement = Extract; +export type WorkerStartingDispatchPlacement = Extract< + WorkerDispatchPlacement, + { state: "starting" } +>; +type WorkerDrainingDispatchPlacement = Extract; +type WorkerReconcilingDispatchPlacement = Extract< + WorkerDispatchPlacement, + { state: "reconciling" } +>; + +export type WorkerDispatchPlacementStore = Pick< + ReturnType, + | "adoptActive" + | "fail" + | "get" + | "listForReconcile" + | "startDispatch" + | "startDrain" + | "startReconcile" + | "transition" +>; + +export type WorkerDispatchEnvironmentService = Pick< + WorkerEnvironmentService, + "attachSession" | "create" | "destroy" | "get" | "reconcileOnce" | "startTunnel" | "stopTunnel" +>; + +export type WorkerActivationBarrier = (params: { + sessionId: string; + sessionKey: string; + agentId: string; + activate: () => WorkerActiveDispatchPlacement; +}) => Promise; + +const RECOVERY_ERROR_LIMIT = 1_024; + +function boundedError(error: unknown): string { + const redacted = redactSensitiveText(formatErrorMessage(error), { mode: "tools" }) + .replace(/\s+/gu, " ") + .trim(); + return truncateUtf16Safe(redacted || "unknown dispatch failure", RECOVERY_ERROR_LIMIT); +} + +export function isUnavailableEnvironment( + environment: NonNullable>, +): boolean { + return ( + environment.state === "draining" || + environment.state === "destroying" || + environment.state === "destroyed" || + environment.state === "failed" || + environment.state === "orphaned" + ); +} + +export function createPlacementFailureActions(deps: { + placements: WorkerDispatchPlacementStore; + environments: WorkerDispatchEnvironmentService; +}) { + const { environments, placements } = deps; + + const updateFailure = ( + placement: WorkerDispatchPlacement, + error: unknown, + ): WorkerDispatchPlacement => + placements.fail({ + sessionId: placement.sessionId, + expectedGeneration: placement.generation, + recoveryError: boundedError(error), + }); + + const cleanupEnvironment = async (params: { + environmentId: string; + ownerEpoch: number | null; + }): Promise => { + const teardownErrors: string[] = []; + try { + await environments.stopTunnel(params.environmentId, params.ownerEpoch ?? undefined); + } catch (error) { + teardownErrors.push(`tunnel stop: ${boundedError(error)}`); + } + try { + await environments.destroy(params.environmentId); + } catch (error) { + teardownErrors.push(`environment destroy: ${boundedError(error)}`); + } + return teardownErrors; + }; + + const teardownEnvironment = async (params: { + placement: WorkerDispatchPlacement; + environmentId: string | null; + ownerEpoch: number | null; + primaryError: unknown; + }): Promise => { + const environmentId = params.environmentId; + const teardownErrors = environmentId + ? await cleanupEnvironment({ + environmentId, + ownerEpoch: params.ownerEpoch, + }) + : []; + const recoveryError = [boundedError(params.primaryError), ...teardownErrors].join("; "); + updateFailure( + params.placement, + new Error(truncateUtf16Safe(recoveryError, RECOVERY_ERROR_LIMIT)), + ); + }; + + const retryFailedTeardown = async (placement: WorkerFailedDispatchPlacement): Promise => { + if (!placement.environmentId) { + return; + } + const environment = environments.get(placement.environmentId); + if ( + !environment || + environment.state === "destroyed" || + environment.state === "failed" || + environment.state === "orphaned" + ) { + return; + } + const teardownErrors = await cleanupEnvironment({ + environmentId: placement.environmentId, + ownerEpoch: placement.activeOwnerEpoch, + }); + if (teardownErrors.length > 0) { + const recoveryError = [placement.recoveryError, ...teardownErrors].filter(Boolean).join("; "); + placements.fail({ + sessionId: placement.sessionId, + expectedGeneration: placement.generation, + recoveryError: truncateUtf16Safe(recoveryError, RECOVERY_ERROR_LIMIT), + }); + } + }; + + const startDrain = ( + placement: WorkerActiveDispatchPlacement, + ): WorkerDrainingDispatchPlacement => { + const draining = placements.startDrain({ + sessionId: placement.sessionId, + environmentId: placement.environmentId, + ownerEpoch: placement.activeOwnerEpoch, + expectedGeneration: placement.generation, + }); + if (draining.state !== "draining") { + throw new Error("Worker placement drain did not produce a draining placement"); + } + return draining; + }; + + const startReconcile = ( + placement: WorkerDrainingDispatchPlacement, + ): WorkerReconcilingDispatchPlacement => { + const reconciling = placements.startReconcile({ + sessionId: placement.sessionId, + environmentId: placement.environmentId, + ownerEpoch: placement.activeOwnerEpoch, + expectedGeneration: placement.generation, + }); + if (reconciling.state !== "reconciling") { + throw new Error("Worker placement reconcile did not produce a reconciling placement"); + } + return reconciling; + }; + + const advanceReclaimed = (placement: WorkerDrainingDispatchPlacement): void => { + // Lost-worker recovery has no live workspace to pull back. Deliberate inbound + // reconciliation remains a separate migration workflow. + const reconciling = startReconcile(placement); + const reclaimed = placements.transition({ + sessionId: reconciling.sessionId, + from: "reconciling", + to: "reclaimed", + expectedGeneration: reconciling.generation, + }); + if (reclaimed.state !== "reclaimed") { + throw new Error("Worker placement reclaim did not produce a reclaimed placement"); + } + }; + + const finishDrainingFailure = ( + placement: WorkerDrainingDispatchPlacement, + error: unknown, + teardownErrors: readonly string[], + ): void => { + const reconciling = startReconcile(placement); + const recoveryError = [boundedError(error), ...teardownErrors].join("; "); + updateFailure(reconciling, new Error(truncateUtf16Safe(recoveryError, RECOVERY_ERROR_LIMIT))); + }; + + const failDraining = async ( + placement: WorkerDrainingDispatchPlacement, + error: unknown, + ): Promise => { + const teardownErrors = await cleanupEnvironment({ + environmentId: placement.environmentId, + ownerEpoch: placement.activeOwnerEpoch, + }); + finishDrainingFailure(placement, error, teardownErrors); + }; + + const reclaimActive = async ( + placement: WorkerActiveDispatchPlacement, + environment: ReturnType, + claimedTurnError: Error, + ): Promise => { + const draining = startDrain(placement); + if (draining.turnClaim) { + await failDraining(draining, claimedTurnError); + return; + } + if (environment && !isUnavailableEnvironment(environment)) { + const teardownErrors = await cleanupEnvironment({ + environmentId: placement.environmentId, + ownerEpoch: placement.activeOwnerEpoch, + }); + if (teardownErrors.length > 0) { + finishDrainingFailure( + draining, + new Error(`Worker reclaim teardown failed: ${teardownErrors.join("; ")}`), + [], + ); + return; + } + } + advanceReclaimed(draining); + }; + + const failActive = async ( + placement: WorkerActiveDispatchPlacement, + error: unknown, + ): Promise => { + const draining = startDrain(placement); + await failDraining(draining, error); + }; + + return { + failActive, + failDraining, + reclaimActive, + retryFailedTeardown, + teardownEnvironment, + }; +} + +export type PlacementFailureActions = ReturnType; diff --git a/src/gateway/worker-environments/placement-dispatch-recovery.ts b/src/gateway/worker-environments/placement-dispatch-recovery.ts new file mode 100644 index 000000000000..66c36dc974be --- /dev/null +++ b/src/gateway/worker-environments/placement-dispatch-recovery.ts @@ -0,0 +1,231 @@ +import { + isUnavailableEnvironment, + type PlacementFailureActions, + type WorkerActivationBarrier, + type WorkerActiveDispatchPlacement, + type WorkerDispatchEnvironmentService, + type WorkerDispatchPlacement, + type WorkerDispatchPlacementStore, + type WorkerFailedDispatchPlacement, + type WorkerStartingDispatchPlacement, +} from "./placement-dispatch-failure.js"; +import type { WorkerEnvironmentService } from "./service.js"; + +function sameActiveEnvironment( + placement: WorkerActiveDispatchPlacement, + environment: ReturnType, +): environment is NonNullable { + return Boolean( + environment && + environment.state === "attached" && + placement.environmentId && + environment.environmentId === placement.environmentId && + placement.activeOwnerEpoch !== null && + environment.ownerEpoch === placement.activeOwnerEpoch && + placement.workerBundleHash && + environment.bootstrapReceipt?.bundleHash === placement.workerBundleHash && + environment.attachedSessionIds.length === 1 && + environment.attachedSessionIds[0] === placement.sessionId, + ); +} + +function isStartingPlacement( + placement: WorkerDispatchPlacement, +): placement is WorkerStartingDispatchPlacement { + return placement.state === "starting"; +} + +function isFailedPlacement( + placement: WorkerDispatchPlacement, +): placement is WorkerFailedDispatchPlacement { + return placement.state === "failed"; +} + +export function createPlacementRecoveryActions(deps: { + placements: WorkerDispatchPlacementStore; + environments: WorkerDispatchEnvironmentService; + runActivationBarrier: WorkerActivationBarrier; + failure: PlacementFailureActions; +}) { + const { environments, failure, placements } = deps; + + const adoptActive = async (placement: WorkerActiveDispatchPlacement): Promise => { + // Worker turns are one-shot SSH children owned by the previous gateway process. A durable + // claim cannot prove that child remains live after restart, so fence the whole placement. + if (placement.turnClaim) { + const error = new Error( + "Active worker turn claim cannot be proven live after gateway restart", + ); + await failure.failActive(placement, error); + return; + } + const environment = placement.environmentId + ? environments.get(placement.environmentId) + : undefined; + if (!environment || isUnavailableEnvironment(environment)) { + await failure.reclaimActive( + placement, + environment, + new Error("Active worker disappeared during restart reconciliation"), + ); + return; + } + if (!sameActiveEnvironment(placement, environment)) { + await failure.reclaimActive( + placement, + environment, + new Error("Active worker placement does not match its environment owner"), + ); + return; + } + try { + await environments.startTunnel({ + environmentId: environment.environmentId, + ownerEpoch: environment.ownerEpoch, + }); + placements.adoptActive({ + sessionId: placement.sessionId, + expectedGeneration: placement.generation, + environmentId: environment.environmentId, + ownerEpoch: environment.ownerEpoch, + }); + } catch (error) { + await failure.failActive(placement, error); + } + }; + + const resumeStarting = async (placement: WorkerStartingDispatchPlacement): Promise => { + const environment = placement.environmentId + ? environments.get(placement.environmentId) + : undefined; + const expectedBundle = placement.workerBundleHash; + const hasSyncedWorkspace = Boolean( + placement.workspaceBaseManifestRef && placement.remoteWorkspaceDir, + ); + const canResume = + environment && + expectedBundle && + environment.bootstrapReceipt?.bundleHash === expectedBundle && + hasSyncedWorkspace; + if (!canResume) { + const error = new Error("Interrupted worker dispatch cannot safely resume"); + await failure.teardownEnvironment({ + placement, + environmentId: placement.environmentId, + ownerEpoch: environment?.ownerEpoch ?? null, + primaryError: error, + }); + return; + } + try { + const ownerEpoch = + environment.state === "attached" && + environment.attachedSessionIds.length === 1 && + environment.attachedSessionIds[0] === placement.sessionId + ? environment.ownerEpoch + : environment.state === "ready" || environment.state === "idle" + ? ( + await environments.attachSession({ + environmentId: environment.environmentId, + ownerEpoch: environment.ownerEpoch, + sessionId: placement.sessionId, + }) + ).ownerEpoch + : undefined; + if (ownerEpoch === undefined) { + throw new Error(`Worker environment cannot resume dispatch from ${environment.state}`); + } + await environments.startTunnel({ environmentId: environment.environmentId, ownerEpoch }); + await deps.runActivationBarrier({ + sessionId: placement.sessionId, + sessionKey: placement.sessionKey, + agentId: placement.agentId, + activate: () => { + const activated = placements.transition({ + sessionId: placement.sessionId, + from: "starting", + to: "active", + expectedGeneration: placement.generation, + patch: { activeOwnerEpoch: ownerEpoch }, + }); + if (activated.state !== "active") { + throw new Error("Worker dispatch activation did not produce an active placement"); + } + return activated; + }, + }); + } catch (error) { + await failure.teardownEnvironment({ + placement, + environmentId: environment.environmentId, + ownerEpoch: environment.ownerEpoch, + primaryError: error, + }); + } + }; + + const reconcile = async (): Promise => { + await environments.reconcileOnce(); + for (const placement of placements.listForReconcile()) { + if (placement.state === "local" || placement.state === "reclaimed") { + continue; + } + if (placement.state === "active") { + await adoptActive(placement); + continue; + } + if (isFailedPlacement(placement)) { + await failure.retryFailedTeardown(placement); + continue; + } + if (isStartingPlacement(placement)) { + await resumeStarting(placement); + continue; + } + const error = new Error(`Worker dispatch interrupted in ${placement.state}`); + if (placement.state === "draining") { + await failure.failDraining(placement, error); + continue; + } + await failure.teardownEnvironment({ + placement, + environmentId: placement.environmentId, + ownerEpoch: placement.activeOwnerEpoch, + primaryError: error, + }); + } + }; + + // Runtime sweeps must not classify a live dispatch preparation as a crash. They only repair + // durable active ownership and retry teardown already fenced by a previous failure. + const reconcileActive = async (): Promise => { + await environments.reconcileOnce(); + for (const placement of placements.listForReconcile()) { + if (isFailedPlacement(placement)) { + await failure.retryFailedTeardown(placement); + continue; + } + if (placement.state !== "active") { + continue; + } + const environment = environments.get(placement.environmentId); + if (!environment || isUnavailableEnvironment(environment)) { + await failure.reclaimActive( + placement, + environment, + new Error("Active worker disappeared during an admitted turn"), + ); + continue; + } + if (!sameActiveEnvironment(placement, environment)) { + await failure.reclaimActive( + placement, + environment, + new Error("Active worker placement does not match its environment owner"), + ); + } + } + }; + + return { reconcile, reconcileActive }; +} diff --git a/src/gateway/worker-environments/placement-dispatch.test.ts b/src/gateway/worker-environments/placement-dispatch.test.ts new file mode 100644 index 000000000000..3e67e99217dd --- /dev/null +++ b/src/gateway/worker-environments/placement-dispatch.test.ts @@ -0,0 +1,694 @@ +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { + closeOpenClawStateDatabaseForTest, + openOpenClawStateDatabase, + type OpenClawStateDatabase, +} from "../../state/openclaw-state-db.js"; +import type { MintedWorkerCredential } from "./credential.js"; +import type { + WorkerDispatchEnvironmentService, + WorkerDispatchPlacementStore, +} from "./placement-dispatch-failure.js"; +import { createWorkerPlacementDispatchService } from "./placement-dispatch.js"; +import { + createWorkerSessionPlacementStore, + type WorkerSessionPlacementRecord, +} from "./placement-store.js"; +import { workerEnvironmentIdForIdempotencyKey } from "./service.js"; +import type { + WorkerEnvironmentBootstrapReceipt, + WorkerEnvironmentProfileSnapshot, + WorkerEnvironmentSshEndpoint, +} from "./store.js"; +import type { WorkerTunnelHandle } from "./tunnel.js"; + +type WorkerDispatchRequest = Parameters< + ReturnType["dispatch"] +>[0]; + +const BUNDLE_HASH = "a".repeat(64); +const MANIFEST_REF = `sha256:${"b".repeat(64)}`; +const HOST_KEY = [["ssh", "ed25519"].join("-"), "AAAA"].join(" "); +const REQUEST: WorkerDispatchRequest = { + sessionId: "session-1", + sessionKey: "agent:main:session-1", + agentId: "main", + profileId: "development", +}; + +type PlacementStore = ReturnType; +type DispatchEnvironmentRecord = Awaited>; +type DispatchStage = + | "barrier" + | "workspace" + | "create" + | "tunnel:ready" + | "sync" + | "attach" + | "tunnel:attached" + | "activation"; + +function seedStartingPlacement( + store: PlacementStore, + environmentId: string, +): WorkerSessionPlacementRecord { + let current = store.startDispatch(REQUEST); + current = store.transition({ + sessionId: REQUEST.sessionId, + from: "requested", + to: "provisioning", + expectedGeneration: current.generation, + patch: { environmentId }, + }); + current = store.transition({ + sessionId: REQUEST.sessionId, + from: "provisioning", + to: "syncing", + expectedGeneration: current.generation, + patch: { workerBundleHash: BUNDLE_HASH }, + }); + current = store.transition({ + sessionId: REQUEST.sessionId, + from: "syncing", + to: "starting", + expectedGeneration: current.generation, + patch: { + workspaceBaseManifestRef: MANIFEST_REF, + remoteWorkspaceDir: "/worker/workspace", + }, + }); + return current; +} + +function seedActivePlacement( + store: PlacementStore, + params: { environmentId: string; ownerEpoch: number }, +): WorkerSessionPlacementRecord { + const current = seedStartingPlacement(store, params.environmentId); + return store.transition({ + sessionId: REQUEST.sessionId, + from: "starting", + to: "active", + expectedGeneration: current.generation, + patch: { activeOwnerEpoch: params.ownerEpoch }, + }); +} + +function createHarness( + placementStore: PlacementStore, + options: { failAt?: DispatchStage; destroyFails?: boolean; claimOnDrain?: boolean } = {}, +) { + const log: string[] = []; + const fail = (stage: DispatchStage) => { + log.push(stage); + if (options.failAt === stage) { + throw new Error(`${stage} failed`); + } + }; + const placements: WorkerDispatchPlacementStore = { + get: (sessionId) => placementStore.get(sessionId), + startDispatch: (params) => { + log.push("placement:requested"); + return placementStore.startDispatch(params); + }, + transition: (params) => { + log.push(`placement:${params.to}`); + return placementStore.transition(params); + }, + fail: (params) => { + log.push("placement:failed"); + return placementStore.fail(params); + }, + listForReconcile: () => placementStore.listForReconcile(), + startDrain: (params) => { + log.push("placement:draining"); + if (options.claimOnDrain) { + placementStore.claimTurn({ + sessionId: params.sessionId, + sessionKey: REQUEST.sessionKey, + agentId: REQUEST.agentId, + claimId: "claim-on-drain", + runId: "run-on-drain", + owner: { + kind: "worker", + environmentId: params.environmentId, + ownerEpoch: params.ownerEpoch, + }, + }); + } + return placementStore.startDrain(params); + }, + startReconcile: (params) => { + log.push("placement:reconciling"); + return placementStore.startReconcile(params); + }, + adoptActive: (params) => { + log.push("placement:adopted"); + return placementStore.adoptActive(params); + }, + }; + const environmentId = workerEnvironmentIdForIdempotencyKey( + `session-dispatch:${REQUEST.sessionId}:1`, + ); + const profileSnapshot: WorkerEnvironmentProfileSnapshot = { + settings: { region: "test" }, + }; + const bootstrapReceipt: WorkerEnvironmentBootstrapReceipt = { + bundleHash: BUNDLE_HASH, + openclawVersion: "2026.7.2", + protocolFeatures: [], + }; + const sshEndpoint: WorkerEnvironmentSshEndpoint = { + host: "worker.example.test", + port: 22, + user: "worker", + hostKey: HOST_KEY, + keyRef: { source: "file", provider: "worker-keys", id: "/key" }, + }; + const environmentBase = { + environmentId, + providerId: "fake", + profileId: "development", + profileSnapshot, + provisionOperationId: "provision-1", + bootstrapReceipt, + teardownTerminalState: null, + lastError: null, + createdAtMs: 1, + updatedAtMs: 1, + stateChangedAtMs: 1, + idleSinceAtMs: null, + destroyRequestedAtMs: null, + leaseId: "lease-1", + sshEndpoint, + }; + const ready = { + ...environmentBase, + state: "ready", + ownerEpoch: 1, + attachedSessionIds: [], + tunnelStatus: "connected", + } satisfies DispatchEnvironmentRecord; + const attached = { + ...environmentBase, + state: "attached", + ownerEpoch: 2, + attachedSessionIds: [REQUEST.sessionId], + tunnelStatus: "connected", + } satisfies DispatchEnvironmentRecord; + let currentEnvironment: ReturnType = ready; + const destroyedEnvironment = (ownerEpoch: number): DispatchEnvironmentRecord => ({ + ...environmentBase, + state: "destroyed", + ownerEpoch, + attachedSessionIds: [], + tunnelStatus: "stopped", + }); + const tunnelHandle = (ownerEpoch: number): WorkerTunnelHandle => ({ + environmentId: ready.environmentId, + ownerEpoch, + remoteSocketPath: "/worker/gateway.sock", + runWorkspaceCommand: vi.fn(async () => ({ + stdout: "", + stderr: "", + code: 0, + signal: null, + killed: false, + termination: "exit" as const, + })), + syncWorkspace: vi.fn(async () => { + fail("sync"); + return { + mode: "git" as const, + remoteWorkspaceDir: "/worker/workspace", + manifestRef: MANIFEST_REF, + }; + }), + stop: vi.fn(async () => {}), + }); + const credential: MintedWorkerCredential = { + credential: ["worker", "credential", "fixture"].join("-"), + deliveryId: "c".repeat(43), + environmentId: ready.environmentId, + bundleHash: BUNDLE_HASH, + sessionId: REQUEST.sessionId, + rpcSetVersion: 1, + ownerEpoch: 2, + expiresAtMs: 10_000, + }; + const environments: WorkerDispatchEnvironmentService = { + create: vi.fn(async () => { + fail("create"); + return ready; + }), + get: vi.fn(() => currentEnvironment), + attachSession: vi.fn(async () => { + fail("attach"); + currentEnvironment = attached; + return credential; + }), + startTunnel: vi.fn(async ({ ownerEpoch }) => { + fail(ownerEpoch === 1 ? "tunnel:ready" : "tunnel:attached"); + return tunnelHandle(ownerEpoch); + }), + stopTunnel: vi.fn(async () => { + log.push("teardown:stop"); + }), + destroy: vi.fn(async () => { + log.push("teardown:destroy"); + if (options.destroyFails) { + throw new Error("destroy pending"); + } + const destroyed = destroyedEnvironment((currentEnvironment?.ownerEpoch ?? 1) + 1); + currentEnvironment = destroyed; + return destroyed; + }), + reconcileOnce: vi.fn(async () => { + log.push("environment:reconcile"); + }), + }; + const service = createWorkerPlacementDispatchService({ + placements, + environments, + runLocalBarrier: async ({ startDispatch }) => { + log.push("barrier"); + const placement = startDispatch(); + if (options.failAt === "barrier") { + throw new Error("barrier failed"); + } + return placement; + }, + runActivationBarrier: async ({ activate }) => { + fail("activation"); + return activate(); + }, + resolveWorkspacePath: async () => { + fail("workspace"); + return "/gateway/workspace"; + }, + }); + return { + log, + placements: { + current: () => placementStore.get(REQUEST.sessionId), + seedStarting: () => seedStartingPlacement(placementStore, environmentId), + seedActive: (ownerEpoch: number) => + seedActivePlacement(placementStore, { environmentId, ownerEpoch }), + seedDraining: (ownerEpoch: number) => { + const active = seedActivePlacement(placementStore, { environmentId, ownerEpoch }); + if (active.state !== "active") { + throw new Error("active placement fixture was not active"); + } + return placementStore.startDrain({ + sessionId: active.sessionId, + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + expectedGeneration: active.generation, + }); + }, + }, + environments, + markEnvironmentDestroyed: () => { + currentEnvironment = destroyedEnvironment((currentEnvironment?.ownerEpoch ?? 1) + 1); + }, + markEnvironmentOwnerEpoch: (ownerEpoch: number) => { + currentEnvironment = { ...attached, ownerEpoch }; + }, + service, + ready, + attached, + }; +} + +describe("worker placement dispatch", () => { + let root: string; + let database: OpenClawStateDatabase; + let placementStore: PlacementStore; + + beforeEach(async () => { + root = await fs.mkdtemp(path.join(await fs.realpath(os.tmpdir()), "openclaw-dispatch-")); + database = openOpenClawStateDatabase({ env: { OPENCLAW_STATE_DIR: root } }); + placementStore = createWorkerSessionPlacementStore({ database, now: () => 1_000 }); + }); + + afterEach(async () => { + closeOpenClawStateDatabaseForTest(); + await fs.rm(root, { recursive: true, force: true }); + }); + + it("orders the migration barrier, provisioning, sync, attachment, and activation", async () => { + const harness = createHarness(placementStore); + + await expect(harness.service.dispatch(REQUEST)).resolves.toMatchObject({ + state: "active", + environmentId: harness.ready.environmentId, + activeOwnerEpoch: 2, + workspaceBaseManifestRef: MANIFEST_REF, + remoteWorkspaceDir: "/worker/workspace", + workerBundleHash: BUNDLE_HASH, + }); + + expect(harness.log).toEqual([ + "barrier", + "placement:requested", + "workspace", + "placement:provisioning", + "create", + "placement:syncing", + "tunnel:ready", + "sync", + "placement:starting", + "attach", + "tunnel:attached", + "activation", + "placement:active", + ]); + }); + + it.each([ + "barrier", + "workspace", + "create", + "tunnel:ready", + "sync", + "attach", + "tunnel:attached", + "activation", + ])("fails closed and tears down acquired resources when %s fails", async (failAt) => { + const harness = createHarness(placementStore, { failAt }); + + await expect(harness.service.dispatch(REQUEST)).rejects.toThrow(`${failAt} failed`); + + expect(harness.placements.current()).toMatchObject({ + state: "failed", + recoveryError: `${failAt} failed`, + }); + const failedAt = harness.log.indexOf("placement:failed"); + expect(failedAt).toBeGreaterThan(-1); + const environmentAcquired = !["barrier", "workspace"].includes(failAt); + expect(harness.log.includes("teardown:stop")).toBe(environmentAcquired); + expect(harness.log.includes("teardown:destroy")).toBe(environmentAcquired); + if (environmentAcquired) { + expect(failedAt).toBeGreaterThan(harness.log.indexOf("teardown:destroy")); + } + }); + + it("does not fail or tear down a dispatch owned by another invocation", async () => { + placementStore.startDispatch(REQUEST); + const harness = createHarness(placementStore); + + await expect(harness.service.dispatch(REQUEST)).rejects.toThrow( + "Cannot dispatch session session-1 from placement requested", + ); + + expect(harness.placements.current()).toMatchObject({ state: "requested" }); + expect(harness.log).not.toContain("placement:failed"); + expect(harness.log).not.toContain("teardown:destroy"); + }); + + it("persists pending teardown evidence after placement is fenced", async () => { + const harness = createHarness(placementStore, { failAt: "sync", destroyFails: true }); + + await expect(harness.service.dispatch(REQUEST)).rejects.toThrow("sync failed"); + + expect(harness.placements.current()).toMatchObject({ + state: "failed", + recoveryError: expect.stringContaining("environment destroy: destroy pending"), + }); + expect(harness.log.filter((entry) => entry === "placement:failed")).toHaveLength(1); + }); + + it("adopts an exact active environment after restart without reprovisioning", async () => { + const harness = createHarness(placementStore); + await harness.environments.attachSession({ + environmentId: harness.ready.environmentId, + ownerEpoch: harness.ready.ownerEpoch, + sessionId: REQUEST.sessionId, + }); + harness.placements.seedActive(harness.attached.ownerEpoch); + harness.log.length = 0; + + await harness.service.reconcile(); + + expect(harness.log).toEqual(["environment:reconcile", "tunnel:attached", "placement:adopted"]); + expect(harness.environments.create).not.toHaveBeenCalled(); + expect(harness.environments.destroy).not.toHaveBeenCalled(); + }); + + it("reclaims an active placement whose environment is already terminal after restart", async () => { + const harness = createHarness(placementStore); + harness.placements.seedActive(harness.attached.ownerEpoch); + harness.markEnvironmentDestroyed(); + harness.log.length = 0; + + await harness.service.reconcile(); + + expect(harness.placements.current()).toMatchObject({ + state: "reclaimed", + environmentId: harness.ready.environmentId, + activeOwnerEpoch: harness.attached.ownerEpoch, + }); + expect(harness.log).toEqual([ + "environment:reconcile", + "placement:draining", + "placement:reconciling", + "placement:reclaimed", + ]); + expect(harness.environments.startTunnel).not.toHaveBeenCalled(); + expect(harness.environments.destroy).not.toHaveBeenCalled(); + }); + + it("fails closed when an active worker turn claim cannot be proven live after restart", async () => { + const harness = createHarness(placementStore); + await harness.environments.attachSession({ + environmentId: harness.ready.environmentId, + ownerEpoch: harness.ready.ownerEpoch, + sessionId: REQUEST.sessionId, + }); + harness.placements.seedActive(harness.attached.ownerEpoch); + placementStore.claimTurn({ + ...REQUEST, + claimId: "claim-1", + runId: "run-1", + owner: { + kind: "worker", + environmentId: harness.attached.environmentId, + ownerEpoch: harness.attached.ownerEpoch, + }, + }); + harness.log.length = 0; + + await harness.service.reconcile(); + + expect(harness.placements.current()).toMatchObject({ + state: "failed", + turnClaim: null, + recoveryError: "Active worker turn claim cannot be proven live after gateway restart", + }); + expect(harness.log).toEqual([ + "environment:reconcile", + "placement:draining", + "teardown:stop", + "teardown:destroy", + "placement:reconciling", + "placement:failed", + ]); + expect(harness.environments.startTunnel).not.toHaveBeenCalled(); + }); + + it("resumes a synced starting placement after restart", async () => { + const harness = createHarness(placementStore); + harness.placements.seedStarting(); + harness.log.length = 0; + + await harness.service.reconcile(); + + expect(harness.placements.current()).toMatchObject({ + state: "active", + environmentId: harness.ready.environmentId, + activeOwnerEpoch: harness.attached.ownerEpoch, + }); + expect(harness.log).toEqual([ + "environment:reconcile", + "attach", + "tunnel:attached", + "activation", + "placement:active", + ]); + expect(harness.environments.create).not.toHaveBeenCalled(); + }); + + it("finishes an interrupted drain through reconciliation before failure", async () => { + const harness = createHarness(placementStore); + harness.placements.seedDraining(harness.attached.ownerEpoch); + harness.log.length = 0; + + await harness.service.reconcile(); + + expect(harness.placements.current()).toMatchObject({ + state: "failed", + turnClaim: null, + recoveryError: "Worker dispatch interrupted in draining", + }); + expect(harness.log).toEqual([ + "environment:reconcile", + "teardown:stop", + "teardown:destroy", + "placement:reconciling", + "placement:failed", + ]); + }); + + it("drains, tears down, and reclaims an idle active placement with a mismatched owner", async () => { + const harness = createHarness(placementStore); + harness.placements.seedActive(99); + + await harness.service.reconcile(); + + expect(harness.placements.current()).toMatchObject({ + state: "reclaimed", + }); + expect(harness.log).toEqual([ + "environment:reconcile", + "placement:draining", + "teardown:stop", + "teardown:destroy", + "placement:reconciling", + "placement:reclaimed", + ]); + + const destroyCalls = vi.mocked(harness.environments.destroy).mock.calls.length; + await harness.service.reconcile(); + expect(harness.environments.destroy).toHaveBeenCalledTimes(destroyCalls); + }); + + it("preserves a live active turn claim during runtime reconciliation", async () => { + const harness = createHarness(placementStore); + await harness.environments.attachSession({ + environmentId: harness.ready.environmentId, + ownerEpoch: harness.ready.ownerEpoch, + sessionId: REQUEST.sessionId, + }); + harness.placements.seedActive(harness.attached.ownerEpoch); + placementStore.claimTurn({ + ...REQUEST, + claimId: "claim-1", + runId: "run-1", + owner: { + kind: "worker", + environmentId: harness.attached.environmentId, + ownerEpoch: harness.attached.ownerEpoch, + }, + }); + harness.log.length = 0; + + await harness.service.reconcileActive(); + + expect(harness.placements.current()).toMatchObject({ + state: "active", + turnClaim: { + claimId: "claim-1", + runId: "run-1", + owner: "worker", + }, + }); + expect(harness.log).toEqual(["environment:reconcile"]); + expect(harness.environments.startTunnel).not.toHaveBeenCalled(); + expect(harness.environments.destroy).not.toHaveBeenCalled(); + }); + + it("fences a live turn before tearing down a mismatched runtime owner", async () => { + const harness = createHarness(placementStore); + await harness.environments.attachSession({ + environmentId: harness.ready.environmentId, + ownerEpoch: harness.ready.ownerEpoch, + sessionId: REQUEST.sessionId, + }); + harness.placements.seedActive(harness.attached.ownerEpoch); + placementStore.claimTurn({ + ...REQUEST, + claimId: "claim-1", + runId: "run-1", + owner: { + kind: "worker", + environmentId: harness.attached.environmentId, + ownerEpoch: harness.attached.ownerEpoch, + }, + }); + harness.markEnvironmentOwnerEpoch(harness.attached.ownerEpoch + 1); + harness.log.length = 0; + + await harness.service.reconcileActive(); + + expect(harness.placements.current()).toMatchObject({ + state: "failed", + turnClaim: null, + recoveryError: "Active worker placement does not match its environment owner", + }); + expect(harness.log).toEqual([ + "environment:reconcile", + "placement:draining", + "teardown:stop", + "teardown:destroy", + "placement:reconciling", + "placement:failed", + ]); + }); + + it("reclaims a terminal active environment during runtime reconciliation", async () => { + const harness = createHarness(placementStore); + harness.placements.seedActive(harness.attached.ownerEpoch); + harness.markEnvironmentDestroyed(); + harness.log.length = 0; + + await harness.service.reconcileActive(); + + expect(harness.placements.current()).toMatchObject({ state: "reclaimed" }); + expect(harness.log).toEqual([ + "environment:reconcile", + "placement:draining", + "placement:reconciling", + "placement:reclaimed", + ]); + expect(harness.environments.destroy).not.toHaveBeenCalled(); + }); + + it("fences a turn admitted immediately before runtime drain", async () => { + const harness = createHarness(placementStore, { claimOnDrain: true }); + harness.placements.seedActive(harness.attached.ownerEpoch); + harness.markEnvironmentDestroyed(); + harness.log.length = 0; + + await harness.service.reconcileActive(); + + expect(harness.placements.current()).toMatchObject({ + state: "failed", + turnClaim: null, + recoveryError: "Active worker disappeared during an admitted turn", + }); + expect(harness.log).toEqual([ + "environment:reconcile", + "placement:draining", + "teardown:stop", + "teardown:destroy", + "placement:reconciling", + "placement:failed", + ]); + }); + + it("leaves in-flight dispatch preparation untouched during runtime reconciliation", async () => { + const harness = createHarness(placementStore); + harness.placements.seedStarting(); + harness.log.length = 0; + + await harness.service.reconcileActive(); + + expect(harness.placements.current()).toMatchObject({ state: "starting" }); + expect(harness.log).toEqual(["environment:reconcile"]); + expect(harness.environments.attachSession).not.toHaveBeenCalled(); + expect(harness.environments.destroy).not.toHaveBeenCalled(); + }); +}); diff --git a/src/gateway/worker-environments/placement-dispatch.ts b/src/gateway/worker-environments/placement-dispatch.ts new file mode 100644 index 000000000000..7b4c7ac17209 --- /dev/null +++ b/src/gateway/worker-environments/placement-dispatch.ts @@ -0,0 +1,178 @@ +import { + createPlacementFailureActions, + type WorkerActivationBarrier, + type WorkerActiveDispatchPlacement, + type WorkerDispatchEnvironmentService, + type WorkerDispatchPlacement, + type WorkerDispatchPlacementStore, +} from "./placement-dispatch-failure.js"; +import { createPlacementRecoveryActions } from "./placement-dispatch-recovery.js"; +import type { WorkerPlacementDispatchRequest } from "./service-contract.js"; +import { type WorkerEnvironmentService, workerEnvironmentIdForIdempotencyKey } from "./service.js"; + +type WorkerLocalDispatchBarrier = (params: { + sessionId: string; + sessionKey: string; + agentId: string; + startDispatch: () => WorkerDispatchPlacement; +}) => Promise; + +type WorkerPlacementDispatchOptions = { + placements: WorkerDispatchPlacementStore; + environments: WorkerDispatchEnvironmentService; + runLocalBarrier: WorkerLocalDispatchBarrier; + runActivationBarrier: WorkerActivationBarrier; + resolveWorkspacePath: (params: { + sessionId: string; + sessionKey: string; + agentId: string; + }) => Promise; +}; + +function requireProvisionedEnvironment( + environment: Awaited>, + expectedEnvironmentId: string, +): { environmentId: string; ownerEpoch: number; bundleHash: string } { + if ( + (environment.state !== "ready" && environment.state !== "idle") || + !environment.bootstrapReceipt || + environment.environmentId !== expectedEnvironmentId + ) { + throw new Error(`Worker environment is not dispatchable: ${environment.state}`); + } + return { + environmentId: environment.environmentId, + ownerEpoch: environment.ownerEpoch, + bundleHash: environment.bootstrapReceipt.bundleHash, + }; +} + +export function createWorkerPlacementDispatchService(options: WorkerPlacementDispatchOptions) { + const { environments, placements } = options; + const failure = createPlacementFailureActions({ environments, placements }); + const recovery = createPlacementRecoveryActions({ + environments, + failure, + placements, + runActivationBarrier: options.runActivationBarrier, + }); + + const dispatch = async ( + request: WorkerPlacementDispatchRequest, + ): Promise => { + let placement: WorkerDispatchPlacement | undefined; + let environmentId: string | null = null; + let ownerEpoch: number | null = null; + try { + placement = await options.runLocalBarrier({ + sessionId: request.sessionId, + sessionKey: request.sessionKey, + agentId: request.agentId, + startDispatch: () => { + placement = placements.startDispatch({ + sessionId: request.sessionId, + sessionKey: request.sessionKey, + agentId: request.agentId, + }); + return placement; + }, + }); + const localPath = await options.resolveWorkspacePath(request); + const idempotencyKey = `session-dispatch:${request.sessionId}:${placement.generation}`; + const expectedEnvironmentId = workerEnvironmentIdForIdempotencyKey(idempotencyKey); + placement = placements.transition({ + sessionId: request.sessionId, + from: "requested", + to: "provisioning", + expectedGeneration: placement.generation, + patch: { environmentId: expectedEnvironmentId }, + }); + const environment = await environments.create(request.profileId, idempotencyKey); + const provisioned = requireProvisionedEnvironment(environment, expectedEnvironmentId); + environmentId = provisioned.environmentId; + ownerEpoch = provisioned.ownerEpoch; + placement = placements.transition({ + sessionId: request.sessionId, + from: "provisioning", + to: "syncing", + expectedGeneration: placement.generation, + patch: { + environmentId, + workerBundleHash: provisioned.bundleHash, + }, + }); + const readyTunnel = await environments.startTunnel({ environmentId, ownerEpoch }); + const synced = await readyTunnel.syncWorkspace({ + localPath, + sessionId: request.sessionId, + generation: placement.generation, + }); + placement = placements.transition({ + sessionId: request.sessionId, + from: "syncing", + to: "starting", + expectedGeneration: placement.generation, + patch: { + workspaceBaseManifestRef: synced.manifestRef, + remoteWorkspaceDir: synced.remoteWorkspaceDir, + }, + }); + const credential = await environments.attachSession({ + environmentId, + ownerEpoch, + sessionId: request.sessionId, + }); + ownerEpoch = credential.ownerEpoch; + await environments.startTunnel({ environmentId, ownerEpoch }); + const startingPlacement = placement; + const activePlacement = await options.runActivationBarrier({ + sessionId: request.sessionId, + sessionKey: request.sessionKey, + agentId: request.agentId, + activate: () => { + const activated = placements.transition({ + sessionId: request.sessionId, + from: "starting", + to: "active", + expectedGeneration: startingPlacement.generation, + patch: { activeOwnerEpoch: ownerEpoch }, + }); + if (activated.state !== "active") { + throw new Error("Worker dispatch activation did not produce an active placement"); + } + return activated; + }, + }); + return activePlacement; + } catch (error) { + const current = placement ? placements.get(request.sessionId) : undefined; + if (current && current.state !== "local" && current.state !== "reclaimed") { + if (current.state === "active") { + await failure.failActive(current, error); + } else { + const currentEnvironmentId = environmentId ?? current.environmentId; + const currentEnvironment = currentEnvironmentId + ? environments.get(currentEnvironmentId) + : undefined; + await failure.teardownEnvironment({ + placement: current, + environmentId: currentEnvironment?.environmentId ?? null, + ownerEpoch: ownerEpoch ?? currentEnvironment?.ownerEpoch ?? null, + primaryError: error, + }); + } + } + throw error; + } + }; + + return { + dispatch, + reconcile: recovery.reconcile, + reconcileActive: recovery.reconcileActive, + }; +} + +export type WorkerPlacementDispatchService = ReturnType< + typeof createWorkerPlacementDispatchService +>; diff --git a/src/gateway/worker-environments/placement-projector.test.ts b/src/gateway/worker-environments/placement-projector.test.ts new file mode 100644 index 000000000000..4b7c52ab46e5 --- /dev/null +++ b/src/gateway/worker-environments/placement-projector.test.ts @@ -0,0 +1,107 @@ +import { Value } from "typebox/value"; +import { describe, expect, it } from "vitest"; +import { SessionPlacementSchema } from "../../../packages/gateway-protocol/src/index.js"; +import { projectWorkerSessionPlacement } from "./placement-projector.js"; +import type { WorkerSessionPlacementRecord } from "./placement-store.js"; + +const BUNDLE_HASH = "a".repeat(64); + +const RECORD_BASE = { + sessionId: "session-1", + agentId: "main", + sessionKey: "agent:main:session-1", + generation: 4, + workspaceBaseManifestRef: null, + remoteWorkspaceDir: null, + workerBundleHash: null, + lastTranscriptAckCursor: null, + lastLiveEventAckCursor: null, + recoveryError: null, + turnClaim: null, + createdAtMs: 100, + updatedAtMs: 200, + stateChangedAtMs: 150, +}; + +describe("worker placement projection", () => { + it("emits only fields valid for each placement discriminator", () => { + const records = [ + { + ...RECORD_BASE, + state: "local", + environmentId: null, + activeOwnerEpoch: null, + }, + { + ...RECORD_BASE, + state: "provisioning", + environmentId: "environment-1", + activeOwnerEpoch: null, + }, + { + ...RECORD_BASE, + state: "reclaimed", + environmentId: "environment-1", + activeOwnerEpoch: 7, + workspaceBaseManifestRef: "manifest-1", + remoteWorkspaceDir: "/workspace", + workerBundleHash: BUNDLE_HASH, + }, + { + ...RECORD_BASE, + state: "failed", + environmentId: "environment-1", + activeOwnerEpoch: 7, + recoveryError: "worker unavailable", + }, + ] satisfies WorkerSessionPlacementRecord[]; + + const projected = records.map(projectWorkerSessionPlacement); + + expect(projected).toEqual([ + { + state: "local", + generation: 4, + createdAtMs: 100, + updatedAtMs: 200, + stateChangedAtMs: 150, + }, + { + state: "provisioning", + generation: 4, + createdAtMs: 100, + updatedAtMs: 200, + stateChangedAtMs: 150, + environmentId: "environment-1", + }, + { + state: "reclaimed", + generation: 4, + createdAtMs: 100, + updatedAtMs: 200, + stateChangedAtMs: 150, + environmentId: "environment-1", + activeOwnerEpoch: 7, + workspaceBaseManifestRef: "manifest-1", + remoteWorkspaceDir: "/workspace", + workerBundleHash: BUNDLE_HASH, + }, + { + state: "failed", + generation: 4, + createdAtMs: 100, + updatedAtMs: 200, + stateChangedAtMs: 150, + environmentId: "environment-1", + activeOwnerEpoch: 7, + recoveryError: "worker unavailable", + }, + ]); + for (const placement of projected) { + expect(Value.Check(SessionPlacementSchema, placement)).toBe(true); + expect(placement).not.toHaveProperty("sessionId"); + expect(placement).not.toHaveProperty("sessionKey"); + expect(placement).not.toHaveProperty("turnClaim"); + } + }); +}); diff --git a/src/gateway/worker-environments/placement-projector.ts b/src/gateway/worker-environments/placement-projector.ts new file mode 100644 index 000000000000..ec831294f5c2 --- /dev/null +++ b/src/gateway/worker-environments/placement-projector.ts @@ -0,0 +1,133 @@ +import type { SessionPlacement } from "../../../packages/gateway-protocol/src/index.js"; +import type { WorkerSessionPlacementRecord } from "./placement-store.js"; + +export type WorkerSessionPlacementReader = { + getMany(sessionIds: readonly string[]): ReadonlyMap; +}; + +/** Removes gateway-only identity and turn-claim fields from the operator projection. */ +export function projectWorkerSessionPlacement( + record: WorkerSessionPlacementRecord, +): SessionPlacement { + const timing = { + generation: record.generation, + createdAtMs: record.createdAtMs, + updatedAtMs: record.updatedAtMs, + stateChangedAtMs: record.stateChangedAtMs, + }; + switch (record.state) { + case "local": + return { state: "local", ...timing }; + case "requested": + return { state: "requested", ...timing }; + case "provisioning": + return { + state: "provisioning", + ...timing, + ...(record.environmentId ? { environmentId: record.environmentId } : {}), + }; + case "syncing": + return { + state: "syncing", + ...timing, + environmentId: record.environmentId, + workerBundleHash: record.workerBundleHash, + }; + case "starting": + return { + state: "starting", + ...timing, + environmentId: record.environmentId, + workerBundleHash: record.workerBundleHash, + workspaceBaseManifestRef: record.workspaceBaseManifestRef, + remoteWorkspaceDir: record.remoteWorkspaceDir, + }; + case "active": + return { + state: "active", + ...timing, + environmentId: record.environmentId, + activeOwnerEpoch: record.activeOwnerEpoch, + workerBundleHash: record.workerBundleHash, + workspaceBaseManifestRef: record.workspaceBaseManifestRef, + remoteWorkspaceDir: record.remoteWorkspaceDir, + ...(record.lastTranscriptAckCursor !== null + ? { lastTranscriptAckCursor: record.lastTranscriptAckCursor } + : {}), + ...(record.lastLiveEventAckCursor !== null + ? { lastLiveEventAckCursor: record.lastLiveEventAckCursor } + : {}), + }; + case "draining": + return { + state: "draining", + ...timing, + environmentId: record.environmentId, + activeOwnerEpoch: record.activeOwnerEpoch, + workerBundleHash: record.workerBundleHash, + workspaceBaseManifestRef: record.workspaceBaseManifestRef, + remoteWorkspaceDir: record.remoteWorkspaceDir, + ...(record.lastTranscriptAckCursor !== null + ? { lastTranscriptAckCursor: record.lastTranscriptAckCursor } + : {}), + ...(record.lastLiveEventAckCursor !== null + ? { lastLiveEventAckCursor: record.lastLiveEventAckCursor } + : {}), + }; + case "reconciling": + return { + state: "reconciling", + ...timing, + environmentId: record.environmentId, + activeOwnerEpoch: record.activeOwnerEpoch, + workerBundleHash: record.workerBundleHash, + workspaceBaseManifestRef: record.workspaceBaseManifestRef, + remoteWorkspaceDir: record.remoteWorkspaceDir, + ...(record.lastTranscriptAckCursor !== null + ? { lastTranscriptAckCursor: record.lastTranscriptAckCursor } + : {}), + ...(record.lastLiveEventAckCursor !== null + ? { lastLiveEventAckCursor: record.lastLiveEventAckCursor } + : {}), + }; + case "reclaimed": + return { + state: "reclaimed", + ...timing, + ...(record.environmentId ? { environmentId: record.environmentId } : {}), + ...(record.activeOwnerEpoch !== null ? { activeOwnerEpoch: record.activeOwnerEpoch } : {}), + ...(record.workspaceBaseManifestRef + ? { workspaceBaseManifestRef: record.workspaceBaseManifestRef } + : {}), + ...(record.remoteWorkspaceDir ? { remoteWorkspaceDir: record.remoteWorkspaceDir } : {}), + ...(record.workerBundleHash ? { workerBundleHash: record.workerBundleHash } : {}), + ...(record.lastTranscriptAckCursor !== null + ? { lastTranscriptAckCursor: record.lastTranscriptAckCursor } + : {}), + ...(record.lastLiveEventAckCursor !== null + ? { lastLiveEventAckCursor: record.lastLiveEventAckCursor } + : {}), + }; + case "failed": + return { + state: "failed", + ...timing, + ...(record.environmentId ? { environmentId: record.environmentId } : {}), + ...(record.activeOwnerEpoch !== null ? { activeOwnerEpoch: record.activeOwnerEpoch } : {}), + ...(record.workspaceBaseManifestRef + ? { workspaceBaseManifestRef: record.workspaceBaseManifestRef } + : {}), + ...(record.remoteWorkspaceDir ? { remoteWorkspaceDir: record.remoteWorkspaceDir } : {}), + ...(record.workerBundleHash ? { workerBundleHash: record.workerBundleHash } : {}), + ...(record.lastTranscriptAckCursor !== null + ? { lastTranscriptAckCursor: record.lastTranscriptAckCursor } + : {}), + ...(record.lastLiveEventAckCursor !== null + ? { lastLiveEventAckCursor: record.lastLiveEventAckCursor } + : {}), + recoveryError: record.recoveryError, + }; + } + // Exhaustive over placement states; the return satisfies consistent-return. + return record satisfies never; +} diff --git a/src/gateway/worker-environments/placement-record.ts b/src/gateway/worker-environments/placement-record.ts new file mode 100644 index 000000000000..de7326a92480 --- /dev/null +++ b/src/gateway/worker-environments/placement-record.ts @@ -0,0 +1,371 @@ +import type { WorkerSessionPlacementState } from "./placement-state.js"; + +export type WorkerSessionPlacementIdentity = { + sessionId: string; + agentId: string; + sessionKey: string; +}; + +export type WorkerSessionTurnOwner = + | { kind: "local" } + | { kind: "worker"; environmentId: string; ownerEpoch: number }; + +export type WorkerSessionTurnClaim = { + sessionId: string; + claimId: string; + runId: string; + placementGeneration: number; + owner: WorkerSessionTurnOwner; +}; + +export type PersistedTurnClaim = + | { + owner: "local"; + claimId: string; + runId: string; + generation: number; + ownerEpoch: null; + } + | { + owner: "worker"; + claimId: string; + runId: string; + generation: number; + ownerEpoch: number; + }; + +type PersistedLocalTurnClaim = Extract; +type PersistedWorkerTurnClaim = Extract; + +type PlacementRecordBase = + WorkerSessionPlacementIdentity & { + generation: number; + turnClaim: TurnClaim; + createdAtMs: number; + updatedAtMs: number; + stateChangedAtMs: number; + }; + +type UnclaimedPlacementRecordBase = PlacementRecordBase; +type LocalClaimablePlacementRecordBase = PlacementRecordBase; +type WorkerClaimablePlacementRecordBase = PlacementRecordBase; + +export type EmptyWorkerPlacementMetadata = { + environmentId: null; + activeOwnerEpoch: null; + workspaceBaseManifestRef: null; + remoteWorkspaceDir: null; + workerBundleHash: null; + lastTranscriptAckCursor: null; + lastLiveEventAckCursor: null; + recoveryError: null; +}; + +type ProvisioningPlacementMetadata = { + environmentId: string | null; + activeOwnerEpoch: null; + workspaceBaseManifestRef: null; + remoteWorkspaceDir: null; + workerBundleHash: null; + lastTranscriptAckCursor: null; + lastLiveEventAckCursor: null; + recoveryError: null; +}; + +type SyncingPlacementMetadata = { + environmentId: string; + activeOwnerEpoch: null; + workspaceBaseManifestRef: null; + remoteWorkspaceDir: null; + workerBundleHash: string; + lastTranscriptAckCursor: null; + lastLiveEventAckCursor: null; + recoveryError: null; +}; + +type StartingPlacementMetadata = { + environmentId: string; + activeOwnerEpoch: null; + workspaceBaseManifestRef: string; + remoteWorkspaceDir: string; + workerBundleHash: string; + lastTranscriptAckCursor: null; + lastLiveEventAckCursor: null; + recoveryError: null; +}; + +export type OwnedWorkerPlacementMetadata = { + environmentId: string; + activeOwnerEpoch: number; + workspaceBaseManifestRef: string; + remoteWorkspaceDir: string; + workerBundleHash: string; + lastTranscriptAckCursor: number | null; + lastLiveEventAckCursor: number | null; + recoveryError: null; +}; + +type TerminalPlacementMetadata = { + environmentId: string | null; + activeOwnerEpoch: number | null; + workspaceBaseManifestRef: string | null; + remoteWorkspaceDir: string | null; + workerBundleHash: string | null; + lastTranscriptAckCursor: number | null; + lastLiveEventAckCursor: number | null; +}; + +type LocalPlacementRecord = LocalClaimablePlacementRecordBase & + EmptyWorkerPlacementMetadata & { + state: "local"; + }; +type RequestedPlacementRecord = LocalClaimablePlacementRecordBase & + EmptyWorkerPlacementMetadata & { + state: "requested"; + }; +type ProvisioningPlacementRecord = UnclaimedPlacementRecordBase & + ProvisioningPlacementMetadata & { + state: "provisioning"; + }; +type SyncingPlacementRecord = UnclaimedPlacementRecordBase & + SyncingPlacementMetadata & { + state: "syncing"; + }; +type StartingPlacementRecord = UnclaimedPlacementRecordBase & + StartingPlacementMetadata & { + state: "starting"; + }; +type ActivePlacementRecord = WorkerClaimablePlacementRecordBase & + OwnedWorkerPlacementMetadata & { + state: "active"; + }; +type DrainingPlacementRecord = WorkerClaimablePlacementRecordBase & + OwnedWorkerPlacementMetadata & { + state: "draining"; + }; +type ReconcilingPlacementRecord = UnclaimedPlacementRecordBase & + OwnedWorkerPlacementMetadata & { + state: "reconciling"; + }; +type ReclaimedPlacementRecord = UnclaimedPlacementRecordBase & + OwnedWorkerPlacementMetadata & { + state: "reclaimed"; + }; +type FailedPlacementRecord = LocalClaimablePlacementRecordBase & + TerminalPlacementMetadata & { + state: "failed"; + recoveryError: string; + }; + +export type WorkerSessionPlacementRecord = + | LocalPlacementRecord + | RequestedPlacementRecord + | ProvisioningPlacementRecord + | SyncingPlacementRecord + | StartingPlacementRecord + | ActivePlacementRecord + | DrainingPlacementRecord + | ReconcilingPlacementRecord + | ReclaimedPlacementRecord + | FailedPlacementRecord; + +export type WorkerSessionPlacementTransitionPatch = { + environmentId?: string | null; + activeOwnerEpoch?: number | null; + workspaceBaseManifestRef?: string | null; + remoteWorkspaceDir?: string | null; + workerBundleHash?: string | null; + lastTranscriptAckCursor?: number | null; + lastLiveEventAckCursor?: number | null; + recoveryError?: string | null; +}; + +export function required(value: string, field: string): string { + const normalized = value.trim(); + if (!normalized) { + throw new Error(`Worker session placement ${field} must be a non-empty string`); + } + return normalized; +} + +export function nullableRequired(value: string | null, field: string): string | null { + return value === null ? null : required(value, field); +} + +export function normalizeEpoch(value: number, field: string): number { + if (!Number.isSafeInteger(value) || value < 1) { + throw new Error(`Worker session placement ${field} must be a positive safe integer`); + } + return value; +} + +export function normalizeCursor(value: number | null, field: string): number | null { + if (value !== null && (!Number.isSafeInteger(value) || value < 0)) { + throw new Error(`Worker session placement ${field} must be a non-negative safe integer`); + } + return value; +} + +export function advanceCursor( + current: number | null, + value: number | undefined, + field: string, +): number | null { + if (value === undefined) { + return current; + } + const next = normalizeCursor(value, field); + if (next === null || current === null) { + return next ?? current; + } + return Math.max(current, next); +} + +export function normalizeIdentity( + input: WorkerSessionPlacementIdentity, +): WorkerSessionPlacementIdentity { + return { + sessionId: required(input.sessionId, "session id"), + agentId: required(input.agentId, "agent id"), + sessionKey: required(input.sessionKey, "session key"), + }; +} + +export function nextGeneration(generation: number): number { + const next = generation + 1; + if (!Number.isSafeInteger(next)) { + throw new Error("Worker session placement generation is exhausted"); + } + return next; +} + +export function localTurnClaimForState( + turnClaim: PersistedTurnClaim | null, + state: "local" | "requested" | "failed", +): PersistedLocalTurnClaim | null { + if (turnClaim?.owner === "worker") { + throw new Error(`Worker turn claim cannot survive placement ${state}`); + } + return turnClaim; +} + +export function workerTurnClaimForState( + turnClaim: PersistedTurnClaim | null, + state: "active" | "draining", +): PersistedWorkerTurnClaim | null { + if (turnClaim?.owner === "local") { + throw new Error(`Local turn claim cannot survive placement ${state}`); + } + return turnClaim; +} + +export function unclaimedTurnForState( + turnClaim: PersistedTurnClaim | null, + state: "provisioning" | "syncing" | "starting" | "reconciling" | "reclaimed", +): null { + if (turnClaim !== null) { + throw new Error(`Turn claim cannot survive placement ${state}`); + } + return null; +} + +export function assertRecordShape(record: { + state: WorkerSessionPlacementState; + environmentId: string | null; + activeOwnerEpoch: number | null; + workspaceBaseManifestRef: string | null; + remoteWorkspaceDir: string | null; + workerBundleHash: string | null; + lastTranscriptAckCursor: number | null; + lastLiveEventAckCursor: number | null; + recoveryError: string | null; + turnClaim: PersistedTurnClaim | null; +}): void { + if (record.state === "local" || record.state === "requested") { + if ( + record.environmentId !== null || + record.activeOwnerEpoch !== null || + record.workspaceBaseManifestRef !== null || + record.remoteWorkspaceDir !== null || + record.workerBundleHash !== null || + record.lastTranscriptAckCursor !== null || + record.lastLiveEventAckCursor !== null || + record.recoveryError !== null + ) { + throw new Error(`Worker session placement ${record.state} cannot retain worker metadata`); + } + } else if (record.state === "provisioning") { + if ( + record.activeOwnerEpoch !== null || + record.workspaceBaseManifestRef !== null || + record.remoteWorkspaceDir !== null || + record.workerBundleHash !== null || + record.lastTranscriptAckCursor !== null || + record.lastLiveEventAckCursor !== null || + record.recoveryError !== null + ) { + throw new Error("Provisioning worker session placement can only retain an environment id"); + } + } else if (record.state === "syncing") { + if ( + !record.environmentId || + record.activeOwnerEpoch !== null || + record.workspaceBaseManifestRef !== null || + record.remoteWorkspaceDir !== null || + !record.workerBundleHash || + record.lastTranscriptAckCursor !== null || + record.lastLiveEventAckCursor !== null || + record.recoveryError !== null + ) { + throw new Error("Syncing worker session placement requires an environment and bundle"); + } + } else if (record.state === "starting") { + if ( + !record.environmentId || + record.activeOwnerEpoch !== null || + !record.workspaceBaseManifestRef || + !record.remoteWorkspaceDir || + !record.workerBundleHash || + record.lastTranscriptAckCursor !== null || + record.lastLiveEventAckCursor !== null || + record.recoveryError !== null + ) { + throw new Error("Starting worker session placement requires complete workspace metadata"); + } + } else if ( + record.state === "active" || + record.state === "draining" || + record.state === "reconciling" || + record.state === "reclaimed" + ) { + if ( + !record.environmentId || + record.activeOwnerEpoch === null || + !record.workspaceBaseManifestRef || + !record.remoteWorkspaceDir || + !record.workerBundleHash || + record.recoveryError !== null + ) { + throw new Error( + `Worker session placement ${record.state} requires complete worker ownership`, + ); + } + normalizeEpoch(record.activeOwnerEpoch, "active owner epoch"); + } else if (!record.recoveryError) { + throw new Error("Failed worker session placement requires a recovery error"); + } + if ( + record.turnClaim?.owner === "local" && + record.state !== "local" && + record.state !== "requested" && + record.state !== "failed" + ) { + throw new Error("Local turn claim requires local, dispatch-barrier, or failed placement"); + } + if (record.turnClaim?.owner === "worker") { + const workerMayFinish = record.state === "active" || record.state === "draining"; + if (!workerMayFinish || record.activeOwnerEpoch !== record.turnClaim.ownerEpoch) { + throw new Error("Worker turn claim requires the active or draining worker owner epoch"); + } + } +} diff --git a/src/gateway/worker-environments/placement-row-codec.ts b/src/gateway/worker-environments/placement-row-codec.ts new file mode 100644 index 000000000000..b68e1d4e0b28 --- /dev/null +++ b/src/gateway/worker-environments/placement-row-codec.ts @@ -0,0 +1,431 @@ +import type { DatabaseSync } from "node:sqlite"; +import type { Selectable } from "kysely"; +import { + executeSqliteQuerySync, + executeSqliteQueryTakeFirstSync, + getNodeSqliteKysely, +} from "../../infra/kysely-sync.js"; +import type { + DB as StateDatabase, + WorkerSessionPlacements, +} from "../../state/openclaw-state-db.generated.js"; +import { + assertRecordShape, + localTurnClaimForState, + nextGeneration, + normalizeCursor, + normalizeEpoch, + nullableRequired, + required, + unclaimedTurnForState, + workerTurnClaimForState, + type EmptyWorkerPlacementMetadata, + type OwnedWorkerPlacementMetadata, + type PersistedTurnClaim, + type WorkerSessionPlacementIdentity, + type WorkerSessionPlacementRecord, + type WorkerSessionPlacementTransitionPatch, +} from "./placement-record.js"; +import { parseWorkerSessionPlacementState } from "./placement-state.js"; + +type PlacementRow = Selectable; +type PlacementDatabase = Pick; + +export const query = (db: DatabaseSync) => getNodeSqliteKysely(db); + +const EMPTY_WORKER_METADATA: EmptyWorkerPlacementMetadata = { + environmentId: null, + activeOwnerEpoch: null, + workspaceBaseManifestRef: null, + remoteWorkspaceDir: null, + workerBundleHash: null, + lastTranscriptAckCursor: null, + lastLiveEventAckCursor: null, + recoveryError: null, +}; + +function parseTurnClaim(row: PlacementRow): PersistedTurnClaim | null { + if (row.turn_claim_owner === null) { + return null; + } + const claimId = required(row.turn_claim_id ?? "", "turn claim id"); + const runId = required(row.turn_claim_run_id ?? "", "turn claim run id"); + const generation = row.turn_claim_generation; + if (generation === null || !Number.isSafeInteger(generation) || generation < 0) { + throw new Error("Worker session placement turn claim generation is invalid"); + } + if (row.turn_claim_owner === "local") { + if (row.turn_claim_owner_epoch !== null) { + throw new Error("Local turn claim cannot retain a worker owner epoch"); + } + return { owner: "local", claimId, runId, generation, ownerEpoch: null }; + } + if (row.turn_claim_owner === "worker") { + return { + owner: "worker", + claimId, + runId, + generation, + ownerEpoch: normalizeEpoch(row.turn_claim_owner_epoch ?? 0, "turn claim owner epoch"), + }; + } + throw new Error(`Invalid worker session turn claim owner: ${row.turn_claim_owner}`); +} + +type ParsedWorkerMetadata = { + environmentId: string | null; + activeOwnerEpoch: number | null; + workspaceBaseManifestRef: string | null; + remoteWorkspaceDir: string | null; + workerBundleHash: string | null; + lastTranscriptAckCursor: number | null; + lastLiveEventAckCursor: number | null; +}; + +function ownedWorkerMetadata( + parsed: ParsedWorkerMetadata, + state: "active" | "draining" | "reconciling" | "reclaimed", +): OwnedWorkerPlacementMetadata { + if ( + parsed.environmentId === null || + parsed.activeOwnerEpoch === null || + parsed.workspaceBaseManifestRef === null || + parsed.remoteWorkspaceDir === null || + parsed.workerBundleHash === null + ) { + throw new Error(`Worker session placement ${state} requires complete worker ownership`); + } + return { + environmentId: parsed.environmentId, + activeOwnerEpoch: parsed.activeOwnerEpoch, + workspaceBaseManifestRef: parsed.workspaceBaseManifestRef, + remoteWorkspaceDir: parsed.remoteWorkspaceDir, + workerBundleHash: parsed.workerBundleHash, + lastTranscriptAckCursor: parsed.lastTranscriptAckCursor, + lastLiveEventAckCursor: parsed.lastLiveEventAckCursor, + recoveryError: null, + }; +} + +export function fromRow(row: PlacementRow): WorkerSessionPlacementRecord { + const state = parseWorkerSessionPlacementState(row.state); + const parsed: ParsedWorkerMetadata = { + environmentId: + row.environment_id === null ? null : required(row.environment_id, "environment id"), + activeOwnerEpoch: + row.active_owner_epoch === null + ? null + : normalizeEpoch(row.active_owner_epoch, "active owner epoch"), + workspaceBaseManifestRef: nullableRequired( + row.workspace_base_manifest_ref, + "workspace base manifest ref", + ), + remoteWorkspaceDir: nullableRequired(row.remote_workspace_dir, "remote workspace directory"), + workerBundleHash: nullableRequired(row.worker_bundle_hash, "worker bundle hash"), + lastTranscriptAckCursor: normalizeCursor( + row.last_transcript_ack_cursor, + "transcript ACK cursor", + ), + lastLiveEventAckCursor: normalizeCursor(row.last_live_event_ack_cursor, "live ACK cursor"), + }; + const recoveryError = nullableRequired(row.recovery_error, "recovery error"); + const turnClaim = parseTurnClaim(row); + const base = { + sessionId: row.session_id, + agentId: row.agent_id, + sessionKey: row.session_key, + generation: row.transition_generation, + createdAtMs: row.created_at_ms, + updatedAtMs: row.updated_at_ms, + stateChangedAtMs: row.state_changed_at_ms, + }; + assertRecordShape({ state, ...parsed, recoveryError, turnClaim }); + switch (state) { + case "local": { + return { + ...base, + state, + turnClaim: localTurnClaimForState(turnClaim, state), + ...EMPTY_WORKER_METADATA, + }; + } + case "requested": { + return { + ...base, + state, + turnClaim: localTurnClaimForState(turnClaim, state), + ...EMPTY_WORKER_METADATA, + }; + } + case "provisioning": { + return { + ...base, + state, + turnClaim: unclaimedTurnForState(turnClaim, state), + ...EMPTY_WORKER_METADATA, + environmentId: parsed.environmentId, + }; + } + case "syncing": { + if (parsed.environmentId === null || parsed.workerBundleHash === null) { + throw new Error("Syncing worker session placement requires an environment and bundle"); + } + return { + ...base, + state, + turnClaim: unclaimedTurnForState(turnClaim, state), + ...EMPTY_WORKER_METADATA, + environmentId: parsed.environmentId, + workerBundleHash: parsed.workerBundleHash, + }; + } + case "starting": { + if ( + parsed.environmentId === null || + parsed.workspaceBaseManifestRef === null || + parsed.remoteWorkspaceDir === null || + parsed.workerBundleHash === null + ) { + throw new Error("Starting worker session placement requires complete workspace metadata"); + } + return { + ...base, + state, + turnClaim: unclaimedTurnForState(turnClaim, state), + ...EMPTY_WORKER_METADATA, + environmentId: parsed.environmentId, + workspaceBaseManifestRef: parsed.workspaceBaseManifestRef, + remoteWorkspaceDir: parsed.remoteWorkspaceDir, + workerBundleHash: parsed.workerBundleHash, + }; + } + case "active": { + return { + ...base, + state, + turnClaim: workerTurnClaimForState(turnClaim, state), + ...ownedWorkerMetadata(parsed, state), + }; + } + case "draining": { + return { + ...base, + state, + turnClaim: workerTurnClaimForState(turnClaim, state), + ...ownedWorkerMetadata(parsed, state), + }; + } + case "reconciling": { + return { + ...base, + state, + turnClaim: unclaimedTurnForState(turnClaim, state), + ...ownedWorkerMetadata(parsed, state), + }; + } + case "reclaimed": { + return { + ...base, + state, + turnClaim: unclaimedTurnForState(turnClaim, state), + ...ownedWorkerMetadata(parsed, state), + }; + } + case "failed": { + if (recoveryError === null) { + throw new Error("Failed worker session placement requires a recovery error"); + } + return { + ...base, + state, + turnClaim: localTurnClaimForState(turnClaim, state), + environmentId: parsed.environmentId, + activeOwnerEpoch: parsed.activeOwnerEpoch, + workspaceBaseManifestRef: parsed.workspaceBaseManifestRef, + remoteWorkspaceDir: parsed.remoteWorkspaceDir, + workerBundleHash: parsed.workerBundleHash, + lastTranscriptAckCursor: parsed.lastTranscriptAckCursor, + lastLiveEventAckCursor: parsed.lastLiveEventAckCursor, + recoveryError, + }; + } + } + // Exhaustive over placement states; the return satisfies consistent-return. + return state satisfies never; +} + +export function find( + db: DatabaseSync, + sessionId: string, +): WorkerSessionPlacementRecord | undefined { + const row = executeSqliteQueryTakeFirstSync( + db, + query(db) + .selectFrom("worker_session_placements") + .selectAll() + .where("session_id", "=", sessionId), + ); + return row ? fromRow(row) : undefined; +} + +export function getRequired(db: DatabaseSync, sessionId: string): WorkerSessionPlacementRecord { + const record = find(db, sessionId); + if (!record) { + throw new Error(`Unknown worker session placement: ${sessionId}`); + } + return record; +} + +function assertIdentity( + record: WorkerSessionPlacementRecord, + identity: WorkerSessionPlacementIdentity, +): void { + if (record.agentId !== identity.agentId || record.sessionKey !== identity.sessionKey) { + throw new Error(`Worker session placement identity changed for ${identity.sessionId}`); + } +} + +function insertLocal( + db: DatabaseSync, + identity: WorkerSessionPlacementIdentity, + nowMs: number, +): WorkerSessionPlacementRecord { + executeSqliteQuerySync( + db, + query(db).insertInto("worker_session_placements").values({ + session_id: identity.sessionId, + agent_id: identity.agentId, + session_key: identity.sessionKey, + state: "local", + environment_id: null, + transition_generation: 0, + active_owner_epoch: null, + workspace_base_manifest_ref: null, + remote_workspace_dir: null, + worker_bundle_hash: null, + last_transcript_ack_cursor: null, + last_live_event_ack_cursor: null, + recovery_error: null, + turn_claim_owner: null, + turn_claim_id: null, + turn_claim_run_id: null, + turn_claim_generation: null, + turn_claim_owner_epoch: null, + created_at_ms: nowMs, + updated_at_ms: nowMs, + state_changed_at_ms: nowMs, + }), + ); + return getRequired(db, identity.sessionId); +} + +export function ensureLocal( + db: DatabaseSync, + identity: WorkerSessionPlacementIdentity, + nowMs: number, +): WorkerSessionPlacementRecord { + const current = find(db, identity.sessionId); + if (current) { + assertIdentity(current, identity); + return current; + } + return insertLocal(db, identity, nowMs); +} + +export function transitionValues( + current: WorkerSessionPlacementRecord, + to: WorkerSessionPlacementRecord["state"], + patch: WorkerSessionPlacementTransitionPatch, + nowMs: number, +): PlacementRow { + const environmentId = + to === "local" || to === "requested" + ? null + : patch.environmentId === undefined + ? current.environmentId + : patch.environmentId === null + ? null + : required(patch.environmentId, "environment id"); + const activeOwnerEpoch = + to === "local" || + to === "requested" || + to === "provisioning" || + to === "syncing" || + to === "starting" + ? null + : patch.activeOwnerEpoch === undefined + ? current.activeOwnerEpoch + : patch.activeOwnerEpoch === null + ? null + : normalizeEpoch(patch.activeOwnerEpoch, "active owner epoch"); + const generation = nextGeneration(current.generation); + const clearsWorkerMetadata = to === "local" || to === "requested"; + const values: PlacementRow = { + session_id: current.sessionId, + agent_id: current.agentId, + session_key: current.sessionKey, + state: to, + environment_id: environmentId, + transition_generation: generation, + active_owner_epoch: activeOwnerEpoch, + workspace_base_manifest_ref: clearsWorkerMetadata + ? null + : patch.workspaceBaseManifestRef === undefined + ? current.workspaceBaseManifestRef + : patch.workspaceBaseManifestRef === null + ? null + : required(patch.workspaceBaseManifestRef, "workspace base manifest ref"), + remote_workspace_dir: clearsWorkerMetadata + ? null + : patch.remoteWorkspaceDir === undefined + ? current.remoteWorkspaceDir + : patch.remoteWorkspaceDir === null + ? null + : required(patch.remoteWorkspaceDir, "remote workspace directory"), + worker_bundle_hash: clearsWorkerMetadata + ? null + : patch.workerBundleHash === undefined + ? current.workerBundleHash + : patch.workerBundleHash === null + ? null + : required(patch.workerBundleHash, "worker bundle hash"), + last_transcript_ack_cursor: clearsWorkerMetadata + ? null + : patch.lastTranscriptAckCursor === undefined + ? current.lastTranscriptAckCursor + : normalizeCursor(patch.lastTranscriptAckCursor, "transcript ACK cursor"), + last_live_event_ack_cursor: clearsWorkerMetadata + ? null + : patch.lastLiveEventAckCursor === undefined + ? current.lastLiveEventAckCursor + : normalizeCursor(patch.lastLiveEventAckCursor, "live ACK cursor"), + recovery_error: clearsWorkerMetadata + ? null + : patch.recoveryError === undefined + ? current.recoveryError + : patch.recoveryError === null + ? null + : required(patch.recoveryError, "recovery error"), + turn_claim_owner: null, + turn_claim_id: null, + turn_claim_run_id: null, + turn_claim_generation: null, + turn_claim_owner_epoch: null, + created_at_ms: current.createdAtMs, + updated_at_ms: nowMs, + state_changed_at_ms: nowMs, + }; + assertRecordShape({ + state: to, + environmentId, + activeOwnerEpoch, + workspaceBaseManifestRef: values.workspace_base_manifest_ref, + remoteWorkspaceDir: values.remote_workspace_dir, + workerBundleHash: values.worker_bundle_hash, + lastTranscriptAckCursor: values.last_transcript_ack_cursor, + lastLiveEventAckCursor: values.last_live_event_ack_cursor, + recoveryError: values.recovery_error, + turnClaim: null, + }); + return values; +} diff --git a/src/gateway/worker-environments/placement-session-runtime.ts b/src/gateway/worker-environments/placement-session-runtime.ts new file mode 100644 index 000000000000..3729416a118a --- /dev/null +++ b/src/gateway/worker-environments/placement-session-runtime.ts @@ -0,0 +1,33 @@ +import { + isDefaultAgentRuntimeId, + OPENCLAW_AGENT_RUNTIME_ID, +} from "../../agents/agent-runtime-id.js"; +import { resolveSessionModelRef } from "../../agents/session-model-ref.js"; +import { resolvePersistedSessionRuntimeId } from "../../agents/session-runtime-compat.js"; +import { resolveEffectiveAgentRuntime } from "../../agents/thinking-runtime.js"; +import type { SessionEntry } from "../../config/sessions.js"; +import type { OpenClawConfig } from "../../config/types.openclaw.js"; + +export function resolveWorkerPlacementSessionRuntime(params: { + cfg: OpenClawConfig; + entry: SessionEntry; + agentId: string; + sessionKey: string; +}): string { + const persistedRuntime = resolvePersistedSessionRuntimeId(params.entry); + if (persistedRuntime && !isDefaultAgentRuntimeId(persistedRuntime)) { + return persistedRuntime; + } + const selectedModel = resolveSessionModelRef(params.cfg, params.entry, params.agentId); + return resolveEffectiveAgentRuntime({ + cfg: params.cfg, + provider: selectedModel.provider, + modelId: selectedModel.model, + agentId: params.agentId, + sessionKey: params.sessionKey, + }); +} + +export function isWorkerPlacementSessionRuntimeSupported(runtime: string): boolean { + return runtime === OPENCLAW_AGENT_RUNTIME_ID; +} diff --git a/src/gateway/worker-environments/placement-state.ts b/src/gateway/worker-environments/placement-state.ts new file mode 100644 index 000000000000..c81204cabbaa --- /dev/null +++ b/src/gateway/worker-environments/placement-state.ts @@ -0,0 +1,47 @@ +const WORKER_SESSION_PLACEMENT_STATES = [ + "local", + "requested", + "provisioning", + "syncing", + "starting", + "active", + "draining", + "reconciling", + "reclaimed", + "failed", +] as const; + +export type WorkerSessionPlacementState = (typeof WORKER_SESSION_PLACEMENT_STATES)[number]; + +type WorkerSessionPlacementTransition = { + [From in WorkerSessionPlacementState]: readonly WorkerSessionPlacementState[]; +}; + +const WORKER_SESSION_PLACEMENT_TRANSITIONS = { + local: ["requested"], + requested: ["provisioning", "failed"], + provisioning: ["syncing", "failed"], + syncing: ["starting", "failed"], + starting: ["active", "failed"], + active: ["draining"], + draining: ["reconciling"], + reconciling: ["local", "reclaimed", "failed"], + reclaimed: ["requested"], + failed: [], +} as const satisfies WorkerSessionPlacementTransition; + +export function parseWorkerSessionPlacementState(value: string): WorkerSessionPlacementState { + if ((WORKER_SESSION_PLACEMENT_STATES as readonly string[]).includes(value)) { + return value as WorkerSessionPlacementState; + } + throw new Error(`Invalid worker session placement state: ${value}`); +} + +export function canTransitionWorkerSessionPlacement( + from: WorkerSessionPlacementState, + to: WorkerSessionPlacementState, +): boolean { + return ( + WORKER_SESSION_PLACEMENT_TRANSITIONS[from] as readonly WorkerSessionPlacementState[] + ).includes(to); +} diff --git a/src/gateway/worker-environments/placement-store.test.ts b/src/gateway/worker-environments/placement-store.test.ts new file mode 100644 index 000000000000..1e17cf1f0ee3 --- /dev/null +++ b/src/gateway/worker-environments/placement-store.test.ts @@ -0,0 +1,695 @@ +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { + closeOpenClawStateDatabaseForTest, + openOpenClawStateDatabase, + type OpenClawStateDatabase, +} from "../../state/openclaw-state-db.js"; +import type { WorkerSessionPlacementIdentity } from "./placement-record.js"; +import { + createWorkerSessionPlacementStore, + type WorkerSessionPlacementStore, +} from "./placement-store.js"; + +const SESSION: WorkerSessionPlacementIdentity = { + sessionId: "session-placement", + agentId: "main", + sessionKey: "agent:main:placement", +}; + +describe("worker session placement store", () => { + let root: string; + let database: OpenClawStateDatabase; + let store: WorkerSessionPlacementStore; + let nowMs: number; + + beforeEach(async () => { + root = await fs.mkdtemp(path.join(await fs.realpath(os.tmpdir()), "openclaw-placement-")); + database = openOpenClawStateDatabase({ env: { OPENCLAW_STATE_DIR: root } }); + nowMs = 1_000; + store = createWorkerSessionPlacementStore({ database, now: () => nowMs }); + }); + + afterEach(async () => { + closeOpenClawStateDatabaseForTest(); + await fs.rm(root, { recursive: true, force: true }); + }); + + function advanceToActive(identity: WorkerSessionPlacementIdentity = SESSION) { + let placement = store.startDispatch(identity); + placement = store.transition({ + sessionId: identity.sessionId, + from: "requested", + to: "provisioning", + expectedGeneration: placement.generation, + patch: { environmentId: `environment-${identity.sessionId}` }, + }); + placement = store.transition({ + sessionId: identity.sessionId, + from: "provisioning", + to: "syncing", + expectedGeneration: placement.generation, + patch: { workerBundleHash: "a".repeat(64) }, + }); + placement = store.transition({ + sessionId: identity.sessionId, + from: "syncing", + to: "starting", + expectedGeneration: placement.generation, + patch: { + workspaceBaseManifestRef: `manifest-${identity.sessionId}`, + remoteWorkspaceDir: `/workspace/${identity.sessionId}`, + }, + }); + const active = store.transition({ + sessionId: identity.sessionId, + from: "starting", + to: "active", + expectedGeneration: placement.generation, + patch: { activeOwnerEpoch: 7 }, + }); + if (active.state !== "active") { + throw new Error("expected active worker placement"); + } + return active; + } + + it("persists the placement lifecycle and rejects stale transition generations", () => { + const requested = store.startDispatch(SESSION); + expect(requested).toMatchObject({ + state: "requested", + generation: 1, + environmentId: null, + activeOwnerEpoch: null, + }); + + const provisioning = store.transition({ + sessionId: SESSION.sessionId, + from: "requested", + to: "provisioning", + expectedGeneration: requested.generation, + patch: { environmentId: "environment-placement" }, + }); + expect(provisioning).toMatchObject({ + state: "provisioning", + generation: 2, + environmentId: "environment-placement", + }); + expect(() => + store.transition({ + sessionId: SESSION.sessionId, + from: "provisioning", + to: "syncing", + expectedGeneration: 1, + }), + ).toThrow("expected provisioning@1, found provisioning@2"); + expect(() => + store.transition({ + sessionId: SESSION.sessionId, + from: "provisioning", + to: "active", + expectedGeneration: provisioning.generation, + }), + ).toThrow("Illegal worker session placement transition"); + + const failed = store.fail({ + sessionId: SESSION.sessionId, + expectedGeneration: provisioning.generation, + recoveryError: "workspace synchronization failed", + }); + expect(failed).toMatchObject({ + state: "failed", + generation: 3, + recoveryError: "workspace synchronization failed", + }); + expect(() => + store.fail({ + sessionId: SESSION.sessionId, + expectedGeneration: failed.generation - 1, + recoveryError: "stale teardown failure", + }), + ).toThrow("changed before failure"); + expect(store.get(SESSION.sessionId)?.recoveryError).toBe("workspace synchronization failed"); + expect( + store.fail({ sessionId: SESSION.sessionId, recoveryError: "teardown retry failed" }), + ).toMatchObject({ + state: "failed", + generation: failed.generation, + recoveryError: "teardown retry failed", + }); + }); + + it("requires each placement phase to persist its complete metadata", () => { + const requested = store.startDispatch(SESSION); + const provisioning = store.transition({ + sessionId: SESSION.sessionId, + from: "requested", + to: "provisioning", + expectedGeneration: requested.generation, + patch: { environmentId: "environment-placement" }, + }); + expect(provisioning).toMatchObject({ + workspaceBaseManifestRef: null, + remoteWorkspaceDir: null, + workerBundleHash: null, + lastTranscriptAckCursor: null, + lastLiveEventAckCursor: null, + }); + + expect(() => + store.transition({ + sessionId: SESSION.sessionId, + from: "provisioning", + to: "syncing", + expectedGeneration: provisioning.generation, + }), + ).toThrow("requires an environment and bundle"); + const syncing = store.transition({ + sessionId: SESSION.sessionId, + from: "provisioning", + to: "syncing", + expectedGeneration: provisioning.generation, + patch: { workerBundleHash: "a".repeat(64) }, + }); + + expect(() => + store.transition({ + sessionId: SESSION.sessionId, + from: "syncing", + to: "starting", + expectedGeneration: syncing.generation, + patch: { workspaceBaseManifestRef: "manifest-placement" }, + }), + ).toThrow("requires complete workspace metadata"); + expect( + store.transition({ + sessionId: SESSION.sessionId, + from: "syncing", + to: "starting", + expectedGeneration: syncing.generation, + patch: { + workspaceBaseManifestRef: "manifest-placement", + remoteWorkspaceDir: "/workspace/placement", + }, + }), + ).toMatchObject({ + state: "starting", + environmentId: "environment-placement", + workerBundleHash: "a".repeat(64), + workspaceBaseManifestRef: "manifest-placement", + remoteWorkspaceDir: "/workspace/placement", + }); + }); + + it("drains and reconciles worker ownership before returning local", () => { + const active = advanceToActive(); + const draining = store.transition({ + sessionId: SESSION.sessionId, + from: "active", + to: "draining", + expectedGeneration: active.generation, + }); + const reconciling = store.startReconcile({ + sessionId: SESSION.sessionId, + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + expectedGeneration: draining.generation, + }); + const local = store.transition({ + sessionId: SESSION.sessionId, + from: "reconciling", + to: "local", + expectedGeneration: reconciling.generation, + }); + expect(local).toMatchObject({ + state: "local", + environmentId: null, + activeOwnerEpoch: null, + }); + }); + + it("rejects reclaim before worker ownership reaches reconciliation", () => { + const requested = store.startDispatch(SESSION); + expect(() => + store.transition({ + sessionId: SESSION.sessionId, + from: "requested", + to: "reclaimed", + expectedGeneration: requested.generation, + }), + ).toThrow("Illegal worker session placement transition"); + expect( + store.fail({ + sessionId: SESSION.sessionId, + expectedGeneration: requested.generation, + recoveryError: "dispatch stopped before provisioning", + }), + ).toMatchObject({ state: "failed" }); + }); + + it("closes local admission before draining the existing local turn", async () => { + const localClaim = store.claimTurn({ + ...SESSION, + owner: { kind: "local" }, + claimId: "local-claim", + runId: "run-local", + }); + const requested = store.startDispatch(SESSION); + expect(requested).toMatchObject({ state: "requested", generation: 1 }); + expect(requested.turnClaim).toMatchObject({ owner: "local", generation: 0 }); + + expect(() => + store.claimTurn({ + ...SESSION, + owner: { kind: "local" }, + claimId: "new-local-claim", + runId: "new-local-run", + }), + ).toThrow("already has an active turn claim"); + expect(() => + store.claimTurn({ + ...SESSION, + owner: localClaim.owner, + claimId: localClaim.claimId, + runId: localClaim.runId, + }), + ).toThrow("already has an active turn claim"); + expect(() => + store.transition({ + sessionId: SESSION.sessionId, + from: "requested", + to: "provisioning", + expectedGeneration: requested.generation, + patch: { environmentId: "environment-placement" }, + }), + ).toThrow("during an active turn"); + + const released = store.waitForTurnClaimRelease(SESSION.sessionId, { timeoutMs: 1_000 }); + store.releaseTurn(localClaim); + await released; + expect( + store.transition({ + sessionId: SESSION.sessionId, + from: "requested", + to: "provisioning", + expectedGeneration: requested.generation, + patch: { environmentId: "environment-placement" }, + }), + ).toMatchObject({ state: "provisioning", turnClaim: null }); + }); + + it("keeps the draining local claim releasable when the dispatch barrier fails", () => { + const localClaim = store.claimTurn({ + ...SESSION, + owner: { kind: "local" }, + claimId: "local-barrier-claim", + runId: "local-barrier-run", + }); + const requested = store.startDispatch(SESSION); + const failed = store.fail({ + sessionId: SESSION.sessionId, + expectedGeneration: requested.generation, + recoveryError: "local drain timed out", + }); + + expect(failed).toMatchObject({ + state: "failed", + recoveryError: "local drain timed out", + turnClaim: { owner: "local", claimId: localClaim.claimId }, + }); + expect(() => + store.claimTurn({ + ...SESSION, + owner: { kind: "local" }, + claimId: "new-local-claim", + runId: "new-local-run", + }), + ).toThrow("already has an active turn claim"); + expect(store.releaseTurn(localClaim)).toMatchObject({ state: "failed", turnClaim: null }); + }); + + it("does not let a stale claim release a later turn that reuses the run id", () => { + const firstClaim = store.claimTurn({ + ...SESSION, + owner: { kind: "local" }, + claimId: "first-claim-token", + runId: "reused-run", + }); + store.releaseTurn(firstClaim); + const secondClaim = store.claimTurn({ + ...SESSION, + owner: { kind: "local" }, + claimId: "second-claim-token", + runId: firstClaim.runId, + }); + + expect(() => store.releaseTurn(firstClaim)).toThrow("turn claim changed before release"); + expect(store.validateTurnClaim(secondClaim)).toBe(true); + expect(store.get(SESSION.sessionId)?.turnClaim).toMatchObject({ + claimId: secondClaim.claimId, + runId: secondClaim.runId, + }); + }); + + it("allows a reset session id to reuse its canonical session key", () => { + const firstClaim = store.claimTurn({ + ...SESSION, + owner: { kind: "local" }, + claimId: "first-session-claim", + runId: "first-session-run", + }); + store.releaseTurn(firstClaim); + + const rotated = store.claimTurn({ + ...SESSION, + sessionId: "session-placement-rotated", + owner: { kind: "local" }, + claimId: "rotated-session-claim", + runId: "rotated-session-run", + }); + expect(rotated.sessionId).toBe("session-placement-rotated"); + expect(store.list().map((record) => record.sessionId)).toEqual([ + SESSION.sessionId, + "session-placement-rotated", + ]); + }); + + it("admits exactly the active placement owner and fences stale worker epochs", () => { + const active = advanceToActive(); + expect(() => + store.claimTurn({ + ...SESSION, + owner: { kind: "local" }, + claimId: "local-after-dispatch", + runId: "local-after-dispatch-run", + }), + ).toThrow("Local turn rejected"); + expect(() => + store.claimTurn({ + ...SESSION, + owner: { + kind: "worker", + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch + 1, + }, + claimId: "stale-worker", + runId: "stale-worker-run", + }), + ).toThrow("stale owner"); + + const workerClaim = store.claimTurn({ + ...SESSION, + owner: { + kind: "worker", + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + }, + claimId: "worker-claim", + runId: "worker-run", + }); + expect(store.validateTurnClaim(workerClaim)).toBe(true); + expect( + store.validateTurnClaim({ + ...workerClaim, + owner: { + kind: "worker", + environmentId: "environment-stale", + ownerEpoch: active.activeOwnerEpoch, + }, + }), + ).toBe(false); + expect( + store.validateWorkerOwner({ + sessionId: SESSION.sessionId, + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + }), + ).toBe(true); + expect(() => + store.claimTurn({ + ...SESSION, + owner: { + kind: "worker", + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + }, + claimId: "competing-worker", + runId: "competing-worker-run", + }), + ).toThrow("already has an active turn claim"); + expect(() => + store.claimTurn({ + ...SESSION, + owner: workerClaim.owner, + claimId: workerClaim.claimId, + runId: workerClaim.runId, + }), + ).toThrow("already has an active turn claim"); + expect(() => + store.fail({ + sessionId: SESSION.sessionId, + expectedGeneration: active.generation, + recoveryError: "active worker disappeared", + }), + ).toThrow("Cannot fail worker session placement from active"); + const draining = store.startDrain({ + sessionId: SESSION.sessionId, + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + expectedGeneration: active.generation, + }); + const reconciling = store.startReconcile({ + sessionId: SESSION.sessionId, + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + expectedGeneration: draining.generation, + }); + expect( + store.transition({ + sessionId: SESSION.sessionId, + from: "reconciling", + to: "failed", + expectedGeneration: reconciling.generation, + patch: { recoveryError: "active worker disappeared" }, + }), + ).toMatchObject({ state: "failed", turnClaim: null }); + expect(store.validateTurnClaim(workerClaim)).toBe(false); + }); + + it("clears dead local claims on restart while adopting active worker ownership", () => { + const localIdentity = { + ...SESSION, + sessionId: "session-local-restart", + sessionKey: "agent:main:local-restart", + }; + store.claimTurn({ + ...localIdentity, + owner: { kind: "local" }, + claimId: "local-before-restart", + runId: "local-restart-run", + }); + const active = advanceToActive(); + const workerClaim = store.claimTurn({ + ...SESSION, + owner: { + kind: "worker", + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + }, + claimId: "worker-before-restart", + runId: "worker-restart-run", + }); + + closeOpenClawStateDatabaseForTest(); + database = openOpenClawStateDatabase({ env: { OPENCLAW_STATE_DIR: root } }); + store = createWorkerSessionPlacementStore({ database, now: () => nowMs }); + + expect(store.clearLocalTurnClaimsAfterRestart()).toBe(1); + expect(store.get(localIdentity.sessionId)?.turnClaim).toBeNull(); + expect(store.validateTurnClaim(workerClaim)).toBe(true); + expect( + store.adoptActive({ + sessionId: SESSION.sessionId, + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + expectedGeneration: active.generation, + }), + ).toMatchObject({ state: "active", turnClaim: { owner: "worker" } }); + expect(store.listForReconcile().map((record) => record.sessionId)).toEqual([SESSION.sessionId]); + expect(store.list().map((record) => record.sessionId)).toEqual([ + localIdentity.sessionId, + SESSION.sessionId, + ]); + }); + + it("closes worker admission before draining the active turn", async () => { + const active = advanceToActive(); + const workerClaim = store.claimTurn({ + ...SESSION, + owner: { + kind: "worker", + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + }, + claimId: "worker-drain-claim", + runId: "worker-drain-run", + }); + + const draining = store.startDrain({ + sessionId: SESSION.sessionId, + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + expectedGeneration: active.generation, + }); + expect(draining).toMatchObject({ + state: "draining", + generation: active.generation + 1, + turnClaim: { owner: "worker", claimId: workerClaim.claimId }, + }); + expect(store.validateTurnClaim(workerClaim)).toBe(true); + + const released = store.waitForTurnClaimRelease(SESSION.sessionId, { timeoutMs: 1_000 }); + expect(store.releaseTurn(workerClaim)).toMatchObject({ state: "draining", turnClaim: null }); + await released; + expect(() => + store.claimTurn({ + ...SESSION, + owner: workerClaim.owner, + claimId: "worker-after-drain", + runId: "worker-after-drain-run", + }), + ).toThrow("stale owner"); + expect( + store.startReconcile({ + sessionId: SESSION.sessionId, + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + expectedGeneration: draining.generation, + }), + ).toMatchObject({ state: "reconciling", turnClaim: null }); + }); + + it("atomically fences a drained claim before its worker is reclaimed", async () => { + const active = advanceToActive(); + const workerClaim = store.claimTurn({ + ...SESSION, + owner: { + kind: "worker", + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + }, + claimId: "worker-reclaim-claim", + runId: "worker-reclaim-run", + }); + const released = store.waitForTurnClaimRelease(SESSION.sessionId, { timeoutMs: 1_000 }); + + const draining = store.startDrain({ + sessionId: SESSION.sessionId, + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + expectedGeneration: active.generation, + }); + expect(() => + store.fail({ + sessionId: SESSION.sessionId, + expectedGeneration: draining.generation, + recoveryError: "worker teardown not yet fenced", + }), + ).toThrow("Cannot fail worker session placement from draining"); + expect(() => + store.transition({ + sessionId: SESSION.sessionId, + from: "draining", + to: "reclaimed", + expectedGeneration: draining.generation, + }), + ).toThrow("Illegal worker session placement transition"); + expect(() => + store.transition({ + sessionId: SESSION.sessionId, + from: "draining", + to: "reconciling", + expectedGeneration: draining.generation, + }), + ).toThrow("Use startReconcile after fencing the drained worker environment"); + expect(() => + store.startReconcile({ + sessionId: SESSION.sessionId, + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + expectedGeneration: draining.generation - 1, + }), + ).toThrow("Cannot reconcile stale worker placement"); + const reconciling = store.startReconcile({ + sessionId: SESSION.sessionId, + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + expectedGeneration: draining.generation, + }); + const reclaimed = store.transition({ + sessionId: SESSION.sessionId, + from: "reconciling", + to: "reclaimed", + expectedGeneration: reconciling.generation, + }); + expect(reclaimed).toMatchObject({ state: "reclaimed", turnClaim: null }); + await released; + expect(store.validateTurnClaim(workerClaim)).toBe(false); + expect(store.startDispatch(SESSION)).toMatchObject({ + state: "requested", + generation: reclaimed.generation + 1, + environmentId: null, + activeOwnerEpoch: null, + workspaceBaseManifestRef: null, + remoteWorkspaceDir: null, + workerBundleHash: null, + }); + }); + + it("binds acknowledged cursors to the exact normalized worker claim", () => { + const active = advanceToActive(); + const firstClaim = store.claimTurn({ + ...SESSION, + owner: { + kind: "worker", + environmentId: ` ${active.environmentId} `, + ownerEpoch: active.activeOwnerEpoch, + }, + claimId: "worker-ack-first", + runId: "worker-ack-first-run", + }); + expect(firstClaim.owner).toEqual({ + kind: "worker", + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + }); + store.releaseTurn(firstClaim); + const currentClaim = store.claimTurn({ + ...SESSION, + owner: firstClaim.owner, + claimId: "worker-ack-current", + runId: "worker-ack-current-run", + }); + + expect(() => store.updateAckCursors({ claim: firstClaim, transcript: 4 })).toThrow( + "Cannot ACK stale worker turn", + ); + expect(store.get(SESSION.sessionId)?.lastTranscriptAckCursor).toBeNull(); + expect( + store.updateAckCursors({ + claim: currentClaim, + transcript: 4, + liveEvent: 9, + }), + ).toMatchObject({ lastTranscriptAckCursor: 4, lastLiveEventAckCursor: 9 }); + expect( + store.updateAckCursors({ + claim: currentClaim, + transcript: 3, + liveEvent: 8, + }), + ).toMatchObject({ lastTranscriptAckCursor: 4, lastLiveEventAckCursor: 9 }); + }); +}); diff --git a/src/gateway/worker-environments/placement-store.ts b/src/gateway/worker-environments/placement-store.ts new file mode 100644 index 000000000000..fc5e0f4685b6 --- /dev/null +++ b/src/gateway/worker-environments/placement-store.ts @@ -0,0 +1,414 @@ +import type { DatabaseSync } from "node:sqlite"; +import { executeSqliteQuerySync } from "../../infra/kysely-sync.js"; +import { + openOpenClawStateDatabase, + runOpenClawStateWriteTransaction, + type OpenClawStateDatabase, +} from "../../state/openclaw-state-db.js"; +import { + assertRecordShape, + nextGeneration, + normalizeEpoch, + normalizeIdentity, + required, + type WorkerSessionPlacementIdentity, + type WorkerSessionPlacementRecord, + type WorkerSessionPlacementTransitionPatch, +} from "./placement-record.js"; +import { + ensureLocal, + find, + fromRow, + getRequired, + query, + transitionValues, +} from "./placement-row-codec.js"; +import { + canTransitionWorkerSessionPlacement, + type WorkerSessionPlacementState, +} from "./placement-state.js"; +import { + createPlacementTurnClaimOps, + signalTurnClaimRelease, + type PlacementStoreRuntime, +} from "./placement-turn-claims.js"; + +export type { WorkerSessionPlacementRecord, WorkerSessionTurnClaim } from "./placement-record.js"; + +function updateTransition( + db: DatabaseSync, + current: WorkerSessionPlacementRecord, + to: WorkerSessionPlacementState, + patch: WorkerSessionPlacementTransitionPatch, + nowMs: number, +): WorkerSessionPlacementRecord { + const values = transitionValues(current, to, patch, nowMs); + const result = executeSqliteQuerySync( + db, + query(db) + .updateTable("worker_session_placements") + .set(values) + .where("session_id", "=", current.sessionId) + .where("state", "=", current.state) + .where("transition_generation", "=", current.generation) + .where("turn_claim_owner", "is", null), + ); + if (result.numAffectedRows !== 1n) { + throw new Error(`Worker session placement ${current.sessionId} changed during transition`); + } + return getRequired(db, current.sessionId); +} + +export function createWorkerSessionPlacementStore( + options: { database?: OpenClawStateDatabase; now?: () => number } = {}, +) { + const path = (options.database ?? openOpenClawStateDatabase()).path; + const now = options.now ?? Date.now; + const runtime: PlacementStoreRuntime = { + path, + now, + read: () => openOpenClawStateDatabase({ path }).db, + write: (operation) => runOpenClawStateWriteTransaction(({ db }) => operation(db), { path }), + }; + const { read, write } = runtime; + + return { + ...createPlacementTurnClaimOps(runtime), + + get(sessionId: string): WorkerSessionPlacementRecord | undefined { + return find(read(), required(sessionId, "session id")); + }, + + getMany(sessionIds: readonly string[]): ReadonlyMap { + const normalizedIds = [ + ...new Set(sessionIds.map((sessionId) => required(sessionId, "session id"))), + ]; + const records = new Map(); + const db = read(); + for (let offset = 0; offset < normalizedIds.length; offset += 250) { + const chunk = normalizedIds.slice(offset, offset + 250); + for (const row of executeSqliteQuerySync( + db, + query(db) + .selectFrom("worker_session_placements") + .selectAll() + .where("session_id", "in", chunk), + ).rows) { + const record = fromRow(row); + records.set(record.sessionId, record); + } + } + return records; + }, + + startDispatch(input: WorkerSessionPlacementIdentity): WorkerSessionPlacementRecord { + const identity = normalizeIdentity(input); + return write((db) => { + const current = ensureLocal(db, identity, now()); + if (current.state !== "local" && current.state !== "reclaimed") { + throw new Error( + `Cannot dispatch session ${identity.sessionId} from placement ${current.state}`, + ); + } + const updatedAtMs = now(); + // Preserve an in-flight local claim while closing admission. Reclaimed + // placement has no live owner and starts a fresh worker generation. + const result = executeSqliteQuerySync( + db, + query(db) + .updateTable("worker_session_placements") + .set({ + state: "requested", + environment_id: null, + transition_generation: nextGeneration(current.generation), + active_owner_epoch: null, + workspace_base_manifest_ref: null, + remote_workspace_dir: null, + worker_bundle_hash: null, + last_transcript_ack_cursor: null, + last_live_event_ack_cursor: null, + recovery_error: null, + updated_at_ms: updatedAtMs, + state_changed_at_ms: updatedAtMs, + }) + .where("session_id", "=", current.sessionId) + .where("state", "=", current.state) + .where("transition_generation", "=", current.generation), + ); + if (result.numAffectedRows !== 1n) { + throw new Error( + `Session ${identity.sessionId} placement changed during dispatch barrier`, + ); + } + return getRequired(db, identity.sessionId); + }); + }, + + transition(input: { + sessionId: string; + from: WorkerSessionPlacementState; + to: WorkerSessionPlacementState; + expectedGeneration: number; + patch?: WorkerSessionPlacementTransitionPatch; + }): WorkerSessionPlacementRecord { + if (!canTransitionWorkerSessionPlacement(input.from, input.to)) { + throw new Error( + `Illegal worker session placement transition: ${input.from} -> ${input.to}`, + ); + } + if (input.from === "draining" && input.to === "reconciling") { + throw new Error("Use startReconcile after fencing the drained worker environment"); + } + const sessionId = required(input.sessionId, "session id"); + return write((db) => { + const current = getRequired(db, sessionId); + if (current.state !== input.from || current.generation !== input.expectedGeneration) { + throw new Error( + `Worker session placement ${sessionId} changed: expected ${input.from}@${input.expectedGeneration}, found ${current.state}@${current.generation}`, + ); + } + if (current.turnClaim) { + throw new Error(`Cannot transition session ${sessionId} during an active turn`); + } + return updateTransition(db, current, input.to, input.patch ?? {}, now()); + }); + }, + + startDrain(input: { + sessionId: string; + environmentId: string; + ownerEpoch: number; + expectedGeneration: number; + }): WorkerSessionPlacementRecord { + const sessionId = required(input.sessionId, "session id"); + const environmentId = required(input.environmentId, "environment id"); + const ownerEpoch = normalizeEpoch(input.ownerEpoch, "active owner epoch"); + return write((db) => { + const current = getRequired(db, sessionId); + if ( + current.state !== "active" || + current.generation !== input.expectedGeneration || + current.environmentId !== environmentId || + current.activeOwnerEpoch !== ownerEpoch + ) { + throw new Error(`Cannot drain stale worker placement for session ${sessionId}`); + } + // Draining closes new admission first. The already-admitted worker may + // finish under its old claim before reconciliation advances ownership. + const values = transitionValues(current, "draining", {}, now()); + const turnClaim = current.turnClaim; + if (turnClaim) { + values.turn_claim_owner = turnClaim.owner; + values.turn_claim_id = turnClaim.claimId; + values.turn_claim_run_id = turnClaim.runId; + values.turn_claim_generation = turnClaim.generation; + values.turn_claim_owner_epoch = turnClaim.ownerEpoch; + } + assertRecordShape({ + state: "draining", + environmentId, + activeOwnerEpoch: ownerEpoch, + workspaceBaseManifestRef: values.workspace_base_manifest_ref, + remoteWorkspaceDir: values.remote_workspace_dir, + workerBundleHash: values.worker_bundle_hash, + lastTranscriptAckCursor: values.last_transcript_ack_cursor, + lastLiveEventAckCursor: values.last_live_event_ack_cursor, + recoveryError: values.recovery_error, + turnClaim, + }); + const result = executeSqliteQuerySync( + db, + query(db) + .updateTable("worker_session_placements") + .set(values) + .where("session_id", "=", sessionId) + .where("state", "=", "active") + .where("transition_generation", "=", current.generation) + .where("environment_id", "=", environmentId) + .where("active_owner_epoch", "=", ownerEpoch), + ); + if (result.numAffectedRows !== 1n) { + throw new Error(`Worker session placement ${sessionId} changed during drain`); + } + return getRequired(db, sessionId); + }); + }, + + startReconcile(input: { + sessionId: string; + environmentId: string; + ownerEpoch: number; + expectedGeneration: number; + }): WorkerSessionPlacementRecord { + const sessionId = required(input.sessionId, "session id"); + const environmentId = required(input.environmentId, "environment id"); + const ownerEpoch = normalizeEpoch(input.ownerEpoch, "active owner epoch"); + const outcome = write((db) => { + const current = getRequired(db, sessionId); + if ( + current.state !== "draining" || + current.generation !== input.expectedGeneration || + current.environmentId !== environmentId || + current.activeOwnerEpoch !== ownerEpoch + ) { + throw new Error(`Cannot reconcile stale worker placement for session ${sessionId}`); + } + // The caller has already fenced the environment. Clear its last claim + // in the same CAS that opens the post-worker reconciliation phase. + const releasedClaim = current.turnClaim !== null; + const values = transitionValues(current, "reconciling", {}, now()); + const update = query(db) + .updateTable("worker_session_placements") + .set(values) + .where("session_id", "=", sessionId) + .where("state", "=", "draining") + .where("transition_generation", "=", current.generation) + .where("environment_id", "=", environmentId) + .where("active_owner_epoch", "=", ownerEpoch); + const guardedUpdate = current.turnClaim + ? update + .where("turn_claim_owner", "=", "worker") + .where("turn_claim_id", "=", current.turnClaim.claimId) + .where("turn_claim_run_id", "=", current.turnClaim.runId) + .where("turn_claim_generation", "=", current.turnClaim.generation) + .where("turn_claim_owner_epoch", "=", current.turnClaim.ownerEpoch) + : update.where("turn_claim_owner", "is", null); + const result = executeSqliteQuerySync(db, guardedUpdate); + if (result.numAffectedRows !== 1n) { + throw new Error(`Worker session placement ${sessionId} changed during reconcile`); + } + return { record: getRequired(db, sessionId), releasedClaim }; + }); + if (outcome.releasedClaim) { + signalTurnClaimRelease(path, sessionId); + } + return outcome.record; + }, + + validateWorkerOwner(input: { + sessionId: string; + environmentId: string; + ownerEpoch: number; + }): boolean { + const current = find(read(), required(input.sessionId, "session id")); + return ( + current?.state === "active" && + current.environmentId === required(input.environmentId, "environment id") && + current.activeOwnerEpoch === normalizeEpoch(input.ownerEpoch, "active owner epoch") + ); + }, + + fail(input: { + sessionId: string; + recoveryError: string; + expectedGeneration?: number; + }): WorkerSessionPlacementRecord { + const sessionId = required(input.sessionId, "session id"); + const recoveryError = required(input.recoveryError, "recovery error"); + const outcome = write((db) => { + const current = getRequired(db, sessionId); + if ( + input.expectedGeneration !== undefined && + current.generation !== input.expectedGeneration + ) { + throw new Error(`Worker session placement ${sessionId} changed before failure`); + } + if (current.state === "failed") { + const result = executeSqliteQuerySync( + db, + query(db) + .updateTable("worker_session_placements") + .set({ recovery_error: recoveryError, updated_at_ms: now() }) + .where("session_id", "=", sessionId) + .where("state", "=", "failed") + .where("transition_generation", "=", current.generation), + ); + if (result.numAffectedRows !== 1n) { + throw new Error(`Worker session placement ${sessionId} changed during failure update`); + } + return { record: getRequired(db, sessionId), releasedClaim: false }; + } + if (!canTransitionWorkerSessionPlacement(current.state, "failed")) { + throw new Error(`Cannot fail worker session placement from ${current.state}`); + } + const localClaim = current.turnClaim?.owner === "local" ? current.turnClaim : null; + const updatedAtMs = now(); + const result = executeSqliteQuerySync( + db, + query(db) + .updateTable("worker_session_placements") + .set({ + state: "failed", + transition_generation: nextGeneration(current.generation), + recovery_error: recoveryError, + turn_claim_owner: localClaim ? "local" : null, + turn_claim_id: localClaim?.claimId ?? null, + turn_claim_run_id: localClaim?.runId ?? null, + turn_claim_generation: localClaim?.generation ?? null, + turn_claim_owner_epoch: null, + updated_at_ms: updatedAtMs, + state_changed_at_ms: updatedAtMs, + }) + .where("session_id", "=", sessionId) + .where("state", "=", current.state) + .where("transition_generation", "=", current.generation), + ); + if (result.numAffectedRows !== 1n) { + throw new Error(`Worker session placement ${sessionId} changed during failure`); + } + return { + record: getRequired(db, sessionId), + releasedClaim: current.turnClaim?.owner === "worker", + }; + }); + if (outcome.releasedClaim) { + signalTurnClaimRelease(path, sessionId); + } + return outcome.record; + }, + + adoptActive(input: { + sessionId: string; + environmentId: string; + ownerEpoch: number; + expectedGeneration?: number; + }): WorkerSessionPlacementRecord { + const sessionId = required(input.sessionId, "session id"); + const environmentId = required(input.environmentId, "environment id"); + const ownerEpoch = normalizeEpoch(input.ownerEpoch, "active owner epoch"); + const current = getRequired(read(), sessionId); + if ( + current.state !== "active" || + current.environmentId !== environmentId || + current.activeOwnerEpoch !== ownerEpoch || + (input.expectedGeneration !== undefined && current.generation !== input.expectedGeneration) + ) { + throw new Error(`Cannot adopt stale worker placement for session ${sessionId}`); + } + return current; + }, + + listForReconcile(): WorkerSessionPlacementRecord[] { + const db = read(); + return executeSqliteQuerySync( + db, + query(db) + .selectFrom("worker_session_placements") + .selectAll() + .where("state", "not in", ["local", "reclaimed"]) + .orderBy("updated_at_ms") + .orderBy("session_id"), + ).rows.map(fromRow); + }, + + list(): WorkerSessionPlacementRecord[] { + const db = read(); + return executeSqliteQuerySync( + db, + query(db).selectFrom("worker_session_placements").selectAll().orderBy("session_id"), + ).rows.map(fromRow); + }, + }; +} + +export type WorkerSessionPlacementStore = ReturnType; diff --git a/src/gateway/worker-environments/placement-turn-claims.ts b/src/gateway/worker-environments/placement-turn-claims.ts new file mode 100644 index 000000000000..cc08f5a5816b --- /dev/null +++ b/src/gateway/worker-environments/placement-turn-claims.ts @@ -0,0 +1,352 @@ +import type { DatabaseSync } from "node:sqlite"; +import { executeSqliteQuerySync } from "../../infra/kysely-sync.js"; +import { + advanceCursor, + normalizeEpoch, + normalizeIdentity, + required, + type WorkerSessionPlacementIdentity, + type WorkerSessionPlacementRecord, + type WorkerSessionTurnClaim, + type WorkerSessionTurnOwner, +} from "./placement-record.js"; +import { ensureLocal, find, getRequired, query } from "./placement-row-codec.js"; + +/** One database path backs many store instances; waits and releases meet here. */ +export type PlacementStoreRuntime = { + path: string; + now: () => number; + read: () => DatabaseSync; + write: (operation: (db: DatabaseSync) => T) => T; +}; + +type TurnClaimReleaseWaiter = () => void; +const turnClaimReleaseWaiters = new Map>>(); + +function waitersFor(path: string, sessionId: string): Set { + let bySession = turnClaimReleaseWaiters.get(path); + if (!bySession) { + bySession = new Map(); + turnClaimReleaseWaiters.set(path, bySession); + } + let waiters = bySession.get(sessionId); + if (!waiters) { + waiters = new Set(); + bySession.set(sessionId, waiters); + } + return waiters; +} + +export function signalTurnClaimRelease(path: string, sessionId: string): void { + const bySession = turnClaimReleaseWaiters.get(path); + const waiters = bySession?.get(sessionId); + if (!waiters) { + return; + } + bySession?.delete(sessionId); + if (bySession?.size === 0) { + turnClaimReleaseWaiters.delete(path); + } + for (const resolve of waiters) { + resolve(); + } +} + +export function createPlacementTurnClaimOps(runtime: PlacementStoreRuntime) { + const { path, now, read, write } = runtime; + + return { + claimTurn( + input: WorkerSessionPlacementIdentity & { + owner: WorkerSessionTurnOwner; + claimId: string; + runId: string; + }, + ): WorkerSessionTurnClaim { + const identity = normalizeIdentity(input); + const claimId = required(input.claimId, "turn claim id"); + const runId = required(input.runId, "turn claim run id"); + const owner: WorkerSessionTurnOwner = + input.owner.kind === "local" + ? { kind: "local" } + : { + kind: "worker", + environmentId: required(input.owner.environmentId, "turn owner environment id"), + ownerEpoch: normalizeEpoch(input.owner.ownerEpoch, "turn owner epoch"), + }; + return write((db) => { + const current = ensureLocal(db, identity, now()); + if (current.turnClaim) { + throw new Error(`Session ${identity.sessionId} already has an active turn claim`); + } + if (owner.kind === "local") { + if (current.state !== "local") { + throw new Error( + `Local turn rejected for session ${identity.sessionId} in placement ${current.state}`, + ); + } + } else if ( + current.state !== "active" || + current.environmentId !== owner.environmentId || + current.activeOwnerEpoch !== owner.ownerEpoch + ) { + throw new Error(`Worker turn rejected for session ${identity.sessionId}: stale owner`); + } + const result = executeSqliteQuerySync( + db, + query(db) + .updateTable("worker_session_placements") + .set({ + turn_claim_owner: owner.kind, + turn_claim_id: claimId, + turn_claim_run_id: runId, + turn_claim_generation: current.generation, + turn_claim_owner_epoch: owner.kind === "worker" ? owner.ownerEpoch : null, + updated_at_ms: now(), + }) + .where("session_id", "=", current.sessionId) + .where("state", "=", current.state) + .where("transition_generation", "=", current.generation) + .where("turn_claim_owner", "is", null), + ); + if (result.numAffectedRows !== 1n) { + throw new Error(`Session ${identity.sessionId} placement changed during turn admission`); + } + return { + sessionId: current.sessionId, + claimId, + runId, + placementGeneration: current.generation, + owner, + }; + }); + }, + + releaseTurn(claim: WorkerSessionTurnClaim): WorkerSessionPlacementRecord { + const sessionId = required(claim.sessionId, "session id"); + const claimId = required(claim.claimId, "turn claim id"); + const runId = required(claim.runId, "turn claim run id"); + const released = write((db) => { + const current = getRequired(db, sessionId); + const persisted = current.turnClaim; + const workerMayFinish = current.state === "active" || current.state === "draining"; + if ( + !persisted || + persisted.claimId !== claimId || + persisted.runId !== runId || + persisted.generation !== claim.placementGeneration || + persisted.owner !== claim.owner.kind || + (claim.owner.kind === "worker" && + (persisted.ownerEpoch !== claim.owner.ownerEpoch || + !workerMayFinish || + current.environmentId !== claim.owner.environmentId || + current.activeOwnerEpoch !== claim.owner.ownerEpoch)) + ) { + throw new Error(`Session ${sessionId} turn claim changed before release`); + } + const result = executeSqliteQuerySync( + db, + query(db) + .updateTable("worker_session_placements") + .set({ + turn_claim_owner: null, + turn_claim_id: null, + turn_claim_run_id: null, + turn_claim_generation: null, + turn_claim_owner_epoch: null, + updated_at_ms: now(), + }) + .where("session_id", "=", sessionId) + .where("turn_claim_id", "=", claimId) + .where("turn_claim_run_id", "=", runId) + .where("turn_claim_generation", "=", claim.placementGeneration), + ); + if (result.numAffectedRows !== 1n) { + throw new Error(`Session ${sessionId} turn claim changed during release`); + } + return getRequired(db, sessionId); + }); + signalTurnClaimRelease(path, sessionId); + return released; + }, + + clearLocalTurnClaimsAfterRestart(): number { + const clearedSessionIds = write((db) => { + const sessionIds = executeSqliteQuerySync( + db, + query(db) + .selectFrom("worker_session_placements") + .select("session_id") + .where("turn_claim_owner", "=", "local"), + ).rows.map((row) => row.session_id); + const result = executeSqliteQuerySync( + db, + query(db) + .updateTable("worker_session_placements") + .set({ + turn_claim_owner: null, + turn_claim_id: null, + turn_claim_run_id: null, + turn_claim_generation: null, + turn_claim_owner_epoch: null, + updated_at_ms: now(), + }) + .where("turn_claim_owner", "=", "local"), + ); + if (result.numAffectedRows !== BigInt(sessionIds.length)) { + throw new Error("Local turn claims changed during restart recovery"); + } + return sessionIds; + }); + for (const sessionId of clearedSessionIds) { + signalTurnClaimRelease(path, sessionId); + } + return clearedSessionIds.length; + }, + + async waitForTurnClaimRelease( + sessionIdInput: string, + waitOptions: { timeoutMs: number; signal?: AbortSignal }, + ): Promise { + const sessionId = required(sessionIdInput, "session id"); + if (!Number.isSafeInteger(waitOptions.timeoutMs) || waitOptions.timeoutMs < 0) { + throw new Error("Worker session turn claim wait timeout must be a non-negative integer"); + } + if (!find(read(), sessionId)?.turnClaim) { + return; + } + await new Promise((resolve, reject) => { + let settled = false; + const waiters = waitersFor(path, sessionId); + const finish = (error?: Error) => { + if (settled) { + return; + } + settled = true; + clearTimeout(timer); + waitOptions.signal?.removeEventListener("abort", onAbort); + waiters.delete(onRelease); + if (waiters.size === 0) { + const bySession = turnClaimReleaseWaiters.get(path); + bySession?.delete(sessionId); + if (bySession?.size === 0) { + turnClaimReleaseWaiters.delete(path); + } + } + if (error) { + reject(error); + } else { + resolve(); + } + }; + const onRelease = () => finish(); + const onAbort = () => finish(new Error(`Turn claim wait aborted for session ${sessionId}`)); + const timer = setTimeout( + () => finish(new Error(`Timed out waiting for session ${sessionId} turn claim release`)), + waitOptions.timeoutMs, + ); + waiters.add(onRelease); + waitOptions.signal?.addEventListener("abort", onAbort, { once: true }); + // Register first, then reread. This closes the release-between-check-and-wait race. + if (!find(read(), sessionId)?.turnClaim) { + finish(); + } else if (waitOptions.signal?.aborted) { + onAbort(); + } + }); + }, + + validateTurnClaim(claim: WorkerSessionTurnClaim): boolean { + const current = find(read(), required(claim.sessionId, "session id")); + const persisted = current?.turnClaim; + return ( + persisted !== undefined && + persisted !== null && + persisted.claimId === claim.claimId && + persisted.runId === claim.runId && + persisted.generation === claim.placementGeneration && + persisted.owner === claim.owner.kind && + (claim.owner.kind === "local" || + (persisted.ownerEpoch === claim.owner.ownerEpoch && + (current?.state === "active" || current?.state === "draining") && + current.environmentId === claim.owner.environmentId && + current.activeOwnerEpoch === claim.owner.ownerEpoch)) + ); + }, + + updateAckCursors(input: { + claim: WorkerSessionTurnClaim; + transcript?: number; + liveEvent?: number; + }): WorkerSessionPlacementRecord { + const sessionId = required(input.claim.sessionId, "session id"); + const claimId = required(input.claim.claimId, "turn claim id"); + const runId = required(input.claim.runId, "turn claim run id"); + if ( + !Number.isSafeInteger(input.claim.placementGeneration) || + input.claim.placementGeneration < 0 + ) { + throw new Error("Worker session placement turn claim generation is invalid"); + } + if (input.claim.owner.kind !== "worker") { + throw new Error("Only a worker turn claim can acknowledge worker cursors"); + } + const placementGeneration = input.claim.placementGeneration; + const environmentId = required(input.claim.owner.environmentId, "environment id"); + const ownerEpoch = normalizeEpoch(input.claim.owner.ownerEpoch, "active owner epoch"); + return write((db) => { + const current = getRequired(db, sessionId); + const persisted = current.turnClaim; + const workerMayFinish = current.state === "active" || current.state === "draining"; + if ( + !workerMayFinish || + current.environmentId !== environmentId || + current.activeOwnerEpoch !== ownerEpoch || + persisted?.owner !== "worker" || + persisted.claimId !== claimId || + persisted.runId !== runId || + persisted.generation !== placementGeneration || + persisted.ownerEpoch !== ownerEpoch + ) { + throw new Error(`Cannot ACK stale worker turn for session ${sessionId}`); + } + // Successful RPC replays can carry an older sequence. Preserve the + // durable high-water mark while acknowledging the idempotent replay. + const transcript = advanceCursor( + current.lastTranscriptAckCursor, + input.transcript, + "transcript ACK cursor", + ); + const liveEvent = advanceCursor( + current.lastLiveEventAckCursor, + input.liveEvent, + "live ACK cursor", + ); + const result = executeSqliteQuerySync( + db, + query(db) + .updateTable("worker_session_placements") + .set({ + last_transcript_ack_cursor: transcript, + last_live_event_ack_cursor: liveEvent, + updated_at_ms: now(), + }) + .where("session_id", "=", sessionId) + .where("state", "=", current.state) + .where("transition_generation", "=", current.generation) + .where("environment_id", "=", environmentId) + .where("active_owner_epoch", "=", ownerEpoch) + .where("turn_claim_owner", "=", "worker") + .where("turn_claim_id", "=", claimId) + .where("turn_claim_run_id", "=", runId) + .where("turn_claim_generation", "=", placementGeneration) + .where("turn_claim_owner_epoch", "=", ownerEpoch), + ); + if (result.numAffectedRows !== 1n) { + throw new Error(`Worker session placement ${sessionId} changed during ACK`); + } + return getRequired(db, sessionId); + }); + }, + }; +} diff --git a/src/gateway/worker-environments/placement-worker-gate.test.ts b/src/gateway/worker-environments/placement-worker-gate.test.ts new file mode 100644 index 000000000000..7963c7d2fc7a --- /dev/null +++ b/src/gateway/worker-environments/placement-worker-gate.test.ts @@ -0,0 +1,153 @@ +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { + closeOpenClawStateDatabaseForTest, + openOpenClawStateDatabase, + type OpenClawStateDatabase, +} from "../../state/openclaw-state-db.js"; +import type { WorkerSessionPlacementIdentity } from "./placement-record.js"; +import { + createWorkerSessionPlacementStore, + type WorkerSessionPlacementStore, +} from "./placement-store.js"; +import { createWorkerSessionPlacementGate } from "./placement-worker-gate.js"; + +const SESSION: WorkerSessionPlacementIdentity = { + sessionId: "session-worker-gate", + agentId: "main", + sessionKey: "agent:main:worker-gate", +}; +const ENVIRONMENT_ID = "environment-worker-gate"; +const OWNER_EPOCH = 7; + +describe("worker session placement gate", () => { + let root: string; + let database: OpenClawStateDatabase; + let store: WorkerSessionPlacementStore; + + beforeEach(async () => { + root = await fs.mkdtemp(path.join(await fs.realpath(os.tmpdir()), "openclaw-worker-gate-")); + database = openOpenClawStateDatabase({ env: { OPENCLAW_STATE_DIR: root } }); + store = createWorkerSessionPlacementStore({ database }); + }); + + afterEach(async () => { + closeOpenClawStateDatabaseForTest(); + await fs.rm(root, { recursive: true, force: true }); + }); + + function activate() { + let placement = store.startDispatch(SESSION); + placement = store.transition({ + sessionId: SESSION.sessionId, + from: "requested", + to: "provisioning", + expectedGeneration: placement.generation, + patch: { environmentId: ENVIRONMENT_ID }, + }); + placement = store.transition({ + sessionId: SESSION.sessionId, + from: "provisioning", + to: "syncing", + expectedGeneration: placement.generation, + patch: { workerBundleHash: "a".repeat(64) }, + }); + placement = store.transition({ + sessionId: SESSION.sessionId, + from: "syncing", + to: "starting", + expectedGeneration: placement.generation, + patch: { + workspaceBaseManifestRef: "manifest-worker-gate", + remoteWorkspaceDir: "/workspace/worker-gate", + }, + }); + return store.transition({ + sessionId: SESSION.sessionId, + from: "starting", + to: "active", + expectedGeneration: placement.generation, + patch: { activeOwnerEpoch: OWNER_EPOCH }, + }); + } + + function preclaim(runId: string) { + const placement = activate(); + return store.claimTurn({ + sessionId: placement.sessionId, + agentId: placement.agentId, + sessionKey: placement.sessionKey, + claimId: `claim:${runId}`, + runId, + owner: { kind: "worker", environmentId: ENVIRONMENT_ID, ownerEpoch: OWNER_EPOCH }, + }); + } + + it("accepts only the exact gateway-preclaimed worker run", () => { + const runId = "run-worker-gate"; + preclaim(runId); + const gate = createWorkerSessionPlacementGate(store); + const binding = { + sessionId: SESSION.sessionId, + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + runId, + }; + + expect(gate.validateWorkerTurn(binding)).toBe(true); + expect(gate.validateWorkerTurn({ ...binding, runId: "run-competing" })).toBe(false); + expect(gate.validateWorkerTurn({ ...binding, ownerEpoch: OWNER_EPOCH + 1 })).toBe(false); + }); + + it("updates exact-owner cursors and rejects stale descriptor replay", () => { + const runId = "run-worker-ack"; + const claim = preclaim(runId); + const gate = createWorkerSessionPlacementGate(store); + const binding = { + sessionId: SESSION.sessionId, + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + runId, + }; + + gate.updateAckCursors({ ...binding, transcriptSeq: 4, liveSeq: 9 }); + expect(store.get(SESSION.sessionId)).toMatchObject({ + generation: claim.placementGeneration, + lastTranscriptAckCursor: 4, + lastLiveEventAckCursor: 9, + }); + store.releaseTurn(claim); + expect(store.get(SESSION.sessionId)?.turnClaim).toBeNull(); + expect(gate.validateWorkerTurn(binding)).toBe(false); + }); + + it("lets the admitted worker finish acknowledgements after draining closes admission", () => { + const runId = "run-worker-draining-ack"; + const claim = preclaim(runId); + const active = store.get(SESSION.sessionId); + if (active?.state !== "active") { + throw new Error("expected active placement"); + } + store.startDrain({ + sessionId: SESSION.sessionId, + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + expectedGeneration: active.generation, + }); + const gate = createWorkerSessionPlacementGate(store); + const binding = { + sessionId: SESSION.sessionId, + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + runId, + }; + + expect(gate.validateWorkerTurn(binding)).toBe(true); + gate.updateAckCursors({ ...binding, transcriptSeq: 5 }); + expect(store.get(SESSION.sessionId)?.lastTranscriptAckCursor).toBe(5); + store.releaseTurn(claim); + expect(gate.validateWorkerTurn(binding)).toBe(false); + }); +}); diff --git a/src/gateway/worker-environments/placement-worker-gate.ts b/src/gateway/worker-environments/placement-worker-gate.ts new file mode 100644 index 000000000000..0e782a3a105a --- /dev/null +++ b/src/gateway/worker-environments/placement-worker-gate.ts @@ -0,0 +1,60 @@ +import type { + WorkerSessionPlacementRecord, + WorkerSessionPlacementStore, + WorkerSessionTurnClaim, +} from "./placement-store.js"; +import type { WorkerPlacementTurnBinding, WorkerSessionPlacementGate } from "./service.js"; + +function claimForBinding( + record: WorkerSessionPlacementRecord | undefined, + binding: WorkerPlacementTurnBinding, +): WorkerSessionTurnClaim | undefined { + const persisted = record?.turnClaim; + if ( + !record || + (record.state !== "active" && record.state !== "draining") || + record.environmentId !== binding.environmentId || + record.activeOwnerEpoch !== binding.ownerEpoch || + persisted?.owner !== "worker" || + persisted.runId !== binding.runId || + persisted.ownerEpoch !== binding.ownerEpoch + ) { + return undefined; + } + return { + sessionId: binding.sessionId, + claimId: persisted.claimId, + runId: persisted.runId, + placementGeneration: persisted.generation, + owner: { + kind: "worker", + environmentId: binding.environmentId, + ownerEpoch: binding.ownerEpoch, + }, + }; +} + +export function createWorkerSessionPlacementGate( + store: WorkerSessionPlacementStore, +): WorkerSessionPlacementGate { + const validateWorkerTurn = (binding: WorkerPlacementTurnBinding): boolean => { + const claim = claimForBinding(store.get(binding.sessionId), binding); + return claim ? store.validateTurnClaim(claim) : false; + }; + + return { + validateWorkerTurn, + + updateAckCursors(binding): void { + const claim = claimForBinding(store.get(binding.sessionId), binding); + if (!claim) { + throw new Error(`Cannot ACK stale worker turn for session ${binding.sessionId}`); + } + store.updateAckCursors({ + claim, + ...(binding.transcriptSeq === undefined ? {} : { transcript: binding.transcriptSeq }), + ...(binding.liveSeq === undefined ? {} : { liveEvent: binding.liveSeq }), + }); + }, + }; +} diff --git a/src/gateway/worker-environments/service-contract.ts b/src/gateway/worker-environments/service-contract.ts index 367c220d4477..f7383f014c54 100644 --- a/src/gateway/worker-environments/service-contract.ts +++ b/src/gateway/worker-environments/service-contract.ts @@ -1,3 +1,4 @@ +import type { WorkerSessionPlacementRecord } from "./placement-record.js"; import type { WorkerEnvironmentState } from "./state.js"; import type { WorkerTunnelHandle, @@ -27,3 +28,18 @@ export type WorkerEnvironmentServiceContract = { startTunnel(request: WorkerTunnelRequest): Promise; stopTunnel(environmentId: string, ownerEpoch?: number): Promise; }; + +export type WorkerPlacementDispatchRequest = { + sessionId: string; + sessionKey: string; + agentId: string; + profileId: string; +}; + +// Leaf dispatch contract: GatewayRequestContext must not import the dispatch +// runtime (it reaches agents/plugins and closes an import cycle through core). +export type WorkerPlacementDispatchContract = { + dispatch( + request: WorkerPlacementDispatchRequest, + ): Promise>; +}; diff --git a/src/gateway/worker-environments/service.test.ts b/src/gateway/worker-environments/service.test.ts index 0b0e1f44867a..b7425abdc71a 100644 --- a/src/gateway/worker-environments/service.test.ts +++ b/src/gateway/worker-environments/service.test.ts @@ -129,6 +129,7 @@ describe("worker environment service", () => { | "tunnelManager" | "generateWorkerCredential" | "liveEvents" + | "placementStore" | "workerCredentialTtlMs" > > = {}, @@ -243,6 +244,7 @@ describe("worker environment service", () => { environmentId, credential: [CREDENTIAL, environmentId].join("-"), sessionId: null, + runId: null, ownerEpoch: 1, rpcSetVersion: 1, handshake: BOOTSTRAP_RECEIPT, @@ -269,6 +271,7 @@ describe("worker environment service", () => { credentialHash: credential.credentialHash, bundleHash: credential.bundleHash, sessionId, + runId: "run-1", ownerEpoch: attached.ownerEpoch, rpcSetVersion: credential.rpcSetVersion, protocolFeatures: [...attached.bootstrapReceipt.protocolFeatures], @@ -282,7 +285,7 @@ describe("worker environment service", () => { return { runEpoch: identity.ownerEpoch, sessionId: identity.sessionId ?? "session-missing", - runId: "run-inference", + runId: identity.runId ?? "run-missing", turnId: "turn-inference", modelRef: { provider: "fake", model: "model-test" }, context: { messages: [] }, @@ -435,6 +438,551 @@ describe("worker environment service", () => { expect(applyTranscriptCommit).toHaveBeenCalledOnce(); }); + it("admits only a gateway-preclaimed worker placement and fences later requests", async () => { + const environmentId = "worker-placement-fence"; + const sessionId = "session-placement-fence"; + const identity = seedAttachedIdentity(environmentId, sessionId); + const placementStore = { + validateWorkerTurn: vi.fn(() => true), + updateAckCursors: vi.fn(), + }; + const workerService = createService(createProvider(), { placementStore }); + const admission = { + environmentId, + credential: [CREDENTIAL, environmentId, sessionId].join("-"), + sessionId, + runId: "run-1", + ownerEpoch: identity.ownerEpoch, + rpcSetVersion: 1, + handshake: BOOTSTRAP_RECEIPT, + }; + + await expect(workerService.admitWorker(admission)).resolves.toMatchObject({ ok: true }); + await expect(workerService.admitWorker(admission)).resolves.toMatchObject({ ok: true }); + expect(placementStore.validateWorkerTurn).toHaveBeenLastCalledWith({ + sessionId, + environmentId, + ownerEpoch: identity.ownerEpoch, + runId: "run-1", + }); + expect(placementStore.validateWorkerTurn).toHaveBeenCalledTimes(2); + expect(workerService.validateWorkerConnection(identity)).toBeNull(); + + const warmEnvironmentId = "worker-placement-warm"; + seedReady(warmEnvironmentId); + const warmAdmission = await workerService.admitWorker(admissionFor(warmEnvironmentId)); + expect(warmAdmission).toMatchObject({ ok: true }); + if (!warmAdmission.ok) { + throw new Error("warm worker admission failed"); + } + expect(workerService.validateWorkerConnection(warmAdmission.identity)).toBeNull(); + expect(placementStore.validateWorkerTurn).toHaveBeenCalledTimes(3); + + placementStore.validateWorkerTurn.mockReturnValue(false); + await expect( + workerService.admitWorker({ ...admission, runId: "run-conflict" }), + ).resolves.toEqual({ ok: false, reason: "placement-mismatch" }); + + placementStore.validateWorkerTurn.mockReturnValue(true); + nowMs += 10_000; + expect(workerService.validateWorkerConnection(identity)).toBeNull(); + expect(workerService.validateWorkerConnection(warmAdmission.identity)).toBe( + "credential-expired", + ); + await expect(workerService.admitWorker(admission)).resolves.toEqual({ + ok: false, + reason: "credential-expired", + }); + + placementStore.validateWorkerTurn.mockReturnValue(false); + expect(workerService.validateWorkerConnection(identity)).toBe("placement-mismatch"); + await expect( + workerService.commitTranscript(identity, { + runEpoch: identity.ownerEpoch, + seq: 1, + baseLeafId: null, + messages: [ + { + role: "user", + content: [{ type: "text", text: "fenced" }], + timestamp: 1, + }, + ], + }), + ).resolves.toEqual({ ok: false, closeReason: "placement-mismatch" }); + }); + + it("persists worker transcript and terminal live ACK cursors", async () => { + const identity = seedAttachedIdentity("worker-placement-ack", "session-placement-ack"); + const placementStore = { + validateWorkerTurn: vi.fn(() => true), + updateAckCursors: vi.fn(), + }; + const applyTranscriptCommit = vi.fn(async () => ({ + ok: true as const, + result: { entryIds: ["entry-placement"], newLeafId: "entry-placement" }, + })); + const liveEvents = createLiveEvents({ + apply: vi.fn( + ({ + request, + }: Parameters["apply"]>[0]) => ({ + ok: true as const, + result: { ackedSeq: request.seq }, + }), + ), + }); + const workerService = createService(createProvider(), { + applyTranscriptCommit, + liveEvents, + placementStore, + }); + const binding = { + sessionId: identity.sessionId ?? "session-missing", + environmentId: identity.environmentId, + ownerEpoch: identity.ownerEpoch, + runId: identity.runId ?? "run-missing", + }; + + await expect( + workerService.commitTranscript(identity, { + runEpoch: identity.ownerEpoch, + seq: 7, + baseLeafId: null, + messages: [ + { + role: "user", + content: [{ type: "text", text: "commit" }], + timestamp: 1, + }, + ], + }), + ).resolves.toMatchObject({ ok: true }); + expect(placementStore.updateAckCursors).toHaveBeenCalledWith({ + ...binding, + transcriptSeq: 7, + }); + + await expect( + workerService.pushLiveEvent(identity, { + runEpoch: identity.ownerEpoch, + lastAckedSeq: 0, + seq: 1, + runId: binding.runId, + event: { kind: "lifecycle", payload: { phase: "end", endedAt: 2 } }, + }), + ).resolves.toEqual({ ok: true, result: { ackedSeq: 1 } }); + expect(placementStore.updateAckCursors).toHaveBeenLastCalledWith({ + ...binding, + liveSeq: 1, + }); + }); + + it("does not ACK a transcript commit after its worker claim is fenced", async () => { + const identity = seedAttachedIdentity("worker-placement-race", "session-placement-race"); + const placementStore = { + validateWorkerTurn: vi.fn(() => true), + updateAckCursors: vi.fn(), + }; + let finishCommit: (() => void) | undefined; + const commitBlocked = new Promise((resolve) => { + finishCommit = resolve; + }); + const applyTranscriptCommit = vi.fn(async () => { + await commitBlocked; + return { + ok: true as const, + result: { entryIds: ["entry-placement-race"], newLeafId: "entry-placement-race" }, + }; + }); + const workerService = createService(createProvider(), { + applyTranscriptCommit, + placementStore, + }); + + const commit = workerService.commitTranscript(identity, { + runEpoch: identity.ownerEpoch, + seq: 1, + baseLeafId: null, + messages: [ + { + role: "user", + content: [{ type: "text", text: "commit before claim fence" }], + timestamp: 1, + }, + ], + }); + await vi.waitFor(() => expect(applyTranscriptCommit).toHaveBeenCalledOnce()); + placementStore.validateWorkerTurn.mockReturnValue(false); + finishCommit?.(); + + await expect(commit).resolves.toEqual({ ok: false, closeReason: "placement-mismatch" }); + expect(placementStore.validateWorkerTurn).toHaveBeenCalledTimes(2); + expect(placementStore.updateAckCursors).not.toHaveBeenCalled(); + }); + + it("advances the transcript cursor when a stale-base commit consumes its sequence", async () => { + const identity = seedAttachedIdentity("worker-placement-stale", "session-placement-stale"); + const placementStore = { + validateWorkerTurn: vi.fn(() => true), + updateAckCursors: vi.fn(), + }; + const applyTranscriptCommit = vi + .fn>() + .mockResolvedValueOnce({ ok: false, reason: "stale-base-leaf" }) + .mockResolvedValueOnce({ ok: false, reason: "invalid-batch" }); + const workerService = createService(createProvider(), { + applyTranscriptCommit, + placementStore, + }); + const request = { + runEpoch: identity.ownerEpoch, + seq: 11, + baseLeafId: "stale-leaf", + messages: [ + { + role: "user" as const, + content: [{ type: "text" as const, text: "stale commit" }], + timestamp: 1, + }, + ], + }; + + await expect(workerService.commitTranscript(identity, request)).resolves.toEqual({ + ok: false, + reason: "stale-base-leaf", + }); + expect(placementStore.updateAckCursors).toHaveBeenCalledWith({ + sessionId: identity.sessionId, + environmentId: identity.environmentId, + ownerEpoch: identity.ownerEpoch, + runId: identity.runId, + transcriptSeq: 11, + }); + + await expect( + workerService.commitTranscript(identity, { ...request, seq: 12 }), + ).resolves.toEqual({ ok: false, reason: "invalid-batch" }); + expect(placementStore.updateAckCursors).toHaveBeenCalledOnce(); + }); + + it("fences after a buffered terminal event becomes acknowledged by a gap fill", async () => { + const identity = seedAttachedIdentity("worker-placement-gap", "session-placement-gap"); + const placementStore = { + validateWorkerTurn: vi.fn(() => true), + updateAckCursors: vi.fn(), + }; + const applyTranscriptCommit = vi.fn(async () => ({ + ok: true as const, + result: { entryIds: ["entry-after-terminal-gap"], newLeafId: "entry-after-terminal-gap" }, + })); + const liveApply = vi.fn( + ({ + request, + }: Parameters["apply"]>[0]) => ({ + ok: true as const, + result: { ackedSeq: request.seq === 1 ? 2 : 0 }, + }), + ); + const workerService = createService(createProvider(), { + applyTranscriptCommit, + liveEvents: createLiveEvents({ apply: liveApply }), + placementStore, + }); + + await expect( + workerService.pushLiveEvent(identity, { + runEpoch: identity.ownerEpoch, + lastAckedSeq: 0, + seq: 2, + runId: identity.runId ?? "run-missing", + event: { kind: "lifecycle", payload: { phase: "end", endedAt: 2 } }, + }), + ).resolves.toEqual({ ok: true, result: { ackedSeq: 0 } }); + expect(placementStore.updateAckCursors).toHaveBeenCalledWith({ + sessionId: identity.sessionId, + environmentId: identity.environmentId, + ownerEpoch: identity.ownerEpoch, + runId: identity.runId, + liveSeq: 0, + }); + + await expect( + workerService.pushLiveEvent(identity, { + runEpoch: identity.ownerEpoch, + lastAckedSeq: 0, + seq: 1, + runId: identity.runId ?? "run-missing", + event: { kind: "assistant", payload: { text: "fills gap", delta: "fills gap" } }, + }), + ).resolves.toEqual({ ok: true, result: { ackedSeq: 2 } }); + await expect( + workerService.commitTranscript(identity, { + runEpoch: identity.ownerEpoch, + seq: 1, + baseLeafId: null, + messages: [ + { + role: "user", + content: [{ type: "text", text: "late transcript" }], + timestamp: 1, + }, + ], + }), + ).resolves.toEqual({ ok: false, closeReason: "placement-mismatch" }); + await expect( + workerService.pushLiveEvent(identity, { + runEpoch: identity.ownerEpoch, + lastAckedSeq: 2, + seq: 3, + runId: identity.runId ?? "run-missing", + event: { kind: "assistant", payload: { text: "late", delta: "late" } }, + }), + ).resolves.toEqual({ ok: false, closeReason: "placement-mismatch" }); + expect(applyTranscriptCommit).not.toHaveBeenCalled(); + expect(liveApply).toHaveBeenCalledTimes(2); + }); + + it("applies a terminal ACK only after its transcript commit finishes", async () => { + const identity = seedAttachedIdentity("worker-placement-order", "session-placement-order"); + const placementStore = { + validateWorkerTurn: vi.fn(() => true), + updateAckCursors: vi.fn(), + }; + let finishCommit: (() => void) | undefined; + const commitBlocked = new Promise((resolve) => { + finishCommit = resolve; + }); + const applyTranscriptCommit = vi.fn(async () => { + await commitBlocked; + return { + ok: true as const, + result: { entryIds: ["entry-order"], newLeafId: "entry-order" }, + }; + }); + const workerService = createService(createProvider(), { + applyTranscriptCommit, + liveEvents: createLiveEvents({ + apply: vi.fn( + ({ + request, + }: Parameters< + NonNullable["apply"] + >[0]) => ({ ok: true as const, result: { ackedSeq: request.seq } }), + ), + }), + placementStore, + }); + + const commit = workerService.commitTranscript(identity, { + runEpoch: identity.ownerEpoch, + seq: 1, + baseLeafId: null, + messages: [ + { + role: "user", + content: [{ type: "text", text: "commit before terminal" }], + timestamp: 1, + }, + ], + }); + await vi.waitFor(() => expect(applyTranscriptCommit).toHaveBeenCalledOnce()); + const terminal = workerService.pushLiveEvent(identity, { + runEpoch: identity.ownerEpoch, + lastAckedSeq: 0, + seq: 1, + runId: identity.runId ?? "run-missing", + event: { kind: "lifecycle", payload: { phase: "end", endedAt: 2 } }, + }); + await Promise.resolve(); + expect(placementStore.updateAckCursors).not.toHaveBeenCalled(); + + finishCommit?.(); + await expect(commit).resolves.toMatchObject({ ok: true }); + await expect(terminal).resolves.toEqual({ ok: true, result: { ackedSeq: 1 } }); + expect(placementStore.updateAckCursors.mock.calls).toEqual([ + [ + { + sessionId: identity.sessionId, + environmentId: identity.environmentId, + ownerEpoch: identity.ownerEpoch, + runId: identity.runId, + transcriptSeq: 1, + }, + ], + [ + { + sessionId: identity.sessionId, + environmentId: identity.environmentId, + ownerEpoch: identity.ownerEpoch, + runId: identity.runId, + liveSeq: 1, + }, + ], + ]); + }); + + it("fences post-terminal mutations while preserving sequenced replays", async () => { + const identity = seedAttachedIdentity("worker-terminal-fence", "session-terminal-fence"); + const placementStore = { + validateWorkerTurn: vi.fn(() => true), + updateAckCursors: vi.fn(), + }; + const applyTranscriptCommit = vi.fn(async () => ({ + ok: true as const, + result: { entryIds: ["entry-terminal"], newLeafId: "entry-terminal" }, + })); + const liveApply = vi.fn( + ({ + request, + }: Parameters["apply"]>[0]) => ({ + ok: true as const, + result: { ackedSeq: request.seq }, + }), + ); + const executeInference = vi.fn( + async () => ({ + type: "error", + reason: "provider-error", + message: "Provider request failed", + }), + ); + const workerService = createService(createProvider(), { + applyTranscriptCommit, + executeInference, + liveEvents: createLiveEvents({ apply: liveApply }), + placementStore, + }); + const transcript = { + runEpoch: identity.ownerEpoch, + seq: 1, + baseLeafId: null, + messages: [ + { + role: "user" as const, + content: [{ type: "text" as const, text: "terminal fence" }], + timestamp: 1, + }, + ], + }; + const terminal = { + runEpoch: identity.ownerEpoch, + lastAckedSeq: 0, + seq: 1, + runId: identity.runId ?? "run-missing", + event: { kind: "lifecycle" as const, payload: { phase: "end" as const, endedAt: 2 } }, + }; + + await expect(workerService.commitTranscript(identity, transcript)).resolves.toMatchObject({ + ok: true, + }); + await expect(workerService.pushLiveEvent(identity, terminal)).resolves.toEqual({ + ok: true, + result: { ackedSeq: 1 }, + }); + + await expect(workerService.commitTranscript(identity, transcript)).resolves.toMatchObject({ + ok: true, + }); + await expect( + workerService.commitTranscript(identity, { ...transcript, seq: 2 }), + ).resolves.toEqual({ ok: false, closeReason: "placement-mismatch" }); + expect(applyTranscriptCommit).toHaveBeenCalledTimes(2); + + await expect(workerService.pushLiveEvent(identity, terminal)).resolves.toEqual({ + ok: true, + result: { ackedSeq: 1 }, + }); + await expect( + workerService.pushLiveEvent(identity, { + ...terminal, + seq: 2, + event: { kind: "assistant", payload: { text: "late", delta: "late" } }, + }), + ).resolves.toEqual({ ok: false, closeReason: "placement-mismatch" }); + expect(liveApply).toHaveBeenCalledTimes(2); + + expect( + workerService.startInference(identity, inferenceRequest(identity), { + connectionId: "connection-terminal-fence", + send: vi.fn(), + }), + ).toEqual({ ok: false, closeReason: "placement-mismatch" }); + expect(workerService.cancelInference(identity, inferenceRequest(identity))).toEqual({ + ok: false, + closeReason: "placement-mismatch", + }); + expect(executeInference).not.toHaveBeenCalled(); + + const rotatedCredentialHash = hashWorkerCredential( + ["rotated", identity.environmentId, identity.sessionId].join("-"), + ); + database.db + .prepare( + "UPDATE worker_environment_credentials SET credential_hash = ? WHERE environment_id = ?", + ) + .run(rotatedCredentialHash, identity.environmentId); + const rotatedIdentity = { ...identity, credentialHash: rotatedCredentialHash }; + await expect( + workerService.commitTranscript(rotatedIdentity, { ...transcript, seq: 2 }), + ).resolves.toMatchObject({ ok: true }); + expect(applyTranscriptCommit).toHaveBeenCalledTimes(3); + }); + + it("does not treat a terminal event on an already ACKed sequence as authoritative", async () => { + const identity = seedAttachedIdentity("worker-terminal-reuse", "session-terminal-reuse"); + const applyTranscriptCommit = vi.fn(async () => ({ + ok: true as const, + result: { entryIds: ["entry-after-reuse"], newLeafId: "entry-after-reuse" }, + })); + const workerService = createService(createProvider(), { + applyTranscriptCommit, + liveEvents: createLiveEvents({ + apply: vi.fn( + ({ + request, + }: Parameters< + NonNullable["apply"] + >[0]) => ({ ok: true as const, result: { ackedSeq: request.seq } }), + ), + }), + placementStore: { + validateWorkerTurn: vi.fn(() => true), + updateAckCursors: vi.fn(), + }, + }); + const event = { + runEpoch: identity.ownerEpoch, + lastAckedSeq: 0, + seq: 1, + runId: identity.runId ?? "run-missing", + event: { kind: "assistant" as const, payload: { text: "first", delta: "first" } }, + }; + + await expect(workerService.pushLiveEvent(identity, event)).resolves.toMatchObject({ ok: true }); + await expect( + workerService.pushLiveEvent(identity, { + ...event, + event: { kind: "lifecycle", payload: { phase: "end", endedAt: 2 } }, + }), + ).resolves.toMatchObject({ ok: true }); + await expect( + workerService.commitTranscript(identity, { + runEpoch: identity.ownerEpoch, + seq: 1, + baseLeafId: null, + messages: [ + { + role: "user", + content: [{ type: "text", text: "still mutable" }], + timestamp: 1, + }, + ], + }), + ).resolves.toMatchObject({ ok: true }); + expect(applyTranscriptCommit).toHaveBeenCalledOnce(); + }); + it("fences inference by epoch and the durable session credential", async () => { const identity = seedAttachedIdentity("worker-inference-fence", "session-inference-fence"); const executeInference = vi.fn( @@ -1526,6 +2074,7 @@ describe("worker environment service", () => { ownerEpoch: request.ownerEpoch, remoteSocketPath: "/tmp/worker/gateway.sock", runWorkspaceCommand: vi.fn(), + syncWorkspace: vi.fn(), stop: async () => {}, }; }), diff --git a/src/gateway/worker-environments/service.ts b/src/gateway/worker-environments/service.ts index 754c3ed8421e..689eed7c77fc 100644 --- a/src/gateway/worker-environments/service.ts +++ b/src/gateway/worker-environments/service.ts @@ -93,6 +93,15 @@ const serviceError = (code: WorkerEnvironmentServiceErrorCode, message: string) new WorkerEnvironmentServiceError(code, message); const ORPHANED_LEASE_ERROR = "Worker provider no longer recognizes the lease"; +function workerEnvironmentIdempotencyDigest(idempotencyKey: string): string { + return createHash("sha256").update(idempotencyKey).digest("hex"); +} + +export function workerEnvironmentIdForIdempotencyKey(idempotencyKey: string): string { + const digest = workerEnvironmentIdempotencyDigest(idempotencyKey); + return `worker:${digest.slice(0, 32)}`; +} + type WorkerEnvironmentServiceOptions = { store: WorkerEnvironmentStore; getConfig: () => OpenClawConfig; @@ -131,6 +140,42 @@ type WorkerEnvironmentServiceOptions = { >; executeInference: WorkerInferenceExecutor; inferenceStore?: WorkerInferenceStore; + placementStore?: WorkerSessionPlacementGate; +}; + +export type WorkerPlacementTurnBinding = Readonly<{ + sessionId: string; + environmentId: string; + ownerEpoch: number; + runId: string; +}>; + +type WorkerProcessTurnBinding = WorkerPlacementTurnBinding & { + credentialHash: string; +}; + +type WorkerTerminalTurnFence = WorkerProcessTurnBinding & { + transcriptSeq: number; + liveSeq: number; +}; + +type WorkerPendingTerminalTurnFence = WorkerProcessTurnBinding & { + terminalLiveSeq: number; +}; + +type WorkerTurnRequest = + | { kind: "inference" } + | { kind: "live"; seq: number } + | { kind: "transcript"; seq: number }; + +export type WorkerSessionPlacementGate = { + validateWorkerTurn(binding: WorkerPlacementTurnBinding): boolean; + updateAckCursors( + binding: WorkerPlacementTurnBinding & { + transcriptSeq?: number; + liveSeq?: number; + }, + ): void; }; type WorkerTranscriptCommitApplicationResult = @@ -211,6 +256,9 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService const operations = new KeyedAsyncQueue(); const activeOperations = new Set>(); const pendingCredentials = new Map(); + const observedAckCursors = new Map(); + const pendingTerminalTurnFences = new Map(); + const terminalTurnFences = new Map(); const now = options.now ?? Date.now; const inference = createWorkerInferenceManager({ execute: options.executeInference, @@ -223,6 +271,86 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService let unsubscribeSessionIdentityMutation: (() => void) | undefined; let stopping = false; + const placementBinding = ( + identity: WorkerConnectionIdentity, + ): WorkerPlacementTurnBinding | undefined => { + if (!identity.sessionId || !identity.runId) { + return undefined; + } + return { + sessionId: identity.sessionId, + environmentId: identity.environmentId, + ownerEpoch: identity.ownerEpoch, + runId: identity.runId, + }; + }; + + const processTurnBinding = ( + identity: WorkerConnectionIdentity, + ): WorkerProcessTurnBinding | undefined => { + const placement = placementBinding(identity); + return placement ? { ...placement, credentialHash: identity.credentialHash } : undefined; + }; + + const matchesTurnBinding = ( + left: WorkerProcessTurnBinding, + right: WorkerProcessTurnBinding, + ): boolean => + left.sessionId === right.sessionId && + left.environmentId === right.environmentId && + left.ownerEpoch === right.ownerEpoch && + left.runId === right.runId && + safeEqualSecret(left.credentialHash, right.credentialHash); + + const recordAckCursor = ( + binding: WorkerProcessTurnBinding, + cursor: { transcriptSeq: number } | { liveSeq: number }, + ): WorkerTerminalTurnFence => { + const current = observedAckCursors.get(binding.sessionId); + const currentTurn = current && matchesTurnBinding(current, binding) ? current : undefined; + const next: WorkerTerminalTurnFence = { + ...binding, + transcriptSeq: + "transcriptSeq" in cursor + ? Math.max(currentTurn?.transcriptSeq ?? 0, cursor.transcriptSeq) + : (currentTurn?.transcriptSeq ?? 0), + liveSeq: + "liveSeq" in cursor + ? Math.max(currentTurn?.liveSeq ?? 0, cursor.liveSeq) + : (currentTurn?.liveSeq ?? 0), + }; + observedAckCursors.set(binding.sessionId, next); + return next; + }; + + const observedAckCursorFor = ( + binding: WorkerProcessTurnBinding, + ): WorkerTerminalTurnFence | undefined => { + const observed = observedAckCursors.get(binding.sessionId); + return observed && matchesTurnBinding(observed, binding) ? observed : undefined; + }; + + const validateWorkerPlacement = ( + identity: WorkerConnectionIdentity, + ): { durableClaim: boolean; valid: boolean } => { + if (!options.placementStore) { + return { durableClaim: false, valid: true }; + } + if (identity.sessionId === null && identity.runId === null) { + return { durableClaim: false, valid: true }; + } + const binding = placementBinding(identity); + const valid = binding ? options.placementStore.validateWorkerTurn(binding) : false; + return { durableClaim: valid, valid }; + }; + + const isTerminalLiveEvent = (request: WorkerLiveEventParams): boolean => + request.event.kind === "lifecycle" && + (request.event.payload.phase === "end" || + (request.event.payload.phase === "error" && + (request.event.payload.aborted === true || + request.event.payload.fallbackExhaustedFailure === true))); + const project = (record: WorkerEnvironmentRecord) => ({ ...record, tunnelStatus: tunnels?.status(record.environmentId) ?? ("stopped" as const), @@ -753,7 +881,7 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService if (!normalizedProfileId || normalizedProfileId !== profileId) { throw serviceError("invalid_profile", "Worker profile id must be non-empty and trimmed"); } - const digest = createHash("sha256").update(idempotencyKey).digest("hex"); + const digest = workerEnvironmentIdempotencyDigest(idempotencyKey); const environmentId = `worker:${digest.slice(0, 32)}`; return withLock(environmentId, async () => { if (stopping) { @@ -1038,6 +1166,9 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService await Promise.allSettled(activeOperations); } pendingCredentials.clear(); + observedAckCursors.clear(); + pendingTerminalTurnFences.clear(); + terminalTurnFences.clear(); options.liveEvents?.clear(); }; @@ -1077,6 +1208,7 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService const validateAttachedWorkerRequest = ( identity: WorkerConnectionIdentity, runEpoch: number, + request: WorkerTurnRequest, ): | { ok: true } | { ok: false; closeReason: WorkerProtocolCloseReason } @@ -1084,11 +1216,29 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService if (stopping) { return { ok: false, closeReason: "environment-unavailable" }; } + const placement = validateWorkerPlacement(identity); + if (!placement.valid) { + return { ok: false, closeReason: "placement-mismatch" }; + } + const turnBinding = processTurnBinding(identity); + const terminalFence = identity.sessionId + ? terminalTurnFences.get(identity.sessionId) + : undefined; + if (turnBinding && terminalFence && matchesTurnBinding(terminalFence, turnBinding)) { + const isReplay = + (request.kind === "transcript" && request.seq <= terminalFence.transcriptSeq) || + (request.kind === "live" && request.seq <= terminalFence.liveSeq); + if (!isReplay) { + return { ok: false, closeReason: "placement-mismatch" }; + } + } const credential = store.getCredential(identity.environmentId); if (!credential || !safeEqualSecret(credential.credentialHash, identity.credentialHash)) { return { ok: false, closeReason: "credential-replaced" }; } - if (now() >= credential.expiresAtMs) { + // TTL limits admission and reconnect. An already-admitted exact durable + // turn stays usable until its terminal ACK or placement fence. + if (now() >= credential.expiresAtMs && !placement.durableClaim) { return { ok: false, closeReason: "credential-expired" }; } const environment = store.get(identity.environmentId); @@ -1111,6 +1261,11 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService ) { return { ok: false, reason: "session-not-attached" }; } + if (turnBinding && terminalFence && !matchesTurnBinding(terminalFence, turnBinding)) { + // Credential rotation identifies a new process turn even when a caller + // intentionally reuses its durable run id (for example, cron sessions). + terminalTurnFences.delete(turnBinding.sessionId); + } return { ok: true }; }; @@ -1119,32 +1274,118 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService request: WorkerTranscriptCommitParams, ): Promise => withLock(identity.environmentId, async () => { - const binding = validateAttachedWorkerRequest(identity, request.runEpoch); + const binding = validateAttachedWorkerRequest(identity, request.runEpoch, { + kind: "transcript", + seq: request.seq, + }); if (!binding.ok) { return binding; } if (!options.applyTranscriptCommit) { return { ok: false, closeReason: "gateway-unavailable" }; } - return await options.applyTranscriptCommit({ identity, request }); + const result = await options.applyTranscriptCommit({ identity, request }); + // Transcript persistence awaits outside the placement transaction. Revalidate the durable + // claim before exposing an ACK so reclamation cannot admit both owners for one session. + const currentBinding = validateAttachedWorkerRequest(identity, request.runEpoch, { + kind: "transcript", + seq: request.seq, + }); + if (!currentBinding.ok) { + return currentBinding; + } + // Stale base is a terminal sequenced outcome. Advance its durable cursor + // so the next worker commit cannot reuse the consumed sequence number. + if (result.ok || result.reason === "stale-base-leaf") { + const placement = placementBinding(identity); + const processTurn = processTurnBinding(identity); + if (!placement || !processTurn) { + return { ok: false, closeReason: "placement-mismatch" }; + } + options.placementStore?.updateAckCursors({ ...placement, transcriptSeq: request.seq }); + recordAckCursor(processTurn, { transcriptSeq: request.seq }); + } + return result; }); - const pushLiveEvent = ( + const applyLiveEvent = ( + identity: WorkerConnectionIdentity, + request: WorkerLiveEventParams, + ): WorkerLiveEventServiceResult => { + const binding = validateAttachedWorkerRequest(identity, request.runEpoch, { + kind: "live", + seq: request.seq, + }); + if (!binding.ok) { + if ("closeReason" in binding) { + return binding; + } + return { ok: false, details: { reason: binding.reason } }; + } + if (request.runId !== identity.runId) { + return { ok: false, closeReason: "placement-mismatch" }; + } + if (!options.liveEvents) { + return { ok: false, closeReason: "gateway-unavailable" }; + } + // The caller holds the environment lock, preserving order with transcript + // commits and the terminal mutation fence while this synchronous receiver runs. + const result = options.liveEvents.apply({ identity, request }); + if (result.ok) { + const placement = placementBinding(identity); + const processTurn = processTurnBinding(identity); + if (!placement || !processTurn) { + return { ok: false, closeReason: "placement-mismatch" }; + } + options.placementStore?.updateAckCursors({ + ...placement, + liveSeq: result.result.ackedSeq, + }); + recordAckCursor(processTurn, { liveSeq: result.result.ackedSeq }); + } + return result; + }; + + const pushLiveEvent = async ( identity: WorkerConnectionIdentity, request: WorkerLiveEventParams, ): Promise => { - const binding = validateAttachedWorkerRequest(identity, request.runEpoch); - if (!binding.ok) { - if ("closeReason" in binding) { - return Promise.resolve(binding); + return await withLock(identity.environmentId, async () => { + const placement = placementBinding(identity); + const processTurn = processTurnBinding(identity); + const observed = processTurn ? observedAckCursorFor(processTurn) : undefined; + const wasNewSequence = request.seq > (observed?.liveSeq ?? 0); + const result = applyLiveEvent(identity, request); + if (!result.ok || !placement || !processTurn) { + return result; } - return Promise.resolve({ ok: false, details: { reason: binding.reason } }); - } - if (!options.liveEvents) { - return Promise.resolve({ ok: false, closeReason: "gateway-unavailable" }); - } - // Publish after authoritative validation without blocking on lifecycle work. - return Promise.resolve(options.liveEvents.apply({ identity, request })); + const pending = pendingTerminalTurnFences.get(placement.sessionId); + if (pending && !matchesTurnBinding(pending, processTurn)) { + pendingTerminalTurnFences.delete(placement.sessionId); + } + if (isTerminalLiveEvent(request) && wasNewSequence) { + pendingTerminalTurnFences.set(placement.sessionId, { + ...processTurn, + terminalLiveSeq: request.seq, + }); + } + const terminal = pendingTerminalTurnFences.get(placement.sessionId); + if ( + terminal && + matchesTurnBinding(terminal, processTurn) && + result.result.ackedSeq >= terminal.terminalLiveSeq + ) { + // A gap fill can ACK a previously buffered terminal event. Fence from + // the observed high-water marks, not only from the request carrying it. + terminalTurnFences.set( + placement.sessionId, + observedAckCursorFor(processTurn) ?? + recordAckCursor(processTurn, { liveSeq: result.result.ackedSeq }), + ); + pendingTerminalTurnFences.delete(placement.sessionId); + } + return result; + }); }; const revalidateInference = ( @@ -1154,7 +1395,9 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService if (request.sessionId !== identity.sessionId) { return "session-not-attached"; } - const binding = validateAttachedWorkerRequest(identity, request.runEpoch); + const binding = validateAttachedWorkerRequest(identity, request.runEpoch, { + kind: "inference", + }); return binding.ok ? null : "reason" in binding ? binding.reason : "session-not-attached"; }; @@ -1163,10 +1406,12 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService request: WorkerInferenceStartParams, sink: WorkerInferenceSink, ): WorkerInferenceStartServiceResult => { - if (request.sessionId !== identity.sessionId) { + if (request.sessionId !== identity.sessionId || request.runId !== identity.runId) { return { ok: false, reason: "session-not-attached" }; } - const binding = validateAttachedWorkerRequest(identity, request.runEpoch); + const binding = validateAttachedWorkerRequest(identity, request.runEpoch, { + kind: "inference", + }); if (!binding.ok) { return binding; } @@ -1182,10 +1427,12 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService identity: WorkerConnectionIdentity, request: WorkerInferenceCancelParams, ): WorkerInferenceCancelServiceResult => { - if (request.sessionId !== identity.sessionId) { + if (request.sessionId !== identity.sessionId || request.runId !== identity.runId) { return { ok: false, reason: "session-not-attached" }; } - const binding = validateAttachedWorkerRequest(identity, request.runEpoch); + const binding = validateAttachedWorkerRequest(identity, request.runEpoch, { + kind: "inference", + }); if (!binding.ok) { return binding; } @@ -1227,12 +1474,41 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService if (stopping) { return { ok: false, reason: "environment-unavailable" } as const; } - return admitWorkerConnection({ store, admission, expectedBuild, nowMs: now() }); + const admitted = admitWorkerConnection({ store, admission, expectedBuild, nowMs: now() }); + if ( + !admitted.ok || + !options.placementStore || + (admitted.identity.sessionId === null && admitted.identity.runId === null) + ) { + return admitted; + } + const placement = placementBinding(admitted.identity); + if (!placement || !options.placementStore.validateWorkerTurn(placement)) { + return { ok: false, reason: "placement-mismatch" } as const; + } + return admitted; + }, + validateWorkerConnection: (identity: WorkerConnectionIdentity) => { + if (stopping) { + return "environment-unavailable" as const; + } + const placement = validateWorkerPlacement(identity); + if (!placement.valid) { + return "placement-mismatch" as const; + } + const environmentFailure = validateWorkerConnectionIdentity({ + store, + identity, + nowMs: now(), + }); + if ( + environmentFailure && + !(environmentFailure === "credential-expired" && placement.durableClaim) + ) { + return environmentFailure; + } + return null; }, - validateWorkerConnection: (identity: WorkerConnectionIdentity) => - stopping - ? ("environment-unavailable" as const) - : validateWorkerConnectionIdentity({ store, identity, nowMs: now() }), commitTranscript, pushLiveEvent, startInference, @@ -1246,6 +1522,36 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService attachSession, takeMintedCredential: (binding: WorkerCredentialBinding) => readPendingCredential(binding)?.grant, + acquireTurnCredential: (binding: WorkerCredentialBinding & { sessionId: string }) => + withLock(binding.environmentId, async () => { + const pending = readPendingCredential(binding)?.grant; + if (pending) { + return pending; + } + const environment = store.get(binding.environmentId); + if ( + !environment || + environment.state !== "attached" || + environment.ownerEpoch !== binding.ownerEpoch || + environment.attachedSessionIds.length !== 1 || + environment.attachedSessionIds[0] !== binding.sessionId + ) { + throw serviceError("invalid_state", "Worker session credential owner is not attached"); + } + const previous = store.getCredential(binding.environmentId); + const minted = mintCredentialLocked(binding); + const grant = stageCredential(minted.grant); + if (previous?.sessionId === binding.sessionId) { + options.liveEvents?.rotateCredential({ + credentialHash: minted.credentialHash, + environmentId: binding.environmentId, + previousCredentialHash: previous.credentialHash, + runEpoch: binding.ownerEpoch, + sessionId: binding.sessionId, + }); + } + return grant; + }), acknowledgeCredentialDelivery: (claim: WorkerCredentialDeliveryClaim): boolean => { const pending = readPendingCredential(claim); if (!pending || pending.grant.deliveryId !== claim.deliveryId) { diff --git a/src/gateway/worker-environments/transcript-commit.test.ts b/src/gateway/worker-environments/transcript-commit.test.ts index b581d295ab4e..b94b1045a25d 100644 --- a/src/gateway/worker-environments/transcript-commit.test.ts +++ b/src/gateway/worker-environments/transcript-commit.test.ts @@ -38,6 +38,7 @@ const IDENTITY: WorkerConnectionIdentity = { credentialHash: ["credential", "hash", "a"].join("-"), bundleHash: "b".repeat(64), sessionId: SESSION_ID, + runId: "run-worker-transcript", ownerEpoch: RUN_EPOCH, rpcSetVersion: 1, protocolFeatures: ["worker-transcript-commit-v1"], diff --git a/src/gateway/worker-environments/tunnel-contract.ts b/src/gateway/worker-environments/tunnel-contract.ts index 636254b87715..34042b310701 100644 --- a/src/gateway/worker-environments/tunnel-contract.ts +++ b/src/gateway/worker-environments/tunnel-contract.ts @@ -7,10 +7,23 @@ export type WorkerTunnelRequest = { ownerEpoch: number; }; -type WorkerWorkspaceCommand = { +export type WorkerWorkspaceCommand = { argv: readonly string[]; input?: string; timeoutMs?: number; + signal?: AbortSignal; +}; + +export type WorkerWorkspaceSyncRequest = { + localPath: string; + sessionId: string; + generation: number; +}; + +export type WorkerWorkspaceSyncResult = { + mode: "git" | "plain"; + remoteWorkspaceDir: string; + manifestRef: string; }; export type WorkerTunnelHandle = { @@ -18,5 +31,6 @@ export type WorkerTunnelHandle = { ownerEpoch: number; remoteSocketPath: string; runWorkspaceCommand(command: WorkerWorkspaceCommand): Promise; + syncWorkspace(request: WorkerWorkspaceSyncRequest): Promise; stop(): Promise; }; diff --git a/src/gateway/worker-environments/tunnel-ssh-runner.ts b/src/gateway/worker-environments/tunnel-ssh-runner.ts new file mode 100644 index 000000000000..4517d4dffa9f --- /dev/null +++ b/src/gateway/worker-environments/tunnel-ssh-runner.ts @@ -0,0 +1,129 @@ +import { spawn } from "node:child_process"; +import { sliceUtf16Safe } from "@openclaw/normalization-core/utf16-slice"; +import { redactSensitiveText } from "../../logging/redact.js"; +import { + runCommandWithTimeout, + type CommandOptions, + type SpawnResult, +} from "../../process/exec.js"; + +export const WORKER_TUNNEL_READY_MARKER = "OPENCLAW_WORKER_TUNNEL_READY"; + +const STOP_GRACE_MS = 1_500; +const STDERR_LIMIT = 4_096; + +type WorkerSshProcessExit = { + code: number | null; + signal: NodeJS.Signals | null; +}; + +export type WorkerSshProcess = { + ready: Promise; + exited: Promise; + stop(): Promise; +}; + +export type WorkerSshRunner = { + start(argv: string[], options: CommandOptions): WorkerSshProcess; + run(argv: string[], options: CommandOptions): Promise; +}; + +export function workerSshProcessError(stderr: string): Error { + const detail = redactSensitiveText(stderr, { mode: "tools" }).replace(/\s+/gu, " ").trim(); + return new Error(detail ? `Worker SSH tunnel failed: ${detail}` : "Worker SSH tunnel failed"); +} + +/** Production runner that treats the remote post-forward marker as connection readiness. */ +export function createWorkerSshRunner(): WorkerSshRunner { + return { + run: runCommandWithTimeout, + start(argv, options) { + const [command, ...args] = argv; + if (!command) { + throw new Error("Worker SSH runner requires a command"); + } + const child = spawn(command, args, { + env: options.baseEnv, + stdio: ["pipe", "pipe", "pipe"], + windowsHide: true, + }); + let closed = false; + let readySettled = false; + let resolveReady!: () => void; + let rejectReady!: (error: Error) => void; + let resolveExited!: (exit: WorkerSshProcessExit) => void; + const ready = new Promise((resolve, reject) => { + resolveReady = resolve; + rejectReady = reject; + }); + const exited = new Promise((resolve) => { + resolveExited = resolve; + }); + let stdout = ""; + let stderr = ""; + const settleReadyError = () => { + if (readySettled) { + return; + } + readySettled = true; + rejectReady(workerSshProcessError(stderr)); + }; + child.stdout.setEncoding("utf8"); + child.stdout.on("error", () => {}); + child.stdout.on("data", (chunk: string) => { + if (readySettled) { + return; + } + stdout = sliceUtf16Safe(`${stdout}${chunk}`, -STDERR_LIMIT); + if (stdout.split(/\r?\n/u).includes(WORKER_TUNNEL_READY_MARKER)) { + readySettled = true; + resolveReady(); + } + }); + child.stderr.setEncoding("utf8"); + child.stderr.on("error", () => {}); + child.stderr.on("data", (chunk: string) => { + stderr = sliceUtf16Safe(`${stderr}${chunk}`, -STDERR_LIMIT); + }); + child.once("error", settleReadyError); + child.once("close", (code, signal) => { + closed = true; + settleReadyError(); + resolveExited({ code, signal }); + }); + child.stdin.on("error", () => {}); + if (options.input !== undefined) { + child.stdin.end(options.input); + } else { + child.stdin.end(); + } + + let stopPromise: Promise | undefined; + return { + ready, + exited, + stop() { + return (stopPromise ??= (async () => { + if (closed) { + return; + } + child.kill("SIGTERM"); + let timer: ReturnType | undefined; + await Promise.race([ + exited, + new Promise((resolve) => { + timer = setTimeout(resolve, STOP_GRACE_MS); + timer.unref?.(); + }), + ]); + clearTimeout(timer); + if (!closed) { + child.kill("SIGKILL"); + await exited; + } + })()); + }, + }; + }, + }; +} diff --git a/src/gateway/worker-environments/tunnel.test.ts b/src/gateway/worker-environments/tunnel.test.ts index 69006fc989a8..7b91ec9fe520 100644 --- a/src/gateway/worker-environments/tunnel.test.ts +++ b/src/gateway/worker-environments/tunnel.test.ts @@ -1,11 +1,20 @@ +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; import { describe, expect, it, vi } from "vitest"; import type { WorkerSshEndpoint } from "../../plugins/types.js"; -import type { CommandOptions, SpawnResult } from "../../process/exec.js"; +import { + runCommandWithTimeout, + type CommandOptions, + type SpawnResult, +} from "../../process/exec.js"; +import { + createWorkerSshRunner, + type WorkerSshProcess, + type WorkerSshRunner, +} from "./tunnel-ssh-runner.js"; import { createWorkerTunnelManager } from "./tunnel.js"; -type WorkerTunnelOptions = NonNullable[0]>; -type WorkerSshRunner = NonNullable; -type WorkerSshProcess = ReturnType; type WorkerSshProcessExit = Awaited; const HOST_KEY = [["ssh", "ed25519"].join("-"), "AAAA"].join(" "); @@ -17,10 +26,10 @@ const SSH: WorkerSshEndpoint = { keyRef: { source: "file", provider: "workers", id: "/identity" }, }; -function success(): SpawnResult { +function success(stdout = "", stderr = ""): SpawnResult { return { - stdout: "", - stderr: "", + stdout, + stderr, code: 0, signal: null, killed: false, @@ -71,7 +80,7 @@ class FakeProcess implements WorkerSshProcess { } } -function fakeRunner() { +function fakeRunner(onRun?: (argv: string[], options: CommandOptions) => SpawnResult | undefined) { const starts: Array<{ argv: string[]; options: CommandOptions; process: FakeProcess }> = []; const runs: Array<{ argv: string[]; options: CommandOptions }> = []; const runner: WorkerSshRunner = { @@ -82,12 +91,77 @@ function fakeRunner() { }, async run(argv, options) { runs.push({ argv, options }); - return success(); + return onRun?.(argv, options) ?? success(); }, }; return { runner, runs, starts }; } +function localWorkspaceRunner(remoteHome: string) { + const starts: Array<{ argv: string[]; options: CommandOptions; process: FakeProcess }> = []; + const runner: WorkerSshRunner = { + start(argv, options) { + const process = new FakeProcess(); + starts.push({ argv, options, process }); + return process; + }, + async run(argv, options) { + if (argv[0] === "git") { + return await runCommandWithTimeout(argv, options); + } + if (argv[0] === "rsync") { + const localArgv = [...argv]; + const remoteShellIndex = localArgv.indexOf("-e"); + if (remoteShellIndex >= 0) { + localArgv.splice(remoteShellIndex, 2); + } + const destination = localArgv.at(-1); + const separator = destination?.indexOf(":") ?? -1; + if (!destination || separator < 0) { + throw new Error("missing test rsync destination"); + } + const remotePath = destination.slice(separator + 1); + // Prod rsync targets the absolute directory returned by the setup script, + // which already lives under the fake remote HOME. + const localDestination = path.isAbsolute(remotePath) + ? remotePath + : path.join(remoteHome, remotePath); + localArgv[localArgv.length - 1] = localDestination; + await fs.mkdir( + destination.endsWith("/") ? localDestination : path.dirname(localDestination), + { recursive: true }, + ); + return await runCommandWithTimeout(localArgv, options); + } + if (argv[0] === "ssh") { + if (options.input?.includes("unsafe worker tunnel directory")) { + return success(); + } + const remoteCommand = argv.at(-1); + if (!remoteCommand) { + throw new Error("missing test SSH remote command"); + } + return await runCommandWithTimeout(["sh", "-c", remoteCommand], { + ...options, + baseEnv: { ...options.baseEnv, HOME: remoteHome }, + }); + } + throw new Error(`unexpected test command: ${argv[0] ?? "missing"}`); + }, + }; + return { runner, starts }; +} + +async function git(root: string, ...args: string[]): Promise { + const result = await runCommandWithTimeout(["git", "-C", root, ...args], { + timeoutMs: 30_000, + }); + if (result.code !== 0) { + throw new Error(result.stderr || result.stdout || `git ${args[0] ?? "command"} failed`); + } + return result.stdout.trim(); +} + const resolveIdentity = async () => ({ kind: "path", path: "/keys/worker" }) as const; async function waitForStarts(starts: unknown[], count: number) { @@ -115,7 +189,7 @@ describe("worker tunnel manager", () => { expect(tunnel?.argv).toContain("StreamLocalBindUnlink=yes"); expect(tunnel?.options.input).not.toContain("rm -f"); expect(tunnel?.argv[tunnel.argv.indexOf("-R") + 1]).toMatch( - /^\/tmp\/ocw-[a-f0-9]+\/gateway\.sock:127\.0\.0\.1:18789$/u, + /^\/tmp\/ocw-[a-f0-9]{16}-3\/gateway\.sock:127\.0\.0\.1:18789$/u, ); tunnel?.process.becomeReady(); const handle = await starting; @@ -134,6 +208,293 @@ describe("worker tunnel manager", () => { expect(manager.status("worker:one")).toBe("stopped"); }); + it("syncs a dirty workspace over pinned rsync and records an immutable manifest", async () => { + const manifestRef = `sha256:${"b".repeat(64)}`; + const remoteWorkspaceDir = "/home/worker/.openclaw-worker/workspaces/env/session/7"; + const localPath = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-worker-sync-test-")); + await fs.writeFile(path.join(localPath, ".worktreeinclude"), "cache/*.bin\n"); + await git(localPath, "init"); + await git(localPath, "config", "user.name", "Worker Sync Test"); + await git(localPath, "config", "user.email", "worker-sync@example.invalid"); + await fs.mkdir(path.join(localPath, "src"), { recursive: true }); + await fs.writeFile(path.join(localPath, "src/tracked.ts"), "tracked\n"); + await git(localPath, "add", ".worktreeinclude", "src/tracked.ts"); + await git(localPath, "commit", "-m", "base"); + const commit = await git(localPath, "rev-parse", "HEAD"); + const fake = fakeRunner((argv, options) => { + if (argv.includes("--show-toplevel")) { + return success(`${localPath}\n`); + } + if (argv.includes("--verify")) { + return success(`${commit}\n`); + } + if (options.input?.includes("unsafe worker workspace directory")) { + return success(`${remoteWorkspaceDir}\n`); + } + if (argv.at(-1)?.includes("worker workspace symlink escapes")) { + return success(`${manifestRef}\n`); + } + return undefined; + }); + const manager = createWorkerTunnelManager({ runner: fake.runner }); + const starting = manager.start({ + environmentId: "worker:sync", + ownerEpoch: 5, + ssh: SSH, + gateway: { host: "127.0.0.1", port: 18789 }, + resolveIdentity, + }); + await waitForStarts(fake.starts, 1); + fake.starts[0]?.process.becomeReady(); + const handle = await starting; + + try { + await expect( + handle.syncWorkspace({ + localPath, + sessionId: "session:one", + generation: 7, + }), + ).resolves.toEqual({ mode: "git", remoteWorkspaceDir, manifestRef }); + + const transfer = fake.runs.findLast((entry) => entry.argv[0] === "rsync"); + expect(transfer?.argv).toContain("--checksum"); + expect(transfer?.argv).toContain(`${localPath}/`); + expect(transfer?.argv.at(-1)).toBe(`worker@worker.example.test:${remoteWorkspaceDir}/`); + expect(transfer?.argv).not.toContain("--protect-args"); + expect(transfer?.argv.some((arg) => arg.startsWith("--files-from="))).toBe(true); + const remoteShell = transfer?.argv[transfer.argv.indexOf("-e") + 1]; + expect(remoteShell).toContain("ClearAllForwardings=yes"); + expect(remoteShell).toContain("ControlMaster=no"); + expect(remoteShell).toContain("ControlPath=none"); + const manifest = fake.runs.find((entry) => + entry.argv.at(-1)?.includes("worker workspace symlink escapes"), + ); + expect(manifest?.argv.at(-1)).toContain(commit); + } finally { + await handle.stop(); + await fs.rm(localPath, { recursive: true }); + } + }); + + it("fails workspace sync before manifest creation when rsync fails", async () => { + const remoteWorkspaceDir = "/home/worker/.openclaw-worker/workspaces/env/session/2"; + const fake = fakeRunner((argv, options) => { + if (argv[0] === "git") { + return { ...success(), code: 128 }; + } + if (argv[0] === "rsync") { + return { ...success("", "transfer denied"), code: 23 }; + } + if (options.input?.includes("unsafe worker workspace directory")) { + return success(`${remoteWorkspaceDir}\n`); + } + return undefined; + }); + const manager = createWorkerTunnelManager({ runner: fake.runner }); + const starting = manager.start({ + environmentId: "worker:sync-failure", + ownerEpoch: 2, + ssh: SSH, + gateway: { host: "127.0.0.1", port: 18789 }, + resolveIdentity, + }); + await waitForStarts(fake.starts, 1); + fake.starts[0]?.process.becomeReady(); + const handle = await starting; + + await expect( + handle.syncWorkspace({ + localPath: "/gateway/worktrees/session-two", + sessionId: "session:two", + generation: 2, + }), + ).rejects.toThrow("Worker workspace sync failed: transfer denied"); + expect( + fake.runs.some((entry) => entry.argv.at(-1)?.includes("worker workspace symlink escapes")), + ).toBe(false); + + await handle.stop(); + }); + + it("materializes a large dirty git workspace as a credential-free commit-capable clone", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-worker-git-sync-")); + const localPath = path.join(root, "local"); + const remoteHome = path.join(root, "remote-home"); + await Promise.all([ + fs.mkdir(path.join(localPath, "generated"), { recursive: true }), + fs.mkdir(remoteHome, { recursive: true }), + ]); + await git(localPath, "init"); + await git(localPath, "config", "user.name", "Worker Sync Test"); + await git(localPath, "config", "user.email", "worker-sync@example.invalid"); + await Promise.all([ + fs.writeFile(path.join(localPath, ".gitignore"), "cache/**\nprivate/**\n"), + fs.writeFile(path.join(localPath, ".worktreeinclude"), "cache/allowed.txt\n"), + fs.writeFile(path.join(localPath, "gone.txt"), "delete me\n"), + fs.writeFile(path.join(localPath, "rename-old.txt"), "rename me\n"), + fs.writeFile(path.join(localPath, "modified.txt"), "before\n"), + ]); + const largeFiles = Array.from( + { length: 1_800 }, + (_, index) => `generated/long-worker-file-name-${String(index).padStart(4, "0")}.txt`, + ); + await Promise.all( + largeFiles.map((file, index) => fs.writeFile(path.join(localPath, file), `${index}\n`)), + ); + await git(localPath, "add", "."); + await git(localPath, "commit", "-m", "base"); + const firstBase = await git(localPath, "rev-parse", "HEAD"); + await fs.mkdir(path.join(localPath, "vendor/sub/.git"), { recursive: true }); + await fs.writeFile(path.join(localPath, "vendor/sub/.git/secret"), "must not transfer\n"); + await git(localPath, "update-index", "--add", "--cacheinfo", `160000,${firstBase},vendor/sub`); + await git(localPath, "commit", "-m", "record submodule"); + const baseCommit = await git(localPath, "rev-parse", "HEAD"); + + await Promise.all([ + fs.rm(path.join(localPath, "gone.txt")), + fs.rename(path.join(localPath, "rename-old.txt"), path.join(localPath, "rename-new.txt")), + fs.writeFile(path.join(localPath, "modified.txt"), "after\n"), + fs.mkdir(path.join(localPath, "cache"), { recursive: true }), + fs.mkdir(path.join(localPath, "private"), { recursive: true }), + ]); + await Promise.all([ + fs.writeFile(path.join(localPath, "cache/allowed.txt"), "allowed\n"), + fs.writeFile(path.join(localPath, "private/ignored.txt"), "private\n"), + ]); + + const fake = localWorkspaceRunner(remoteHome); + const manager = createWorkerTunnelManager({ runner: fake.runner }); + const starting = manager.start({ + environmentId: "worker:real-git-sync", + ownerEpoch: 11, + ssh: SSH, + gateway: { host: "127.0.0.1", port: 18789 }, + resolveIdentity, + }); + await waitForStarts(fake.starts, 1); + fake.starts[0]?.process.becomeReady(); + const handle = await starting; + + try { + const result = await handle.syncWorkspace({ + localPath, + sessionId: "session:real-git-sync", + generation: 1, + }); + expect(result.mode).toBe("git"); + expect(result.manifestRef).toMatch(/^sha256:[a-f0-9]{64}$/u); + await expect( + fs.readFile(path.join(result.remoteWorkspaceDir, largeFiles[0] ?? ""), "utf8"), + ).resolves.toBe("0\n"); + await expect( + fs.readFile(path.join(result.remoteWorkspaceDir, largeFiles.at(-1) ?? ""), "utf8"), + ).resolves.toBe("1799\n"); + await expect(fs.access(path.join(result.remoteWorkspaceDir, "gone.txt"))).rejects.toThrow(); + await expect( + fs.readFile(path.join(result.remoteWorkspaceDir, "rename-new.txt"), "utf8"), + ).resolves.toBe("rename me\n"); + await expect( + fs.readFile(path.join(result.remoteWorkspaceDir, "cache/allowed.txt"), "utf8"), + ).resolves.toBe("allowed\n"); + await expect( + fs.access(path.join(result.remoteWorkspaceDir, "private/ignored.txt")), + ).rejects.toThrow(); + await expect( + fs.access(path.join(result.remoteWorkspaceDir, "vendor/sub/.git/secret")), + ).rejects.toThrow(); + expect(await git(result.remoteWorkspaceDir, "rev-parse", "HEAD")).toBe(baseCommit); + expect(await git(result.remoteWorkspaceDir, "rev-list", "--count", "HEAD")).toBe("1"); + expect(await git(result.remoteWorkspaceDir, "remote")).toBe(""); + const status = await runCommandWithTimeout( + ["git", "-C", result.remoteWorkspaceDir, "status", "--porcelain"], + { timeoutMs: 30_000 }, + ); + const statusLines = status.stdout.split("\n").filter(Boolean); + expect(statusLines).toContain(" D gone.txt"); + expect(statusLines).toContain("?? rename-new.txt"); + await git(result.remoteWorkspaceDir, "add", "-A"); + await git(result.remoteWorkspaceDir, "commit", "-m", "worker commit"); + await git(result.remoteWorkspaceDir, "merge-base", "--is-ancestor", baseCommit, "HEAD"); + + const manifestPath = path.join( + remoteHome, + ".openclaw-worker/manifests", + `${result.manifestRef.slice("sha256:".length)}.json`, + ); + const manifest = JSON.parse(await fs.readFile(manifestPath, "utf8")) as { + entries: Array<{ path: string }>; + }; + expect(manifest.entries.some((entry) => entry.path === ".git")).toBe(false); + expect(manifest.entries.some((entry) => entry.path.startsWith(".git/"))).toBe(false); + } finally { + await handle.stop(); + await fs.rm(root, { recursive: true }); + } + }, 60_000); + + it("mirrors plain workspaces and rejects escaping symlinks in a git overlay", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-worker-sync-modes-")); + const plainPath = path.join(root, "plain"); + const gitPath = path.join(root, "git"); + const remoteHome = path.join(root, "remote-home"); + await Promise.all([ + fs.mkdir(path.join(plainPath, "nested/.git"), { recursive: true }), + fs.mkdir(gitPath, { recursive: true }), + fs.mkdir(remoteHome, { recursive: true }), + ]); + await Promise.all([ + fs.writeFile(path.join(plainPath, "hello.txt"), "plain\n"), + fs.writeFile(path.join(plainPath, "nested/.git/config"), "private metadata\n"), + ]); + await git(gitPath, "init"); + await git(gitPath, "config", "user.name", "Worker Sync Test"); + await git(gitPath, "config", "user.email", "worker-sync@example.invalid"); + await fs.writeFile(path.join(gitPath, "tracked.txt"), "tracked\n"); + await git(gitPath, "add", "tracked.txt"); + await git(gitPath, "commit", "-m", "base"); + await fs.symlink(path.join(root, "outside"), path.join(gitPath, "escape")); + + const fake = localWorkspaceRunner(remoteHome); + const manager = createWorkerTunnelManager({ runner: fake.runner }); + const starting = manager.start({ + environmentId: "worker:real-sync-modes", + ownerEpoch: 12, + ssh: SSH, + gateway: { host: "127.0.0.1", port: 18789 }, + resolveIdentity, + }); + await waitForStarts(fake.starts, 1); + fake.starts[0]?.process.becomeReady(); + const handle = await starting; + + try { + const plain = await handle.syncWorkspace({ + localPath: plainPath, + sessionId: "session:plain-sync", + generation: 1, + }); + expect(plain.mode).toBe("plain"); + await expect( + fs.readFile(path.join(plain.remoteWorkspaceDir, "hello.txt"), "utf8"), + ).resolves.toBe("plain\n"); + await expect( + fs.access(path.join(plain.remoteWorkspaceDir, "nested/.git/config")), + ).rejects.toThrow(); + + await expect( + handle.syncWorkspace({ + localPath: gitPath, + sessionId: "session:symlink-sync", + generation: 2, + }), + ).rejects.toThrow("worker workspace symlink escapes the sync root"); + } finally { + await handle.stop(); + await fs.rm(root, { recursive: true }); + } + }, 60_000); + it("reconnects with capped backoff after unexpected exits and failed attempts", async () => { const fake = fakeRunner(); const delays: number[] = []; @@ -309,3 +670,16 @@ describe("worker tunnel manager", () => { expect(fake.starts).toHaveLength(1); }); }); + +describe("createWorkerSshRunner diagnostic tails", () => { + it("keeps SSH tunnel failure stderr on a valid UTF-16 boundary", async () => { + const retained = "b".repeat(4095); + const child = createWorkerSshRunner().start( + [process.execPath, "-e", `process.stderr.write(${JSON.stringify(`a😀${retained}`)})`], + { timeoutMs: 10_000, baseEnv: process.env }, + ); + + await expect(child.ready).rejects.toThrow(`Worker SSH tunnel failed: ${retained}`); + await child.exited; + }); +}); diff --git a/src/gateway/worker-environments/tunnel.ts b/src/gateway/worker-environments/tunnel.ts index cf75ab3ecfa0..c7d7c2227151 100644 --- a/src/gateway/worker-environments/tunnel.ts +++ b/src/gateway/worker-environments/tunnel.ts @@ -1,15 +1,7 @@ -import { spawn } from "node:child_process"; -import { randomBytes } from "node:crypto"; -import { sliceUtf16Safe } from "@openclaw/normalization-core/utf16-slice"; import { RetrySupervisor } from "../../../packages/retry/src/index.js"; import { sleepWithAbort, type BackoffPolicy } from "../../infra/backoff.js"; -import { redactSensitiveText } from "../../logging/redact.js"; import type { WorkerSshEndpoint } from "../../plugins/types.js"; -import { - runCommandWithTimeout, - type CommandOptions, - type SpawnResult, -} from "../../process/exec.js"; +import type { SpawnResult } from "../../process/exec.js"; import { prepareWorkerSsh, type PreparedWorkerSsh, @@ -23,14 +15,18 @@ import type { WorkerTunnelRequest, WorkerTunnelStatus, } from "./tunnel-contract.js"; +import { + createWorkerSshRunner, + type WorkerSshProcess, + type WorkerSshRunner, + workerSshProcessError, + WORKER_TUNNEL_READY_MARKER, +} from "./tunnel-ssh-runner.js"; +import { createWorkerWorkspaceActions, stableWorkerPathComponent } from "./workspace-sync.js"; export type { WorkerTunnelHandle } from "./tunnel-contract.js"; -const READY_MARKER = "OPENCLAW_WORKER_TUNNEL_READY"; const REMOTE_SOCKET_NAME = "gateway.sock"; const REMOTE_SETUP_TIMEOUT_MS = 20_000; -const WORKSPACE_TIMEOUT_MS = 10 * 60_000; -const STOP_GRACE_MS = 1_500; -const STDERR_LIMIT = 4_096; const DEFAULT_STABLE_CONNECTION_MS = 30_000; const DEFAULT_BACKOFF: BackoffPolicy = { initialMs: 250, @@ -43,7 +39,14 @@ const REMOTE_SOCKET_SETUP_SCRIPT = String.raw`set -eu directory=$1 socket=$2 umask 077 -mkdir -p -- "$directory" +if [ -e "$directory" ] || [ -L "$directory" ]; then + if [ ! -d "$directory" ] || [ -L "$directory" ]; then + printf '%s\n' 'unsafe worker tunnel directory' >&2 + exit 2 + fi +else + mkdir -- "$directory" +fi chmod 700 -- "$directory" rm -f -- "$socket" `; @@ -51,7 +54,7 @@ rm -f -- "$socket" const REMOTE_TUNNEL_READY_SCRIPT = String.raw`set -eu socket=$1 test -S "$socket" -printf '%s\n' '${READY_MARKER}' +printf '%s\n' '${WORKER_TUNNEL_READY_MARKER}' trap 'exit 0' HUP INT TERM while :; do sleep 3600; done `; @@ -63,22 +66,6 @@ rm -f -- "$socket" rmdir -- "$directory" 2>/dev/null || true `; -type WorkerSshProcessExit = { - code: number | null; - signal: NodeJS.Signals | null; -}; - -type WorkerSshProcess = { - ready: Promise; - exited: Promise; - stop(): Promise; -}; - -type WorkerSshRunner = { - start(argv: string[], options: CommandOptions): WorkerSshProcess; - run(argv: string[], options: CommandOptions): Promise; -}; - type WorkerTunnelStartRequest = WorkerTunnelRequest & { gateway: { host: "127.0.0.1" | "::1"; port: number }; ssh: WorkerSshEndpoint; @@ -113,106 +100,6 @@ type WorkerTunnelManagerOptions = { stableConnectionMs?: number; }; -function processError(stderr: string): Error { - const detail = redactSensitiveText(stderr, { mode: "tools" }).replace(/\s+/gu, " ").trim(); - return new Error(detail ? `Worker SSH tunnel failed: ${detail}` : "Worker SSH tunnel failed"); -} - -/** Production runner that treats the remote post-forward marker as connection readiness. */ -function createWorkerSshRunner(): WorkerSshRunner { - return { - run: runCommandWithTimeout, - start(argv, options) { - const [command, ...args] = argv; - if (!command) { - throw new Error("Worker SSH runner requires a command"); - } - const child = spawn(command, args, { - env: options.baseEnv, - stdio: ["pipe", "pipe", "pipe"], - windowsHide: true, - }); - let closed = false; - let readySettled = false; - let resolveReady!: () => void; - let rejectReady!: (error: Error) => void; - let resolveExited!: (exit: WorkerSshProcessExit) => void; - const ready = new Promise((resolve, reject) => { - resolveReady = resolve; - rejectReady = reject; - }); - const exited = new Promise((resolve) => { - resolveExited = resolve; - }); - let stdout = ""; - let stderr = ""; - const settleReadyError = () => { - if (readySettled) { - return; - } - readySettled = true; - rejectReady(processError(stderr)); - }; - child.stdout.setEncoding("utf8"); - child.stdout.on("error", () => {}); - child.stdout.on("data", (chunk: string) => { - if (readySettled) { - return; - } - stdout = sliceUtf16Safe(`${stdout}${chunk}`, -STDERR_LIMIT); - if (stdout.split(/\r?\n/u).includes(READY_MARKER)) { - readySettled = true; - resolveReady(); - } - }); - child.stderr.setEncoding("utf8"); - child.stderr.on("error", () => {}); - child.stderr.on("data", (chunk: string) => { - stderr = sliceUtf16Safe(`${stderr}${chunk}`, -STDERR_LIMIT); - }); - child.once("error", settleReadyError); - child.once("close", (code, signal) => { - closed = true; - settleReadyError(); - resolveExited({ code, signal }); - }); - child.stdin.on("error", () => {}); - if (options.input !== undefined) { - child.stdin.end(options.input); - } else { - child.stdin.end(); - } - - let stopPromise: Promise | undefined; - return { - ready, - exited, - stop() { - return (stopPromise ??= (async () => { - if (closed) { - return; - } - child.kill("SIGTERM"); - let timer: ReturnType | undefined; - await Promise.race([ - exited, - new Promise((resolve) => { - timer = setTimeout(resolve, STOP_GRACE_MS); - timer.unref?.(); - }), - ]); - clearTimeout(timer); - if (!closed) { - child.kill("SIGKILL"); - await exited; - } - })()); - }, - }; - }, - }; -} - function success(result: SpawnResult): boolean { return result.termination === "exit" && result.code === 0; } @@ -285,7 +172,7 @@ export function createWorkerTunnelManager(options: WorkerTunnelManagerOptions = }); const result = await runner.run(command.argv, command.options); if (!success(result)) { - throw processError(result.stderr || result.stdout); + throw workerSshProcessError(result.stderr || result.stdout); } }; @@ -304,36 +191,14 @@ export function createWorkerTunnelManager(options: WorkerTunnelManagerOptions = environmentId: entry.environmentId, ownerEpoch: entry.ownerEpoch, remoteSocketPath: entry.remoteSocketPath, - async runWorkspaceCommand(command) { - if (!isCurrent(entry) || !entry.prepared || entry.status !== "connected") { - throw new Error("Worker tunnel owner is no longer connected"); - } - const task = runner.run( - [ - "ssh", - ...workerSshOptions(entry.prepared, { forwarding: "disabled" }), - "-a", - "-x", - "-T", - "-p", - String(entry.prepared.port), - "--", - entry.prepared.sshTarget, - workerSshRemoteCommand(command.argv), - ], - workerSshCommandOptions({ - input: command.input, - timeoutMs: command.timeoutMs ?? WORKSPACE_TIMEOUT_MS, - signal: entry.abortController.signal, - }), - ); - entry.workspaceTasks.add(task); - void task.then( - () => entry.workspaceTasks.delete(task), - () => entry.workspaceTasks.delete(task), - ); - return await task; - }, + ...createWorkerWorkspaceActions({ + environmentId: entry.environmentId, + ownerSignal: entry.abortController.signal, + isConnected: () => isCurrent(entry) && entry.status === "connected", + getPrepared: () => entry.prepared, + runner, + tasks: entry.workspaceTasks, + }), stop: () => stop(entry.environmentId, entry.ownerEpoch), }); @@ -469,7 +334,8 @@ export function createWorkerTunnelManager(options: WorkerTunnelManagerOptions = rejectReady = reject; }); void ready.catch(() => undefined); - const remoteDirectory = `/tmp/ocw-${randomBytes(8).toString("hex")}`; + const environmentKey = stableWorkerPathComponent(request.environmentId, 16); + const remoteDirectory = `/tmp/ocw-${environmentKey}-${request.ownerEpoch}`; const entry: TunnelEntry = { environmentId: request.environmentId, ownerEpoch: request.ownerEpoch, diff --git a/src/gateway/worker-environments/worker-turn-launcher.test.ts b/src/gateway/worker-environments/worker-turn-launcher.test.ts new file mode 100644 index 000000000000..0a0fba95be53 --- /dev/null +++ b/src/gateway/worker-environments/worker-turn-launcher.test.ts @@ -0,0 +1,1305 @@ +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { SessionManager } from "../../agents/sessions/session-manager.js"; +import { + makeAgentAssistantMessage, + makeAgentUserMessage, +} from "../../agents/test-helpers/agent-message-fixtures.js"; +import type { SpawnResult } from "../../process/exec.js"; +import { createDeferred } from "../../shared/deferred.js"; +import { + closeOpenClawStateDatabaseForTest, + openOpenClawStateDatabase, + type OpenClawStateDatabase, +} from "../../state/openclaw-state-db.js"; +import { + parseWorkerLaunchDescriptor, + type WorkerLaunchDescriptor, +} from "../../worker/launch-descriptor.js"; +import type { MintedWorkerCredential } from "./credential.js"; +import { + createWorkerSessionPlacementStore, + type WorkerSessionPlacementStore, +} from "./placement-store.js"; +import { createWorkerSessionPlacementGate } from "./placement-worker-gate.js"; +import type { WorkerTunnelHandle } from "./tunnel-contract.js"; +import { createWorkerSessionTurnPlacementProvider } from "./worker-turn-launcher.js"; + +type WorkerTurnLauncherOptions = Parameters[0]; +type WorkerTurnEnvironmentService = WorkerTurnLauncherOptions["environments"]; + +const SESSION_ID = "session-worker-turn"; +const SESSION_KEY = "agent:main:worker-turn"; +const ENVIRONMENT_ID = "environment-worker-turn"; +const OWNER_EPOCH = 3; +const BUNDLE_HASH = "a".repeat(64); +const HOST_KEY = [["ssh", "ed25519"].join("-"), "AAAA"].join(" "); +type WorkerTurnEnvironmentRecord = NonNullable>; + +describe("worker turn launcher", () => { + let root: string; + let database: OpenClawStateDatabase; + let placements: WorkerSessionPlacementStore; + let sessionFile: string; + + beforeEach(async () => { + root = await fs.mkdtemp(path.join(await fs.realpath(os.tmpdir()), "openclaw-worker-turn-")); + database = openOpenClawStateDatabase({ env: { OPENCLAW_STATE_DIR: root } }); + placements = createWorkerSessionPlacementStore({ database }); + const manager = SessionManager.create(path.join(root, "sessions"), path.join(root, "sessions")); + const file = manager.getSessionFile(); + if (!file) { + throw new Error("expected file-backed session manager"); + } + sessionFile = file; + }); + + afterEach(async () => { + closeOpenClawStateDatabaseForTest(); + await fs.rm(root, { recursive: true, force: true }); + }); + + function seedActivePlacement(): void { + let placement = placements.startDispatch({ + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + }); + placement = placements.transition({ + sessionId: SESSION_ID, + from: "requested", + to: "provisioning", + expectedGeneration: placement.generation, + patch: { environmentId: ENVIRONMENT_ID }, + }); + placement = placements.transition({ + sessionId: SESSION_ID, + from: "provisioning", + to: "syncing", + expectedGeneration: placement.generation, + patch: { workerBundleHash: BUNDLE_HASH }, + }); + placement = placements.transition({ + sessionId: SESSION_ID, + from: "syncing", + to: "starting", + expectedGeneration: placement.generation, + patch: { + remoteWorkspaceDir: "/worker/workspace", + workspaceBaseManifestRef: `sha256:${"b".repeat(64)}`, + }, + }); + placements.transition({ + sessionId: SESSION_ID, + from: "starting", + to: "active", + expectedGeneration: placement.generation, + patch: { activeOwnerEpoch: OWNER_EPOCH }, + }); + } + + function seedReclaimedPlacement() { + seedActivePlacement(); + const active = placements.get(SESSION_ID); + if (active?.state !== "active") { + throw new Error("expected active placement to reclaim"); + } + const draining = placements.startDrain({ + sessionId: SESSION_ID, + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + expectedGeneration: active.generation, + }); + const reconciling = placements.startReconcile({ + sessionId: SESSION_ID, + environmentId: active.environmentId, + ownerEpoch: active.activeOwnerEpoch, + expectedGeneration: draining.generation, + }); + const reclaimed = placements.transition({ + sessionId: SESSION_ID, + from: "reconciling", + to: "reclaimed", + expectedGeneration: reconciling.generation, + }); + if (reclaimed.state !== "reclaimed") { + throw new Error("expected reclaimed placement"); + } + return reclaimed; + } + + function attachedEnvironment(): WorkerTurnEnvironmentRecord { + return { + environmentId: ENVIRONMENT_ID, + providerId: "fake", + profileId: "development", + profileSnapshot: { settings: { region: "test" } }, + provisionOperationId: "provision-worker-turn", + bootstrapReceipt: { + bundleHash: BUNDLE_HASH, + openclawVersion: "2026.7.2", + protocolFeatures: [], + }, + ownerEpoch: OWNER_EPOCH, + teardownTerminalState: null, + attachedSessionIds: [SESSION_ID], + lastError: null, + createdAtMs: 1, + updatedAtMs: 1, + stateChangedAtMs: 1, + idleSinceAtMs: null, + destroyRequestedAtMs: null, + tunnelStatus: "connected", + state: "attached", + leaseId: "lease-worker-turn", + sshEndpoint: { + host: "worker.example.test", + port: 22, + user: "worker", + hostKey: HOST_KEY, + keyRef: { source: "file", provider: "worker-keys", id: "/worker/key" }, + }, + }; + } + + function credential(deliveryId = "c".repeat(43)): MintedWorkerCredential { + return { + credential: ["worker", "turn", "credential"].join("-"), + deliveryId, + environmentId: ENVIRONMENT_ID, + bundleHash: BUNDLE_HASH, + sessionId: SESSION_ID, + rpcSetVersion: 1, + ownerEpoch: OWNER_EPOCH, + expiresAtMs: Date.now() + 60_000, + }; + } + + function unusedEnvironments(): WorkerTurnEnvironmentService { + const unexpected = () => new Error("unexpected worker environment call"); + return { + get: vi.fn(() => undefined), + acquireTurnCredential: vi.fn(async () => { + throw unexpected(); + }), + acknowledgeCredentialDelivery: vi.fn(() => { + throw unexpected(); + }), + startTunnel: vi.fn(async () => { + throw unexpected(); + }), + stopTunnel: vi.fn(async () => { + throw unexpected(); + }), + destroy: vi.fn(async () => { + throw unexpected(); + }), + }; + } + + function turn(runId = "run-worker-turn") { + return { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + sessionFile, + workspaceDir: root, + prompt: "Inspect this workspace", + timeoutMs: 5_000, + runId, + provider: "openai", + model: "gpt-test", + config: { + agents: { + defaults: { + models: { + "openai/gpt-test": { agentRuntime: { id: "openclaw" } }, + }, + }, + }, + }, + }; + } + + it("atomically claims and releases a local turn around the local loop", async () => { + const environments = unusedEnvironments(); + const provider = createWorkerSessionTurnPlacementProvider({ environments, placements }); + + const result = await provider.executeTurn( + { sessionId: SESSION_ID, sessionKey: SESSION_KEY, agentId: "main", runId: "run-local" }, + turn("run-local"), + async () => { + expect(placements.get(SESSION_ID)?.turnClaim).toMatchObject({ + owner: "local", + runId: "run-local", + }); + return { payloads: [{ text: "local" }], meta: { durationMs: 1 } }; + }, + ); + + expect(result.payloads).toEqual([{ text: "local" }]); + expect(placements.get(SESSION_ID)?.turnClaim).toBeNull(); + }); + + it("leaves no placement row for an auxiliary model run without a session key", async () => { + const provider = createWorkerSessionTurnPlacementProvider({ + environments: unusedEnvironments(), + placements, + }); + const runLocal = vi.fn(async () => ({ meta: { durationMs: 1 } })); + + await provider.executeTurn( + { sessionId: SESSION_ID, agentId: "main", runId: "run-model-probe" }, + { ...turn("run-model-probe"), modelRun: true }, + runLocal, + ); + + expect(runLocal).toHaveBeenCalledOnce(); + expect(placements.list()).toEqual([]); + }); + + it("keeps recovery-only admission invisible for sessions without durable placement", async () => { + const provider = createWorkerSessionTurnPlacementProvider({ + admitNewPlacements: false, + environments: unusedEnvironments(), + placements, + }); + + await provider.executeTurn( + { sessionId: SESSION_ID, sessionKey: SESSION_KEY, agentId: "main", runId: "run-local" }, + turn("run-local"), + async () => ({ meta: { durationMs: 1 } }), + ); + await provider.executeLocalTurn( + { sessionId: SESSION_ID, sessionKey: SESSION_KEY, agentId: "main", runId: "run-cli" }, + async () => ({ kind: "cli" }), + ); + + expect(placements.list()).toEqual([]); + }); + + it("still admits an existing local placement in recovery-only mode", async () => { + const seedClaim = placements.claimTurn({ + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + claimId: "seed-local-placement", + runId: "seed-local-placement", + owner: { kind: "local" }, + }); + placements.releaseTurn(seedClaim); + const provider = createWorkerSessionTurnPlacementProvider({ + admitNewPlacements: false, + environments: unusedEnvironments(), + placements, + }); + + await provider.executeTurn( + { sessionId: SESSION_ID, runId: "run-existing-local" }, + turn("run-existing-local"), + async () => { + expect(placements.get(SESSION_ID)?.turnClaim).toMatchObject({ + owner: "local", + runId: "run-existing-local", + }); + return { meta: { durationMs: 1 } }; + }, + ); + + expect(placements.get(SESSION_ID)).toMatchObject({ state: "local", turnClaim: null }); + }); + + it("holds a local placement claim around CLI execution", async () => { + const environments = unusedEnvironments(); + const provider = createWorkerSessionTurnPlacementProvider({ environments, placements }); + + const result = await provider.executeLocalTurn( + { sessionId: SESSION_ID, sessionKey: SESSION_KEY, agentId: "main", runId: "run-cli" }, + async () => { + expect(placements.get(SESSION_ID)?.turnClaim).toMatchObject({ + owner: "local", + runId: "run-cli", + }); + return { kind: "cli" }; + }, + ); + + expect(result).toEqual({ kind: "cli" }); + expect(placements.get(SESSION_ID)?.turnClaim).toBeNull(); + }); + + it("mints a fresh claim token when a later turn reuses the run id", async () => { + const environments = unusedEnvironments(); + const provider = createWorkerSessionTurnPlacementProvider({ environments, placements }); + const claimIds: string[] = []; + const claim = { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-reused", + }; + + for (let index = 0; index < 2; index += 1) { + await provider.executeLocalTurn(claim, async () => { + const claimId = placements.get(SESSION_ID)?.turnClaim?.claimId; + if (!claimId) { + throw new Error("expected active placement claim"); + } + claimIds.push(claimId); + }); + } + + expect(claimIds).toHaveLength(2); + expect(claimIds[0]).not.toBe(claimIds[1]); + expect(placements.get(SESSION_ID)?.turnClaim).toBeNull(); + }); + + it("does not let a stale local finally release a reclaimed run id", async () => { + const provider = createWorkerSessionTurnPlacementProvider({ + environments: unusedEnvironments(), + placements, + }); + const firstStarted = createDeferred(); + const releaseFirst = createDeferred(); + const secondStarted = createDeferred(); + const releaseSecond = createDeferred(); + const claim = { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-restarted", + }; + + const first = provider.executeLocalTurn(claim, async () => { + firstStarted.resolve(); + await releaseFirst.promise; + }); + await firstStarted.promise; + const firstClaimId = placements.get(SESSION_ID)?.turnClaim?.claimId; + expect(placements.clearLocalTurnClaimsAfterRestart()).toBe(1); + + const second = provider.executeLocalTurn(claim, async () => { + secondStarted.resolve(); + await releaseSecond.promise; + }); + await secondStarted.promise; + const secondClaimId = placements.get(SESSION_ID)?.turnClaim?.claimId; + expect(secondClaimId).toBeTruthy(); + expect(secondClaimId).not.toBe(firstClaimId); + + releaseFirst.resolve(); + await first; + expect(placements.get(SESSION_ID)?.turnClaim?.claimId).toBe(secondClaimId); + + releaseSecond.resolve(); + await second; + expect(placements.get(SESSION_ID)?.turnClaim).toBeNull(); + }); + + it("rejects local CLI execution after worker activation", async () => { + seedActivePlacement(); + const environments = unusedEnvironments(); + const provider = createWorkerSessionTurnPlacementProvider({ environments, placements }); + const runLocal = vi.fn(async () => ({ kind: "cli" })); + + await expect( + provider.executeLocalTurn( + { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-local-after-dispatch", + }, + runLocal, + ), + ).rejects.toThrow(`Local turn rejected for session ${SESSION_ID} in placement active`); + + expect(runLocal).not.toHaveBeenCalled(); + expect(placements.get(SESSION_ID)).toMatchObject({ state: "active", turnClaim: null }); + }); + + it.each([ + ["CLI", "claude-cli"], + ["plugin", "test-harness"], + ])( + "rejects an active worker turn assigned to a configured %s runtime", + async (_kind, runtimeId) => { + seedActivePlacement(); + const getEnvironment = vi.fn(() => undefined); + const environments: WorkerTurnEnvironmentService = { + ...unusedEnvironments(), + get: getEnvironment, + }; + const provider = createWorkerSessionTurnPlacementProvider({ environments, placements }); + const runLocal = vi.fn(async () => ({ meta: { durationMs: 1 } })); + const runId = `run-${runtimeId}`; + + await expect( + provider.executeTurn( + { sessionId: SESSION_ID, sessionKey: SESSION_KEY, agentId: "main", runId }, + { + ...turn(runId), + config: { + agents: { + defaults: { + models: { + "openai/gpt-test": { agentRuntime: { id: runtimeId } }, + }, + }, + }, + }, + }, + runLocal, + ), + ).rejects.toThrow(`Cloud worker turns require the OpenClaw runtime, not ${runtimeId}`); + + expect(runLocal).not.toHaveBeenCalled(); + expect(getEnvironment).not.toHaveBeenCalled(); + expect(placements.get(SESSION_ID)).toMatchObject({ state: "active", turnClaim: null }); + }, + ); + + it("launches the active worker with projected history and releases its claim", async () => { + seedActivePlacement(); + const manager = SessionManager.open(sessionFile); + const earlierRequestId = manager.appendMessage( + makeAgentUserMessage({ content: "Earlier request", timestamp: 10 }), + ); + manager.appendMessage( + makeAgentAssistantMessage({ + content: [{ type: "toolCall", id: "call-1", name: "read", arguments: {} }], + timestamp: 11, + }), + ); + manager.appendCustomMessageEntry("context", "Custom durable context", true, {}); + manager.appendCompaction("Compacted durable context", earlierRequestId, 100); + manager.appendMessage({ + role: "toolResult", + toolCallId: "call-1", + toolName: "read", + content: [{ type: "text", text: "result" }], + isError: false, + timestamp: 12, + }); + let descriptor: WorkerLaunchDescriptor | undefined; + const acknowledgeCredentialDelivery = vi.fn(() => true); + const tunnel: WorkerTunnelHandle = { + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + remoteSocketPath: "/worker/gateway.sock", + runWorkspaceCommand: vi.fn(async (command): Promise => { + expect(placements.get(SESSION_ID)?.turnClaim).toMatchObject({ + owner: "worker", + runId: "run-worker-turn", + ownerEpoch: OWNER_EPOCH, + }); + descriptor = parseWorkerLaunchDescriptor(JSON.parse(command.input ?? "")); + expect(command.argv).toEqual([ + "sh", + "-c", + 'exec node "$HOME/.openclaw-worker/$1/openclaw.mjs" worker', + "openclaw-worker", + BUNDLE_HASH, + ]); + expect(command.argv.join(" ")).not.toContain(credential().credential); + await Promise.resolve(); + expect(acknowledgeCredentialDelivery).toHaveBeenCalledOnce(); + const completed = SessionManager.open(sessionFile); + const leafId = completed.appendMessage( + makeAgentAssistantMessage({ + content: [{ type: "text", text: "Worker reply" }], + timestamp: 21, + }), + ); + createWorkerSessionPlacementGate(placements).updateAckCursors({ + sessionId: SESSION_ID, + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + runId: "run-worker-turn", + transcriptSeq: 1, + }); + return { + stdout: JSON.stringify({ + status: "completed", + transcriptLeafId: leafId, + transcriptNextSeq: 2, + }), + stderr: "", + code: 0, + signal: null, + killed: false, + termination: "exit", + }; + }), + syncWorkspace: vi.fn(async () => { + throw new Error("unexpected workspace sync"); + }), + stop: vi.fn(async () => {}), + }; + const environments: WorkerTurnEnvironmentService = { + get: vi.fn(() => attachedEnvironment()), + acquireTurnCredential: vi.fn(async () => credential()), + acknowledgeCredentialDelivery, + startTunnel: vi.fn(async () => tunnel), + stopTunnel: vi.fn(async () => {}), + destroy: vi.fn(async () => attachedEnvironment()), + }; + const provider = createWorkerSessionTurnPlacementProvider({ environments, placements }); + const runLocal = vi.fn(async () => ({ meta: { durationMs: 1 } })); + + const result = await provider.executeTurn( + { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-worker-turn", + }, + { ...turn(), transcriptPrompt: "Canonical transcript request" }, + runLocal, + ); + + expect(runLocal).not.toHaveBeenCalled(); + expect(result.payloads).toEqual([{ text: "Worker reply" }]); + expect(placements.get(SESSION_ID)?.turnClaim).toBeNull(); + expect(descriptor?.assignment.prompt).toBe("Inspect this workspace"); + expect(descriptor?.assignment.suppressPromptTranscript).toBe(true); + expect(descriptor?.assignment.initialMessages).toEqual([ + { + role: "user", + content: [ + { + type: "text", + text: expect.stringContaining("Compacted durable context"), + }, + ], + timestamp: expect.any(Number), + }, + { + role: "user", + content: [{ type: "text", text: "Earlier request" }], + timestamp: 10, + }, + expect.objectContaining({ role: "assistant" }), + { + role: "user", + content: [{ type: "text", text: "Custom durable context" }], + timestamp: expect.any(Number), + }, + { + role: "toolResult", + toolCallId: "call-1", + toolName: "read", + content: [{ type: "text", text: "result" }], + isError: false, + timestamp: 12, + }, + ]); + expect( + SessionManager.open(sessionFile) + .getEntries() + .flatMap((entry) => + entry.type === "message" && entry.message.role === "user" ? [entry.message.content] : [], + ), + ).toContainEqual([{ type: "text", text: "Canonical transcript request" }]); + }); + + it("does not replay an already-persisted current user message into worker history", async () => { + seedActivePlacement(); + const manager = SessionManager.open(sessionFile); + manager.appendMessage(makeAgentUserMessage({ content: "Earlier request", timestamp: 18 })); + manager.appendMessage( + makeAgentAssistantMessage({ + content: [{ type: "text", text: "Earlier reply" }], + timestamp: 19, + }), + ); + manager.appendMessage( + makeAgentUserMessage({ content: "Inspect this workspace", timestamp: 20 }), + ); + let descriptor: WorkerLaunchDescriptor | undefined; + const tunnel: WorkerTunnelHandle = { + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + remoteSocketPath: "/worker/gateway.sock", + runWorkspaceCommand: vi.fn(async (command): Promise => { + descriptor = parseWorkerLaunchDescriptor(JSON.parse(command.input ?? "")); + const completed = SessionManager.open(sessionFile); + const leafId = completed.appendMessage( + makeAgentAssistantMessage({ + content: [{ type: "text", text: "Worker reply" }], + timestamp: 21, + }), + ); + createWorkerSessionPlacementGate(placements).updateAckCursors({ + sessionId: SESSION_ID, + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + runId: "run-persisted-user", + transcriptSeq: 1, + }); + return { + stdout: JSON.stringify({ + status: "completed", + transcriptLeafId: leafId, + transcriptNextSeq: 2, + }), + stderr: "", + code: 0, + signal: null, + killed: false, + termination: "exit", + }; + }), + syncWorkspace: vi.fn(async () => { + throw new Error("unexpected workspace sync"); + }), + stop: vi.fn(async () => {}), + }; + const environments: WorkerTurnEnvironmentService = { + get: vi.fn(() => attachedEnvironment()), + acquireTurnCredential: vi.fn(async () => credential()), + acknowledgeCredentialDelivery: vi.fn(() => true), + startTunnel: vi.fn(async () => tunnel), + stopTunnel: vi.fn(async () => {}), + destroy: vi.fn(async () => attachedEnvironment()), + }; + const provider = createWorkerSessionTurnPlacementProvider({ environments, placements }); + + await provider.executeTurn( + { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-persisted-user", + }, + { + ...turn("run-persisted-user"), + suppressNextUserMessagePersistence: true, + }, + async () => ({ meta: { durationMs: 1 } }), + ); + + expect(descriptor?.assignment.prompt).toBe("Inspect this workspace"); + expect(descriptor?.assignment.initialMessages).toMatchObject([ + { role: "user" }, + { role: "assistant" }, + ]); + const persistedEntries = (await fs.readFile(sessionFile, "utf8")) + .split("\n") + .filter(Boolean) + .map((line) => JSON.parse(line) as unknown); + const persistedCurrentUsers = persistedEntries.filter((entry) => { + if (typeof entry !== "object" || entry === null || !("message" in entry)) { + return false; + } + const message = entry.message; + if ( + typeof message !== "object" || + message === null || + !("role" in message) || + !("content" in message) + ) { + return false; + } + return ( + message.role === "user" && + (message.content === "Inspect this workspace" || + (Array.isArray(message.content) && + message.content.some( + (part) => + typeof part === "object" && + part !== null && + "text" in part && + part.text === "Inspect this workspace", + ))) + ); + }); + expect(persistedCurrentUsers).toHaveLength(1); + }); + + it("reports canonical multi-call usage and the terminal provider model", async () => { + seedActivePlacement(); + const environments: WorkerTurnEnvironmentService = { + get: vi.fn(() => attachedEnvironment()), + acquireTurnCredential: vi.fn(async () => credential()), + acknowledgeCredentialDelivery: vi.fn(() => true), + startTunnel: vi.fn(async () => ({ + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + remoteSocketPath: "/worker/gateway.sock", + runWorkspaceCommand: vi.fn(async (): Promise => { + const completed = SessionManager.open(sessionFile); + completed.appendMessage( + makeAgentAssistantMessage({ + content: [{ type: "toolCall", id: "call-usage", name: "read", arguments: {} }], + provider: "openai", + model: "gpt-first-call", + stopReason: "toolUse", + timestamp: 21, + usage: { + input: 100, + output: 10, + cacheRead: 20, + cacheWrite: 5, + totalTokens: 135, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }), + ); + completed.appendMessage({ + role: "toolResult", + toolCallId: "call-usage", + toolName: "read", + content: [{ type: "text", text: "usage result" }], + isError: false, + timestamp: 22, + }); + const leafId = completed.appendMessage( + makeAgentAssistantMessage({ + content: [{ type: "text", text: "Usage reply" }], + provider: "anthropic", + model: "claude-reported", + timestamp: 23, + usage: { + input: 200, + output: 30, + cacheRead: 40, + cacheWrite: 0, + contextUsage: { + state: "available", + promptTokens: 240, + totalTokens: 270, + }, + totalTokens: 270, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }), + ); + createWorkerSessionPlacementGate(placements).updateAckCursors({ + sessionId: SESSION_ID, + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + runId: "run-worker-usage", + transcriptSeq: 1, + }); + return { + stdout: JSON.stringify({ + status: "completed", + transcriptLeafId: leafId, + transcriptNextSeq: 2, + }), + stderr: "", + code: 0, + signal: null, + killed: false, + termination: "exit", + }; + }), + syncWorkspace: vi.fn(async () => { + throw new Error("unexpected workspace sync"); + }), + stop: vi.fn(async () => {}), + })), + stopTunnel: vi.fn(async () => {}), + destroy: vi.fn(async () => attachedEnvironment()), + }; + const provider = createWorkerSessionTurnPlacementProvider({ environments, placements }); + + const result = await provider.executeTurn( + { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-worker-usage", + }, + turn("run-worker-usage"), + async () => ({ meta: { durationMs: 1 } }), + ); + + expect(result.meta.agentMeta).toEqual({ + sessionId: SESSION_ID, + sessionFile, + provider: "anthropic", + model: "claude-reported", + usage: { + input: 300, + output: 40, + cacheRead: 60, + cacheWrite: 5, + total: 270, + }, + lastCallUsage: { + input: 200, + output: 30, + cacheRead: 40, + cacheWrite: 0, + contextUsage: { + state: "available", + promptTokens: 240, + totalTokens: 270, + }, + total: 270, + }, + promptTokens: 240, + }); + }); + + it("keeps an active placement when tunnel startup fails before remote handoff", async () => { + seedActivePlacement(); + const acknowledgeCredentialDelivery = vi.fn(() => true); + const stopTunnel = vi.fn(async () => {}); + const destroy = vi.fn(async () => attachedEnvironment()); + const environments: WorkerTurnEnvironmentService = { + get: vi.fn(() => attachedEnvironment()), + acquireTurnCredential: vi.fn(async () => credential()), + acknowledgeCredentialDelivery, + startTunnel: vi.fn(async () => { + throw new Error("tunnel unavailable"); + }), + stopTunnel, + destroy, + }; + const provider = createWorkerSessionTurnPlacementProvider({ environments, placements }); + const runLocal = vi.fn(async () => ({ meta: { durationMs: 1 } })); + + await expect( + provider.executeTurn( + { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-tunnel-unavailable", + }, + turn("run-tunnel-unavailable"), + runLocal, + ), + ).rejects.toThrow("tunnel unavailable"); + + expect(runLocal).not.toHaveBeenCalled(); + expect(acknowledgeCredentialDelivery).not.toHaveBeenCalled(); + expect(stopTunnel).not.toHaveBeenCalled(); + expect(destroy).not.toHaveBeenCalled(); + expect(placements.get(SESSION_ID)).toMatchObject({ state: "active", turnClaim: null }); + }); + + it("fails placement and tears down after an ambiguous remote launch failure", async () => { + seedActivePlacement(); + const teardownStates: string[] = []; + const observedPlacements: WorkerSessionPlacementStore = { + ...placements, + startReconcile: (input) => { + teardownStates.push(`reconcile-before:${placements.get(SESSION_ID)?.state ?? "missing"}`); + const reconciling = placements.startReconcile(input); + teardownStates.push(`reconcile-after:${reconciling.state}`); + expect(reconciling.turnClaim).toBeNull(); + return reconciling; + }, + }; + const stopTunnel = vi.fn(async () => { + const placement = placements.get(SESSION_ID); + teardownStates.push(`stop:${placement?.state ?? "missing"}`); + expect(placement).toMatchObject({ state: "draining", turnClaim: { owner: "worker" } }); + }); + const destroy = vi.fn(async () => { + teardownStates.push(`destroy:${placements.get(SESSION_ID)?.state ?? "missing"}`); + return attachedEnvironment(); + }); + const environments: WorkerTurnEnvironmentService = { + get: vi.fn(() => attachedEnvironment()), + acquireTurnCredential: vi.fn(async () => credential()), + acknowledgeCredentialDelivery: vi.fn(() => true), + startTunnel: vi.fn(async () => ({ + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + remoteSocketPath: "/worker/gateway.sock", + runWorkspaceCommand: vi.fn(async () => { + throw new Error("remote launch failed"); + }), + syncWorkspace: vi.fn(async () => { + throw new Error("unexpected workspace sync"); + }), + stop: vi.fn(async () => {}), + })), + stopTunnel, + destroy, + }; + const provider = createWorkerSessionTurnPlacementProvider({ + environments, + placements: observedPlacements, + }); + const runLocal = vi.fn(async () => ({ meta: { durationMs: 1 } })); + + await expect( + provider.executeTurn( + { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-failed", + }, + turn("run-failed"), + runLocal, + ), + ).rejects.toThrow("remote launch failed"); + expect(runLocal).not.toHaveBeenCalled(); + expect(placements.get(SESSION_ID)).toMatchObject({ + state: "failed", + turnClaim: null, + recoveryError: "remote launch failed", + }); + expect(stopTunnel).toHaveBeenCalledWith(ENVIRONMENT_ID, OWNER_EPOCH); + expect(destroy).toHaveBeenCalledWith(ENVIRONMENT_ID); + expect(teardownStates).toEqual([ + "stop:draining", + "destroy:draining", + "reconcile-before:draining", + "reconcile-after:reconciling", + ]); + }); + + it("launches only one worker loop for concurrent admission of the same run", async () => { + seedActivePlacement(); + const commandStarted = createDeferred(); + const commandFinished = createDeferred<{ + stdout: string; + stderr: string; + code: number; + signal: null; + killed: false; + termination: "exit"; + }>(); + const runWorkspaceCommand = vi.fn(() => { + commandStarted.resolve(); + return commandFinished.promise; + }); + const environments: WorkerTurnEnvironmentService = { + get: vi.fn(() => attachedEnvironment()), + acquireTurnCredential: vi.fn(async () => credential()), + acknowledgeCredentialDelivery: vi.fn(() => true), + startTunnel: vi.fn(async () => ({ + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + remoteSocketPath: "/worker/gateway.sock", + runWorkspaceCommand, + syncWorkspace: vi.fn(async () => { + throw new Error("unexpected workspace sync"); + }), + stop: vi.fn(async () => {}), + })), + stopTunnel: vi.fn(async () => {}), + destroy: vi.fn(async () => attachedEnvironment()), + }; + const provider = createWorkerSessionTurnPlacementProvider({ environments, placements }); + const claim = { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-overlap", + }; + const first = provider.executeTurn(claim, turn("run-overlap"), async () => ({ + meta: { durationMs: 1 }, + })); + await commandStarted.promise; + + await expect( + provider.executeTurn(claim, turn("run-overlap"), async () => ({ + meta: { durationMs: 1 }, + })), + ).rejects.toThrow("already has an active turn claim"); + expect(runWorkspaceCommand).toHaveBeenCalledOnce(); + + const completed = SessionManager.open(sessionFile); + const leafId = completed.appendMessage( + makeAgentAssistantMessage({ + content: [{ type: "text", text: "Only worker reply" }], + timestamp: 31, + }), + ); + createWorkerSessionPlacementGate(placements).updateAckCursors({ + sessionId: SESSION_ID, + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + runId: "run-overlap", + transcriptSeq: 1, + }); + commandFinished.resolve({ + stdout: JSON.stringify({ + status: "completed", + transcriptLeafId: leafId, + transcriptNextSeq: 2, + }), + stderr: "", + code: 0, + signal: null, + killed: false, + termination: "exit", + }); + await expect(first).resolves.toMatchObject({ payloads: [{ text: "Only worker reply" }] }); + expect(placements.get(SESSION_ID)?.turnClaim).toBeNull(); + }); + + it("keeps an active placement after an acknowledged turn failure and admits the next turn", async () => { + seedActivePlacement(); + const turnIds: string[] = []; + let launchCount = 0; + const stopTunnel = vi.fn(async () => {}); + const destroy = vi.fn(async () => attachedEnvironment()); + const environments: WorkerTurnEnvironmentService = { + get: vi.fn(() => attachedEnvironment()), + acquireTurnCredential: vi.fn(async () => credential(String(launchCount + 1).repeat(43))), + acknowledgeCredentialDelivery: vi.fn(() => true), + startTunnel: vi.fn(async () => ({ + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + remoteSocketPath: "/worker/gateway.sock", + runWorkspaceCommand: vi.fn(async (command): Promise => { + launchCount += 1; + const descriptor = parseWorkerLaunchDescriptor(JSON.parse(command.input ?? "")); + turnIds.push(descriptor.assignment.turnId); + if (launchCount === 1) { + return { + stdout: JSON.stringify({ status: "failed", reason: "turn-failed" }), + stderr: "", + code: 0, + signal: null, + killed: false, + termination: "exit", + }; + } + const completed = SessionManager.open(sessionFile); + const leafId = completed.appendMessage( + makeAgentAssistantMessage({ + content: [{ type: "text", text: "Recovered worker reply" }], + timestamp: 41, + }), + ); + createWorkerSessionPlacementGate(placements).updateAckCursors({ + sessionId: SESSION_ID, + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + runId: "run-model-recovered", + transcriptSeq: 1, + }); + return { + stdout: JSON.stringify({ + status: "completed", + transcriptLeafId: leafId, + transcriptNextSeq: 2, + }), + stderr: "", + code: 0, + signal: null, + killed: false, + termination: "exit", + }; + }), + syncWorkspace: vi.fn(async () => { + throw new Error("unexpected workspace sync"); + }), + stop: vi.fn(async () => {}), + })), + stopTunnel, + destroy, + }; + const provider = createWorkerSessionTurnPlacementProvider({ environments, placements }); + + await expect( + provider.executeTurn( + { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-model-failed", + }, + turn("run-model-failed"), + async () => ({ meta: { durationMs: 1 } }), + ), + ).rejects.toThrow("Cloud worker turn failed"); + expect(placements.get(SESSION_ID)).toMatchObject({ state: "active", turnClaim: null }); + + await expect( + provider.executeTurn( + { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-model-recovered", + }, + turn("run-model-recovered"), + async () => ({ meta: { durationMs: 1 } }), + ), + ).resolves.toMatchObject({ payloads: [{ text: "Recovered worker reply" }] }); + expect(turnIds).toHaveLength(2); + expect(turnIds[0]).not.toBe(turnIds[1]); + expect(stopTunnel).not.toHaveBeenCalled(); + expect(destroy).not.toHaveBeenCalled(); + expect(placements.get(SESSION_ID)).toMatchObject({ state: "active", turnClaim: null }); + }); + + it("redispatches a reclaimed placement before launching the worker turn", async () => { + const reclaimed = seedReclaimedPlacement(); + const runId = "run-reclaimed-worker"; + let redispatchCalls = 0; + const redispatchReclaimed: NonNullable< + WorkerTurnLauncherOptions["redispatchReclaimed"] + > = async (placement) => { + redispatchCalls += 1; + expect(placement).toEqual(reclaimed); + expect(placements.get(SESSION_ID)?.turnClaim).toBeNull(); + seedActivePlacement(); + const active = placements.get(SESSION_ID); + if (active?.state !== "active") { + throw new Error("expected active redispatched placement"); + } + return active; + }; + const runWorkspaceCommand = vi.fn(async (): Promise => { + expect(placements.get(SESSION_ID)).toMatchObject({ + state: "active", + turnClaim: { owner: "worker", runId }, + }); + const completed = SessionManager.open(sessionFile); + const leafId = completed.appendMessage( + makeAgentAssistantMessage({ + content: [{ type: "text", text: "Redispatched worker reply" }], + timestamp: 51, + }), + ); + createWorkerSessionPlacementGate(placements).updateAckCursors({ + sessionId: SESSION_ID, + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + runId, + transcriptSeq: 1, + }); + return { + stdout: JSON.stringify({ + status: "completed", + transcriptLeafId: leafId, + transcriptNextSeq: 2, + }), + stderr: "", + code: 0, + signal: null, + killed: false, + termination: "exit", + }; + }); + const environments: WorkerTurnEnvironmentService = { + get: vi.fn(() => attachedEnvironment()), + acquireTurnCredential: vi.fn(async () => credential()), + acknowledgeCredentialDelivery: vi.fn(() => true), + startTunnel: vi.fn(async () => ({ + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + remoteSocketPath: "/worker/gateway.sock", + runWorkspaceCommand, + syncWorkspace: vi.fn(async () => { + throw new Error("unexpected workspace sync"); + }), + stop: vi.fn(async () => {}), + })), + stopTunnel: vi.fn(async () => {}), + destroy: vi.fn(async () => attachedEnvironment()), + }; + const provider = createWorkerSessionTurnPlacementProvider({ + environments, + placements, + redispatchReclaimed, + }); + const runLocal = vi.fn(async () => ({ meta: { durationMs: 1 } })); + + const result = await provider.executeTurn( + { sessionId: SESSION_ID, sessionKey: SESSION_KEY, agentId: "main", runId }, + turn(runId), + runLocal, + ); + + expect(result.payloads).toEqual([{ text: "Redispatched worker reply" }]); + expect(redispatchCalls).toBe(1); + expect(runWorkspaceCommand).toHaveBeenCalledOnce(); + expect(runLocal).not.toHaveBeenCalled(); + expect(placements.get(SESSION_ID)).toMatchObject({ state: "active", turnClaim: null }); + }); + + it("rejects a reclaimed placement when redispatch is unavailable", async () => { + seedReclaimedPlacement(); + const provider = createWorkerSessionTurnPlacementProvider({ + environments: unusedEnvironments(), + placements, + }); + const runLocal = vi.fn(async () => ({ meta: { durationMs: 1 } })); + + await expect( + provider.executeTurn( + { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-reclaimed-unavailable", + }, + turn("run-reclaimed-unavailable"), + runLocal, + ), + ).rejects.toThrow("Reclaimed worker placement requires redispatch"); + expect(runLocal).not.toHaveBeenCalled(); + expect(placements.get(SESSION_ID)).toMatchObject({ state: "reclaimed", turnClaim: null }); + }); + + it("does not fall back locally when reclaimed redispatch fails", async () => { + seedReclaimedPlacement(); + const provider = createWorkerSessionTurnPlacementProvider({ + environments: unusedEnvironments(), + placements, + redispatchReclaimed: async () => { + throw new Error("reclaimed redispatch failed"); + }, + }); + const runLocal = vi.fn(async () => ({ meta: { durationMs: 1 } })); + + await expect( + provider.executeTurn( + { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-reclaimed-failed", + }, + turn("run-reclaimed-failed"), + runLocal, + ), + ).rejects.toThrow("reclaimed redispatch failed"); + expect(runLocal).not.toHaveBeenCalled(); + expect(placements.get(SESSION_ID)).toMatchObject({ state: "reclaimed", turnClaim: null }); + }); + + it("rejects non-active placement without falling back to the local loop", async () => { + placements.startDispatch({ + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + }); + const provider = createWorkerSessionTurnPlacementProvider({ + environments: unusedEnvironments(), + placements, + }); + const runLocal = vi.fn(async () => ({ meta: { durationMs: 1 } })); + + await expect( + provider.executeTurn( + { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-requested", + }, + turn("run-requested"), + runLocal, + ), + ).rejects.toThrow("Worker turn rejected in placement requested"); + expect(runLocal).not.toHaveBeenCalled(); + expect(placements.get(SESSION_ID)?.turnClaim).toBeNull(); + }); +}); diff --git a/src/gateway/worker-environments/worker-turn-launcher.ts b/src/gateway/worker-environments/worker-turn-launcher.ts new file mode 100644 index 000000000000..1bb5ee9d92f3 --- /dev/null +++ b/src/gateway/worker-environments/worker-turn-launcher.ts @@ -0,0 +1,464 @@ +import { randomUUID } from "node:crypto"; +import { truncateUtf16Safe } from "@openclaw/normalization-core/utf16-slice"; +import { mapThinkingLevelForProvider } from "../../agents/embedded-agent-runner/utils.js"; +import type { + LocalTurnPlacementClaim, + SessionPlacementAdmissionProvider, + SessionPlacementTurnParams, +} from "../../agents/session-placement-admission.js"; +import { convertToLlm } from "../../agents/sessions/messages.js"; +import { SessionManager } from "../../agents/sessions/session-manager.js"; +import { formatErrorMessage } from "../../infra/errors.js"; +import { redactSensitiveText } from "../../logging/redact.js"; +import { parseWorkerLaunchDescriptor } from "../../worker/launch-descriptor.js"; +import type { + WorkerSessionPlacementRecord, + WorkerSessionPlacementStore, + WorkerSessionTurnClaim, +} from "./placement-store.js"; +import type { WorkerEnvironmentService } from "./service.js"; +import { + assertSupportedTurn, + assistantText, + buildWorkerAgentMeta, + fitLaunchDescriptor, + parseRuntimeResult, + windowInitialMessages, +} from "./worker-turn-payload.js"; + +const WORKER_LAUNCH_SCRIPT = 'exec node "$HOME/.openclaw-worker/$1/openclaw.mjs" worker'; + +type WorkerTurnEnvironmentService = Pick< + WorkerEnvironmentService, + | "acknowledgeCredentialDelivery" + | "acquireTurnCredential" + | "destroy" + | "get" + | "startTunnel" + | "stopTunnel" +>; + +type ActiveWorkerPlacement = Extract; +type ReclaimedWorkerPlacement = Extract; + +type WorkerTurnLauncherOptions = { + admitNewPlacements?: boolean; + environments: WorkerTurnEnvironmentService; + placements: WorkerSessionPlacementStore; + redispatchReclaimed?: (placement: ReclaimedWorkerPlacement) => Promise; +}; + +class WorkerTurnExecutionError extends Error {} + +function required(value: string | undefined, field: string): string { + const normalized = value?.trim(); + if (!normalized) { + throw new Error(`Worker turn ${field} is required`); + } + return normalized; +} + +async function waitForTurnOperation(params: { + operation: Promise; + signal?: AbortSignal; + timeoutMs: number; +}): Promise { + const timeout = AbortSignal.timeout(params.timeoutMs); + const signal = params.signal ? AbortSignal.any([params.signal, timeout]) : timeout; + const abortError = () => + signal.reason instanceof Error + ? signal.reason + : new Error("Cloud worker operation aborted", { cause: signal.reason }); + if (signal.aborted) { + throw abortError(); + } + return await new Promise((resolve, reject) => { + const onAbort = () => reject(abortError()); + signal.addEventListener("abort", onAbort, { once: true }); + params.operation.then(resolve, reject).finally(() => { + signal.removeEventListener("abort", onAbort); + }); + }); +} + +function resolvePlacementIdentity( + claim: LocalTurnPlacementClaim, + placement: WorkerSessionPlacementRecord | undefined, +) { + return { + sessionId: claim.sessionId, + agentId: placement?.agentId ?? required(claim.agentId, "agent id"), + sessionKey: placement?.sessionKey ?? required(claim.sessionKey, "session key"), + }; +} + +function requireActivePlacement(placement: WorkerSessionPlacementRecord): ActiveWorkerPlacement { + if ( + placement.state !== "active" || + !placement.remoteWorkspaceDir || + !placement.workerBundleHash + ) { + throw new Error(`Worker turn rejected in placement ${placement.state}`); + } + return placement; +} + +function releaseClaimIfOwned( + placements: WorkerSessionPlacementStore, + turnClaim: WorkerSessionTurnClaim, +): void { + if (placements.validateTurnClaim(turnClaim)) { + placements.releaseTurn(turnClaim); + } +} + +async function executeLocalTurn(params: { + claim: LocalTurnPlacementClaim; + placements: WorkerSessionPlacementStore; + runLocal: () => Promise; +}): Promise { + const current = params.placements.get(params.claim.sessionId); + const turnClaim = params.placements.claimTurn({ + ...resolvePlacementIdentity(params.claim, current), + claimId: randomUUID(), + runId: params.claim.runId, + owner: { kind: "local" }, + }); + try { + return await params.runLocal(); + } finally { + releaseClaimIfOwned(params.placements, turnClaim); + } +} + +function recoveryError(error: unknown): string { + const message = redactSensitiveText(formatErrorMessage(error), { mode: "tools" }) + .replace(/\s+/gu, " ") + .trim(); + return truncateUtf16Safe(message || "cloud worker turn failed", 1_024); +} + +async function failHandedOffTurn(params: { + environments: WorkerTurnEnvironmentService; + placements: WorkerSessionPlacementStore; + placement: ActiveWorkerPlacement; + error: unknown; +}): Promise { + const primaryFailure = recoveryError(params.error); + const failures = [primaryFailure]; + let draining: WorkerSessionPlacementRecord; + try { + draining = params.placements.startDrain({ + sessionId: params.placement.sessionId, + environmentId: params.placement.environmentId, + ownerEpoch: params.placement.activeOwnerEpoch, + expectedGeneration: params.placement.generation, + }); + } catch { + // Exact drain ownership failed. Do not tear down an environment that may + // now belong to a newer placement generation. + return; + } + if (draining.state !== "draining") { + return; + } + try { + await params.environments.stopTunnel( + params.placement.environmentId, + params.placement.activeOwnerEpoch, + ); + } catch (error) { + failures.push(`tunnel stop: ${recoveryError(error)}`); + } + try { + await params.environments.destroy(params.placement.environmentId); + } catch (error) { + failures.push(`environment destroy: ${recoveryError(error)}`); + } + try { + // Both teardown calls returned through the environment queue. Fence stale + // worker RPC durably now; failed teardown remains eligible for retry. + const reconciling = params.placements.startReconcile({ + sessionId: draining.sessionId, + environmentId: draining.environmentId, + ownerEpoch: draining.activeOwnerEpoch, + expectedGeneration: draining.generation, + }); + if (reconciling.state !== "reconciling") { + return; + } + params.placements.fail({ + sessionId: reconciling.sessionId, + expectedGeneration: reconciling.generation, + recoveryError: truncateUtf16Safe(failures.join("; "), 1_024), + }); + } catch { + // Leave the durable draining or reconciling row for startup reconciliation. + } +} + +async function executeWorkerTurn(params: { + environments: WorkerTurnEnvironmentService; + onHandoff: () => void; + placement: ActiveWorkerPlacement; + placements: WorkerSessionPlacementStore; + turn: SessionPlacementTurnParams; +}) { + const { placement, turn } = params; + const modelRef = assertSupportedTurn(turn); + const environment = params.environments.get(placement.environmentId); + if ( + !environment || + environment.state !== "attached" || + environment.ownerEpoch !== placement.activeOwnerEpoch || + environment.bootstrapReceipt?.bundleHash !== placement.workerBundleHash || + environment.attachedSessionIds.length !== 1 || + environment.attachedSessionIds[0] !== placement.sessionId + ) { + throw new Error("Active worker placement does not match its attached environment"); + } + + const startedAt = Date.now(); + turn.onExecutionStarted?.({ lifecycleGeneration: turn.lifecycleGeneration }); + turn.onExecutionPhase?.({ phase: "runner_entered", backend: "cloud-worker" }); + const manager = SessionManager.open(turn.sessionFile); + const userMessageAlreadyPersisted = + turn.suppressNextUserMessagePersistence === true || + turn.userTurnTranscriptRecorder?.hasPersisted() === true; + const contextMessages = convertToLlm(manager.buildSessionContext().messages); + const leaf = manager.getLeafEntry(); + const initialMessages = windowInitialMessages( + userMessageAlreadyPersisted && leaf?.type === "message" && leaf.message.role === "user" + ? contextMessages.slice(0, -1) + : contextMessages, + ); + let baseLeafId = manager.getLeafId(); + if (!userMessageAlreadyPersisted) { + const persisted = turn.userTurnTranscriptRecorder + ? await turn.userTurnTranscriptRecorder.persistApproved({ cwd: turn.workspaceDir }) + : undefined; + if (persisted) { + baseLeafId = persisted.messageId; + turn.userTurnTranscriptRecorder?.markRuntimePersisted(persisted.message); + turn.onUserMessagePersisted?.(persisted.message); + } else if (turn.userTurnTranscriptRecorder?.hasPersisted()) { + baseLeafId = SessionManager.open(turn.sessionFile).getLeafId(); + } else if (!turn.userTurnTranscriptRecorder) { + const message = { + role: "user" as const, + content: [{ type: "text" as const, text: turn.transcriptPrompt ?? turn.prompt }], + timestamp: Date.now(), + }; + baseLeafId = manager.appendMessage(message); + turn.onUserMessagePersisted?.(message); + } else { + throw new Error("Cloud worker turn could not persist its canonical user message"); + } + } + turn.onExecutionPhase?.({ + phase: "model_resolution", + backend: "cloud-worker", + provider: modelRef.provider, + model: modelRef.model, + }); + + const credential = await params.environments.acquireTurnCredential({ + environmentId: placement.environmentId, + ownerEpoch: placement.activeOwnerEpoch, + sessionId: placement.sessionId, + }); + const tunnel = await waitForTurnOperation({ + operation: params.environments.startTunnel({ + environmentId: placement.environmentId, + ownerEpoch: placement.activeOwnerEpoch, + }), + ...(turn.abortSignal ? { signal: turn.abortSignal } : {}), + timeoutMs: turn.timeoutMs, + }); + const reasoning = mapThinkingLevelForProvider(turn.thinkLevel); + const descriptor = fitLaunchDescriptor( + (windowedMessages) => + parseWorkerLaunchDescriptor({ + version: 1, + socketPath: tunnel.remoteSocketPath, + admission: { + environmentId: placement.environmentId, + credential: credential.credential, + sessionId: placement.sessionId, + ownerEpoch: placement.activeOwnerEpoch, + rpcSetVersion: credential.rpcSetVersion, + handshake: environment.bootstrapReceipt, + }, + assignment: { + runId: turn.runId, + turnId: randomUUID(), + prompt: turn.prompt, + suppressPromptTranscript: true, + workspaceDir: placement.remoteWorkspaceDir, + modelRef, + inferenceOptions: reasoning ? { reasoning } : {}, + ...(turn.extraSystemPrompt === undefined ? {} : { systemPrompt: turn.extraSystemPrompt }), + initialMessages: windowedMessages, + transcript: { + baseLeafId, + nextSeq: (placement.lastTranscriptAckCursor ?? 0) + 1, + }, + liveEvents: { + ackedSeq: placement.lastLiveEventAckCursor ?? 0, + nextSeq: (placement.lastLiveEventAckCursor ?? 0) + 1, + }, + }, + }), + initialMessages, + ); + turn.userTurnTranscriptRecorder?.markSentToProvider?.(); + turn.onExecutionPhase?.({ phase: "attempt_dispatch", backend: "cloud-worker" }); + const handoffAbort = new AbortController(); + params.onHandoff(); + const processPromise = tunnel.runWorkspaceCommand({ + argv: ["sh", "-c", WORKER_LAUNCH_SCRIPT, "openclaw-worker", placement.workerBundleHash], + input: JSON.stringify(descriptor), + timeoutMs: turn.timeoutMs, + signal: turn.abortSignal + ? AbortSignal.any([turn.abortSignal, handoffAbort.signal]) + : handoffAbort.signal, + }); + turn.onExecutionPhase?.({ phase: "process_spawned", backend: "cloud-worker" }); + let credentialDelivered: boolean; + try { + credentialDelivered = params.environments.acknowledgeCredentialDelivery(credential); + } catch (error) { + handoffAbort.abort(); + await processPromise.catch(() => undefined); + throw new Error("Cloud worker credential handoff failed", { cause: error }); + } + if (!credentialDelivered) { + handoffAbort.abort(); + await processPromise.catch(() => undefined); + throw new Error("Cloud worker credential owner changed during process handoff"); + } + const processResult = await processPromise; + if (processResult.code !== 0 || processResult.signal !== null || processResult.killed) { + throw new Error("Cloud worker process failed before completing the turn"); + } + const runtimeResult = parseRuntimeResult(processResult.stdout); + if (runtimeResult.status === "fenced") { + throw new Error(`Cloud worker turn was fenced: ${runtimeResult.reason}`); + } + if (runtimeResult.status === "failed") { + throw new WorkerTurnExecutionError("Cloud worker turn failed"); + } + + const completed = SessionManager.open(turn.sessionFile); + const currentPlacement = params.placements.get(placement.sessionId); + if ( + runtimeResult.transcriptLeafId !== completed.getLeafId() || + runtimeResult.transcriptNextSeq !== (currentPlacement?.lastTranscriptAckCursor ?? 0) + 1 + ) { + throw new Error("Cloud worker result does not match its committed transcript acknowledgement"); + } + const terminal = runtimeResult.transcriptLeafId + ? completed.getEntry(runtimeResult.transcriptLeafId) + : undefined; + if (!terminal || terminal.type !== "message" || terminal.message.role !== "assistant") { + throw new Error("Cloud worker completed without a terminal assistant transcript message"); + } + const text = assistantText(terminal.message); + const baseIndex = completed.getBranch().findIndex((entry) => entry.id === baseLeafId); + const workerMessages = completed + .getBranch() + .slice(baseIndex + 1) + .flatMap((entry) => (entry.type === "message" ? [entry.message] : [])); + return { + ...(text ? { payloads: [{ text }] } : {}), + meta: { + durationMs: Date.now() - startedAt, + agentMeta: { + sessionId: placement.sessionId, + sessionFile: turn.sessionFile, + ...buildWorkerAgentMeta({ messages: workerMessages, modelRef }), + }, + stopReason: terminal.message.stopReason, + }, + }; +} + +export function createWorkerSessionTurnPlacementProvider( + options: WorkerTurnLauncherOptions, +): SessionPlacementAdmissionProvider { + return { + async executeLocalTurn(claim: LocalTurnPlacementClaim, runLocal: () => Promise) { + if (!options.placements.get(claim.sessionId) && options.admitNewPlacements === false) { + return await runLocal(); + } + return await executeLocalTurn({ claim, placements: options.placements, runLocal }); + }, + async executeTurn(claim, turn, runLocal) { + const current = options.placements.get(claim.sessionId); + if ( + !current && + (options.admitNewPlacements === false || + (turn.modelRun === true && !claim.sessionKey?.trim())) + ) { + return await runLocal(); + } + if (!current || current.state === "local") { + return await executeLocalTurn({ claim, placements: options.placements, runLocal }); + } + let routablePlacement = current; + if (routablePlacement.state === "reclaimed") { + if (!options.redispatchReclaimed) { + throw new Error("Reclaimed worker placement requires redispatch"); + } + routablePlacement = await options.redispatchReclaimed(routablePlacement); + } + const identity = resolvePlacementIdentity(claim, routablePlacement); + const placement = requireActivePlacement(routablePlacement); + const turnClaim = options.placements.claimTurn({ + ...identity, + claimId: randomUUID(), + runId: claim.runId, + owner: { + kind: "worker", + environmentId: placement.environmentId, + ownerEpoch: placement.activeOwnerEpoch, + }, + }); + let handedOff = false; + try { + const result = await executeWorkerTurn({ + environments: options.environments, + onHandoff: () => { + handedOff = true; + }, + placement, + placements: options.placements, + turn, + }); + if (!options.placements.validateTurnClaim(turnClaim)) { + throw new Error("Cloud worker turn ownership changed before result reconciliation"); + } + options.placements.releaseTurn(turnClaim); + return result; + } catch (error) { + if (error instanceof WorkerTurnExecutionError) { + if (options.placements.validateTurnClaim(turnClaim)) { + options.placements.releaseTurn(turnClaim); + throw error; + } + } + if (handedOff) { + await failHandedOffTurn({ + environments: options.environments, + placements: options.placements, + placement, + error, + }); + } else { + releaseClaimIfOwned(options.placements, turnClaim); + } + throw error; + } + }, + }; +} diff --git a/src/gateway/worker-environments/worker-turn-payload.ts b/src/gateway/worker-environments/worker-turn-payload.ts new file mode 100644 index 000000000000..24e3ff515f42 --- /dev/null +++ b/src/gateway/worker-environments/worker-turn-payload.ts @@ -0,0 +1,197 @@ +import type { WorkerTranscriptMessage } from "../../../packages/gateway-protocol/src/schema/worker-admission.js"; +import { + WORKER_INFERENCE_MAX_CONTEXT_MESSAGES, + WORKER_PROTOCOL_MAX_INFERENCE_PAYLOAD_BYTES, +} from "../../../packages/gateway-protocol/src/schema/worker-inference.js"; +import { + isDefaultAgentRuntimeId, + normalizeOptionalAgentRuntimeId, + OPENCLAW_AGENT_RUNTIME_ID, +} from "../../agents/agent-runtime-id.js"; +import { + buildUsageAgentMetaFields, + resolveReportedModelRef, +} from "../../agents/embedded-agent-runner/run/helpers.js"; +import { + createUsageAccumulator, + mergeUsageIntoAccumulator, +} from "../../agents/embedded-agent-runner/usage-accumulator.js"; +import { resolveDefaultModelForAgent } from "../../agents/model-selection-config.js"; +import type { AgentMessage } from "../../agents/runtime/index.js"; +import type { SessionPlacementTurnParams } from "../../agents/session-placement-admission.js"; +import { resolveEffectiveAgentRuntime } from "../../agents/thinking-runtime.js"; +import { hasNonzeroUsage, normalizeUsage } from "../../agents/usage.js"; +import type { WorkerLaunchDescriptor } from "../../worker/launch-descriptor.js"; +import { toWorkerTranscriptMessage } from "../../worker/transcript-message.js"; +import type { WorkerRuntimeResult } from "../../worker/worker.runtime.js"; + +export function windowInitialMessages(messages: AgentMessage[]): WorkerTranscriptMessage[] { + const projected = messages.flatMap((message) => { + const value = toWorkerTranscriptMessage(message); + return value ? [value] : []; + }); + if (projected.length <= WORKER_INFERENCE_MAX_CONTEXT_MESSAGES) { + return projected; + } + const minimumStart = projected.length - WORKER_INFERENCE_MAX_CONTEXT_MESSAGES; + const completeTurnStart = projected.findIndex( + (message, index) => index >= minimumStart && message.role === "user", + ); + if (completeTurnStart < 0) { + throw new Error("Worker turn transcript has no complete context window"); + } + return projected.slice(completeTurnStart); +} + +export function fitLaunchDescriptor( + build: (initialMessages: WorkerTranscriptMessage[]) => WorkerLaunchDescriptor, + messages: WorkerTranscriptMessage[], +): WorkerLaunchDescriptor { + let initialMessages = messages; + while (true) { + const descriptor = build(initialMessages); + if ( + Buffer.byteLength(JSON.stringify(descriptor), "utf8") <= + WORKER_PROTOCOL_MAX_INFERENCE_PAYLOAD_BYTES + ) { + return descriptor; + } + const nextTurn = initialMessages.findIndex( + (message, index) => index > 0 && message.role === "user", + ); + if (nextTurn < 0) { + throw new Error("Worker turn context exceeds the launch descriptor payload limit"); + } + initialMessages = initialMessages.slice(nextTurn); + } +} + +export function parseRuntimeResult(stdout: string): WorkerRuntimeResult { + let value: unknown; + try { + value = JSON.parse(stdout.trim()) as unknown; + } catch (error) { + throw new Error("Worker process returned invalid output", { cause: error }); + } + if (!value || typeof value !== "object" || Array.isArray(value)) { + throw new Error("Worker process returned invalid output"); + } + const result = value as Record; + if ( + result.status === "failed" && + result.reason === "turn-failed" && + Object.keys(result).every((key) => ["status", "reason"].includes(key)) + ) { + return result as WorkerRuntimeResult; + } + if ( + result.status === "completed" && + (result.transcriptLeafId === null || typeof result.transcriptLeafId === "string") && + typeof result.transcriptNextSeq === "number" && + Number.isSafeInteger(result.transcriptNextSeq) && + result.transcriptNextSeq >= 1 && + Object.keys(result).every((key) => + ["status", "transcriptLeafId", "transcriptNextSeq"].includes(key), + ) + ) { + return result as WorkerRuntimeResult; + } + if ( + result.status === "fenced" && + (result.reason === "credential-replaced" || result.reason === "owner-epoch-mismatch") && + Object.keys(result).every((key) => ["status", "reason"].includes(key)) + ) { + return result as WorkerRuntimeResult; + } + throw new Error("Worker process returned invalid output"); +} + +export function assistantText(message: AgentMessage): string { + if (message.role !== "assistant") { + return ""; + } + return message.content.flatMap((part) => (part.type === "text" ? [part.text] : [])).join(""); +} + +export function buildWorkerAgentMeta(params: { + messages: AgentMessage[]; + modelRef: { provider: string; model: string }; +}) { + const usageAccumulator = createUsageAccumulator(); + const assistants = params.messages.filter( + (message): message is Extract => + message.role === "assistant", + ); + let lastRunPromptUsage: ReturnType; + for (const assistant of assistants) { + const usage = normalizeUsage(assistant.usage); + mergeUsageIntoAccumulator(usageAccumulator, usage); + if (hasNonzeroUsage(usage)) { + lastRunPromptUsage = usage; + } + } + const lastAssistant = assistants.at(-1); + const usageMeta = buildUsageAgentMetaFields({ + usageAccumulator, + lastAssistantUsage: lastAssistant?.usage, + lastRunPromptUsage, + lastTurnTotal: lastRunPromptUsage?.total, + }); + const reportedModelRef = resolveReportedModelRef({ + ...params.modelRef, + assistant: lastAssistant, + }); + return { + provider: reportedModelRef.provider, + model: reportedModelRef.model, + usage: usageMeta.usage, + lastCallUsage: usageMeta.lastCallUsage, + promptTokens: usageMeta.promptTokens, + }; +} + +function resolveTurnModelRef(params: SessionPlacementTurnParams): { + provider: string; + model: string; +} { + const explicitProvider = params.provider?.trim(); + const explicitModel = params.model?.trim(); + const defaults = + explicitProvider && explicitModel + ? undefined + : resolveDefaultModelForAgent({ cfg: params.config ?? {}, agentId: params.agentId }); + return { + provider: explicitProvider ?? defaults?.provider ?? "", + model: explicitModel ?? defaults?.model ?? "", + }; +} + +export function assertSupportedTurn(params: SessionPlacementTurnParams): { + provider: string; + model: string; +} { + if (params.images?.length || params.imageOrder?.length) { + throw new Error("Cloud worker turns do not yet support current-turn image input"); + } + if (params.clientTools?.length) { + throw new Error("Cloud worker turns do not support client-provided tools"); + } + const modelRef = resolveTurnModelRef(params); + const explicitRuntime = + normalizeOptionalAgentRuntimeId(params.agentHarnessId) ?? + normalizeOptionalAgentRuntimeId(params.agentHarnessRuntimeOverride); + const runtime = + explicitRuntime && !isDefaultAgentRuntimeId(explicitRuntime) + ? explicitRuntime + : resolveEffectiveAgentRuntime({ + cfg: params.config ?? {}, + provider: modelRef.provider, + modelId: modelRef.model, + agentId: params.agentId, + sessionKey: params.sessionKey, + }); + if (runtime !== OPENCLAW_AGENT_RUNTIME_ID) { + throw new Error(`Cloud worker turns require the OpenClaw runtime, not ${runtime}`); + } + return modelRef; +} diff --git a/src/gateway/worker-environments/workspace-sync-local.test.ts b/src/gateway/worker-environments/workspace-sync-local.test.ts new file mode 100644 index 000000000000..fab2781ebecb --- /dev/null +++ b/src/gateway/worker-environments/workspace-sync-local.test.ts @@ -0,0 +1,57 @@ +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { describe, expect, it } from "vitest"; +import { runLocalCommandToFile } from "./workspace-sync-local.js"; + +async function waitForFile(filePath: string): Promise { + const deadline = Date.now() + 5_000; + while (Date.now() < deadline) { + try { + await fs.access(filePath); + return; + } catch { + await new Promise((resolve) => { + setTimeout(resolve, 10); + }); + } + } + throw new Error(`Timed out waiting for ${filePath}`); +} + +describe("runLocalCommandToFile", () => { + it("force-kills a command that ignores abort termination", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-sync-")); + const outputPath = path.join(root, "output"); + const readyPath = path.join(root, "ready"); + const controller = new AbortController(); + const operation = runLocalCommandToFile({ + argv: [ + process.execPath, + "-e", + [ + 'const fs = require("node:fs");', + 'process.on("SIGTERM", () => {});', + 'fs.writeFileSync(process.argv[1], "ready");', + "setInterval(() => {}, 1000);", + ].join(""), + readyPath, + ], + outputPath, + signal: controller.signal, + timeoutMs: 10_000, + }); + + try { + await waitForFile(readyPath); + const abortedAt = Date.now(); + controller.abort(); + await expect(operation).rejects.toThrow("Worker workspace file enumeration was aborted"); + expect(Date.now() - abortedAt).toBeLessThan(3_000); + } finally { + controller.abort(); + await operation.catch(() => undefined); + await fs.rm(root, { recursive: true, force: true }); + } + }); +}); diff --git a/src/gateway/worker-environments/workspace-sync-local.ts b/src/gateway/worker-environments/workspace-sync-local.ts new file mode 100644 index 000000000000..acb9f9e0e355 --- /dev/null +++ b/src/gateway/worker-environments/workspace-sync-local.ts @@ -0,0 +1,235 @@ +import { spawn } from "node:child_process"; +import { createReadStream } from "node:fs"; +import fs from "node:fs/promises"; +import path from "node:path"; +import { sliceUtf16Safe } from "@openclaw/normalization-core/utf16-slice"; +import { killProcessTree } from "../../process/kill-tree.js"; +import { workerSshCommandOptions } from "./ssh.js"; + +const STDERR_LIMIT = 4_096; +const COMMAND_KILL_GRACE_MS = 300; +const COMMAND_CLOSE_GRACE_MS = 1_000; + +function validateGitRelativePath(file: string): string { + if ( + !file || + path.posix.isAbsolute(file) || + path.posix.normalize(file) !== file || + file === ".." || + file.startsWith("../") + ) { + throw new Error("Worker workspace git file list contains an unsafe path"); + } + return file; +} + +async function* readNulFile(filePath: string): AsyncGenerator { + let pending = Buffer.alloc(0); + for await (const value of createReadStream(filePath)) { + const chunk = Buffer.isBuffer(value) ? value : Buffer.from(value); + const buffer = pending.length === 0 ? chunk : Buffer.concat([pending, chunk]); + let offset = 0; + for (;;) { + const separator = buffer.indexOf(0, offset); + if (separator < 0) { + break; + } + yield validateGitRelativePath(buffer.subarray(offset, separator).toString("utf8")); + offset = separator + 1; + } + pending = Buffer.from(buffer.subarray(offset)); + } + if (pending.length > 0) { + throw new Error("Worker workspace git file list is not NUL terminated"); + } +} + +export async function runLocalCommandToFile(params: { + argv: string[]; + inputPath?: string; + outputPath: string; + signal: AbortSignal; + timeoutMs: number; +}): Promise { + const [command, ...args] = params.argv; + if (!command) { + throw new Error("Worker workspace command requires an executable"); + } + const output = await fs.open(params.outputPath, "wx", 0o600); + const input = params.inputPath ? await fs.open(params.inputPath, "r") : undefined; + let stderr = ""; + let timer: ReturnType | undefined; + let terminationTimer: ReturnType | undefined; + let abort: (() => void) | undefined; + try { + if (params.signal.aborted) { + throw new Error("Worker workspace file enumeration was aborted"); + } + const child = spawn(command, args, { + env: workerSshCommandOptions({ timeoutMs: params.timeoutMs }).baseEnv, + stdio: [input?.fd ?? "ignore", output.fd, "pipe"], + ...(process.platform !== "win32" ? { detached: true } : {}), + windowsHide: true, + }); + const childStderr = child.stderr; + if (!childStderr) { + throw new Error("Worker workspace command has no stderr pipe"); + } + childStderr.setEncoding("utf8"); + childStderr.on("data", (chunk: string) => { + stderr = sliceUtf16Safe(`${stderr}${chunk}`, -STDERR_LIMIT); + }); + const result = await new Promise<{ code: number | null; error?: Error }>((resolve) => { + let settled = false; + const finish = (value: { code: number | null; error?: Error }) => { + if (settled) { + return; + } + settled = true; + resolve(value); + }; + let terminationStarted = false; + const terminate = () => { + if (settled || terminationStarted) { + return; + } + terminationStarted = true; + const pid = child.pid; + if (typeof pid === "number" && pid > 0) { + killProcessTree(pid, { graceMs: COMMAND_KILL_GRACE_MS }); + } else { + child.kill("SIGTERM"); + } + // A descendant can retain stderr even after the direct child exits. Bound + // shutdown so placement replacement cannot wait forever on that pipe. + terminationTimer = setTimeout(() => { + if (typeof pid === "number" && pid > 0) { + killProcessTree(pid, { force: true }); + } else { + child.kill("SIGKILL"); + } + childStderr.destroy(); + finish({ code: null }); + }, COMMAND_KILL_GRACE_MS + COMMAND_CLOSE_GRACE_MS); + terminationTimer.unref?.(); + }; + child.once("error", (error) => finish({ code: null, error })); + child.once("close", (code) => finish({ code })); + abort = terminate; + params.signal.addEventListener("abort", abort, { once: true }); + timer = setTimeout(terminate, params.timeoutMs); + timer.unref?.(); + if (params.signal.aborted) { + terminate(); + } + }); + if (result.error) { + throw result.error; + } + if (params.signal.aborted) { + throw new Error("Worker workspace file enumeration was aborted"); + } + if (result.code !== 0) { + throw new Error( + stderr.trim() + ? `Worker workspace file enumeration failed: ${stderr.trim()}` + : "Worker workspace file enumeration failed", + ); + } + } finally { + clearTimeout(timer); + clearTimeout(terminationTimer); + if (abort) { + params.signal.removeEventListener("abort", abort); + } + await output.close(); + await input?.close(); + } +} + +function hasErrorCode(error: unknown, code: string): boolean { + return ( + typeof error === "object" && + error !== null && + "code" in error && + typeof error.code === "string" && + error.code === code + ); +} + +export async function writeEligibleGitFiles(params: { + gitRoot: string; + eligiblePath: string; + ignoredPath: string; + selectedPath: string; + outputPath: string; +}): Promise { + const output = await fs.open(params.outputPath, "wx", 0o600); + const canonicalRoot = await fs.realpath(params.gitRoot); + let buffered: string[] = []; + let bufferedBytes = 0; + const flush = async () => { + if (buffered.length === 0) { + return; + } + await output.write(buffered.join("")); + buffered = []; + bufferedBytes = 0; + }; + const appendIfTransferable = async (file: string) => { + const absolute = path.join(canonicalRoot, file); + const stats = await fs.lstat(absolute).catch((error: unknown) => { + if (hasErrorCode(error, "ENOENT")) { + return undefined; + } + throw error; + }); + // Gitlinks are directories. Keep their commit in the base repository without + // recursively copying nested repositories or their credential-bearing metadata. + if (!stats || (!stats.isFile() && !stats.isSymbolicLink())) { + return; + } + if (stats.isSymbolicLink()) { + // Mirrors the remote manifest guard, but before transfer: macOS openrsync + // stat-fails escaping links with an opaque error instead of copying them. + const target = await fs.readlink(absolute); + const resolvedTarget = path.resolve(path.dirname(absolute), target); + if ( + resolvedTarget !== canonicalRoot && + !resolvedTarget.startsWith(canonicalRoot + path.sep) + ) { + throw new Error(`worker workspace symlink escapes the sync root: ${file}`); + } + } + const record = `${file}\0`; + buffered.push(record); + bufferedBytes += Buffer.byteLength(record); + if (bufferedBytes >= 64 * 1024) { + await flush(); + } + }; + try { + for await (const file of readNulFile(params.eligiblePath)) { + await appendIfTransferable(file); + } + const ignored = readNulFile(params.ignoredPath)[Symbol.asyncIterator](); + const selected = readNulFile(params.selectedPath)[Symbol.asyncIterator](); + let ignoredItem = await ignored.next(); + let selectedItem = await selected.next(); + while (!ignoredItem.done && !selectedItem.done) { + const order = Buffer.compare(Buffer.from(ignoredItem.value), Buffer.from(selectedItem.value)); + if (order === 0) { + await appendIfTransferable(ignoredItem.value); + ignoredItem = await ignored.next(); + selectedItem = await selected.next(); + } else if (order < 0) { + ignoredItem = await ignored.next(); + } else { + selectedItem = await selected.next(); + } + } + await flush(); + } finally { + await output.close(); + } +} diff --git a/src/gateway/worker-environments/workspace-sync-scripts.ts b/src/gateway/worker-environments/workspace-sync-scripts.ts new file mode 100644 index 000000000000..15f034e0716f --- /dev/null +++ b/src/gateway/worker-environments/workspace-sync-scripts.ts @@ -0,0 +1,185 @@ +export const REMOTE_WORKSPACE_SETUP_SCRIPT = String.raw`set -eu +relative=$1 +root=$HOME/.openclaw-worker + +ensure_private_directory() { + directory=$1 + if [ -e "$directory" ] || [ -L "$directory" ]; then + if [ ! -d "$directory" ] || [ -L "$directory" ]; then + printf '%s\n' 'unsafe worker workspace directory' >&2 + exit 2 + fi + else + mkdir "$directory" + fi + chmod 700 "$directory" +} + +ensure_private_directory "$root" +current=$root +old_ifs=$IFS +IFS=/ +set -- $relative +IFS=$old_ifs +for segment in "$@"; do + current=$current/$segment + ensure_private_directory "$current" +done +cd "$current" +find . -mindepth 1 -maxdepth 1 -exec rm -rf -- {} + +pwd -P +`; + +export const REMOTE_GIT_WORKSPACE_SETUP_SCRIPT = String.raw`set -eu +workspace=$1 +pack=$2 +base=$3 +author_name=$4 +author_email=$5 +cd "$workspace" +if ! command -v git >/dev/null 2>&1; then + printf '%s\n' 'git is required for a git worker workspace' >&2 + exit 2 +fi +case ${"${"}#base} in + 40) git init -q . ;; + 64) git init -q --object-format=sha256 . ;; + *) printf '%s\n' 'invalid worker git base object id' >&2; exit 2 ;; +esac +git index-pack --stdin < "$pack" >/dev/null +printf '%s\n' "$base" > .git/shallow +actual=$(git rev-parse --verify "$base^{commit}") +if [ "$actual" != "$base" ]; then + printf '%s\n' 'worker git base does not match the synced pack' >&2 + exit 2 +fi +git update-ref refs/heads/openclaw-worker "$base" +git symbolic-ref HEAD refs/heads/openclaw-worker +git read-tree "$base" +git ls-files --stage -z | node -e ' +const childProcess = require("node:child_process"); +const chunks = []; +process.stdin.on("data", (chunk) => chunks.push(chunk)); +process.stdin.on("end", () => { + const paths = Buffer.concat(chunks) + .toString("utf8") + .split("\0") + .filter(Boolean) + .flatMap((record) => { + const separator = record.indexOf("\t"); + return separator >= 0 && record.startsWith("160000 ") ? [record.slice(separator + 1)] : []; + }); + if (paths.length > 0) { + childProcess.execFileSync("git", ["update-index", "--skip-worktree", "--", ...paths]); + } +});' +rm -f -- "$pack" +if [ -n "$author_name" ]; then git config user.name "$author_name"; fi +if [ -n "$author_email" ]; then git config user.email "$author_email"; fi +`; + +export const REMOTE_WORKSPACE_MANIFEST_JS = String.raw`const crypto = require("node:crypto"); +const fs = require("node:fs"); +const path = require("node:path"); +const root = fs.realpathSync(process.argv[1]); +const baseCommit = process.argv[2] || null; +const entries = []; +function fail(message) { + throw new Error(message); +} +function walk(relativeDirectory) { + const absoluteDirectory = relativeDirectory ? path.join(root, relativeDirectory) : root; + for (const name of fs.readdirSync(absoluteDirectory).sort()) { + if (!relativeDirectory && name === ".git") { + continue; + } + const relative = relativeDirectory ? relativeDirectory + "/" + name : name; + const absolute = path.join(root, relative); + const stats = fs.lstatSync(absolute); + const mode = stats.mode & 0o777; + if (stats.isDirectory()) { + entries.push({ path: relative, type: "directory", mode }); + walk(relative); + } else if (stats.isFile()) { + entries.push({ + path: relative, + type: "file", + mode, + size: stats.size, + sha256: null, + }); + } else if (stats.isSymbolicLink()) { + const target = fs.readlinkSync(absolute); + const resolvedTarget = path.resolve(path.dirname(absolute), target); + if (resolvedTarget !== root && !resolvedTarget.startsWith(root + path.sep)) { + fail("worker workspace symlink escapes the sync root: " + relative); + } + entries.push({ path: relative, type: "symlink", mode, target }); + } else { + fail("unsupported worker workspace entry: " + relative); + } + } +} +async function hashFiles() { + for (const entry of entries) { + if (entry.type !== "file") { + continue; + } + const hash = crypto.createHash("sha256"); + const stream = fs.createReadStream(path.join(root, entry.path)); + for await (const chunk of stream) { + hash.update(chunk); + } + entry.sha256 = hash.digest("hex"); + } +} +function ensurePrivateDirectory(directory) { + try { + const stats = fs.lstatSync(directory); + if (stats.isSymbolicLink() || !stats.isDirectory()) { + fail("unsafe worker manifest directory"); + } + } catch (error) { + if (error && error.code === "ENOENT") { + fs.mkdirSync(directory, { mode: 0o700 }); + } else { + throw error; + } + } + fs.chmodSync(directory, 0o700); +} +async function main() { + walk(""); + await hashFiles(); + const manifest = JSON.stringify({ version: 1, baseCommit, entries }); + const digest = crypto.createHash("sha256").update(manifest).digest("hex"); + const workerRoot = path.join(process.env.HOME, ".openclaw-worker"); + const manifestRoot = path.join(workerRoot, "manifests"); + ensurePrivateDirectory(workerRoot); + ensurePrivateDirectory(manifestRoot); + const manifestPath = path.join(manifestRoot, digest + ".json"); + const temporaryPath = manifestPath + "." + process.pid + "." + crypto.randomBytes(4).toString("hex"); + fs.writeFileSync(temporaryPath, manifest, { encoding: "utf8", flag: "wx", mode: 0o600 }); + try { + try { + fs.linkSync(temporaryPath, manifestPath); + } catch (error) { + const existing = error && error.code === "EEXIST" ? fs.lstatSync(manifestPath) : null; + if ( + !existing || + existing.isSymbolicLink() || + !existing.isFile() || + fs.readFileSync(manifestPath, "utf8") !== manifest + ) { + throw error; + } + } + } finally { + fs.rmSync(temporaryPath, { force: true }); + } + process.stdout.write("sha256:" + digest + "\n"); +} +main().catch((error) => { + process.stderr.write(String(error && error.stack ? error.stack : error) + "\n"); + process.exitCode = 1; +});`; diff --git a/src/gateway/worker-environments/workspace-sync.ts b/src/gateway/worker-environments/workspace-sync.ts new file mode 100644 index 000000000000..6e9a126f5dea --- /dev/null +++ b/src/gateway/worker-environments/workspace-sync.ts @@ -0,0 +1,407 @@ +import { createHash } from "node:crypto"; +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { redactSensitiveText } from "../../logging/redact.js"; +import type { CommandOptions, SpawnResult } from "../../process/exec.js"; +import { + type PreparedWorkerSsh, + workerSshCommandOptions, + workerSshOptions, + workerSshRemoteCommand, +} from "./ssh.js"; +import type { + WorkerTunnelHandle, + WorkerWorkspaceCommand, + WorkerWorkspaceSyncRequest, + WorkerWorkspaceSyncResult, +} from "./tunnel-contract.js"; +import { runLocalCommandToFile, writeEligibleGitFiles } from "./workspace-sync-local.js"; +import { + REMOTE_GIT_WORKSPACE_SETUP_SCRIPT, + REMOTE_WORKSPACE_MANIFEST_JS, + REMOTE_WORKSPACE_SETUP_SCRIPT, +} from "./workspace-sync-scripts.js"; + +const REMOTE_SETUP_TIMEOUT_MS = 20_000; +const WORKSPACE_TIMEOUT_MS = 10 * 60_000; +// Relative to the $HOME/.openclaw-worker root owned by REMOTE_WORKSPACE_SETUP_SCRIPT; +// rsync targets must use the returned absolute directory, never this relative path. +const REMOTE_WORKSPACE_ROOT = "workspaces"; +const REMOTE_GIT_PACK_NAME = ".openclaw-base.pack"; +const GIT_COMMIT_PATTERN = /^[a-f0-9]{40}(?:[a-f0-9]{24})?$/u; +const MANIFEST_REF_PATTERN = /^sha256:[a-f0-9]{64}$/u; + +type WorkerWorkspaceRunner = { + run(argv: string[], options: CommandOptions): Promise; +}; + +type WorkerWorkspaceActionsOptions = { + environmentId: string; + ownerSignal: AbortSignal; + isConnected: () => boolean; + getPrepared: () => PreparedWorkerSsh | undefined; + runner: WorkerWorkspaceRunner; + tasks: Set>; +}; + +function success(result: SpawnResult): boolean { + return result.termination === "exit" && result.code === 0; +} + +function workspaceSyncError(result: SpawnResult): Error { + const detail = redactSensitiveText(result.stderr || result.stdout, { mode: "tools" }) + .replace(/\s+/gu, " ") + .trim(); + return new Error( + detail ? `Worker workspace sync failed: ${detail}` : "Worker workspace sync failed", + ); +} + +export function stableWorkerPathComponent(value: string, length: number): string { + return createHash("sha256").update(value).digest("hex").slice(0, length); +} + +function validateWorkspaceSyncRequest(request: WorkerWorkspaceSyncRequest): void { + if (!request.sessionId.trim()) { + throw new Error("Worker workspace session id must be non-empty"); + } + if (!path.isAbsolute(request.localPath)) { + throw new Error("Worker workspace local path must be absolute"); + } + if (!Number.isSafeInteger(request.generation) || request.generation < 0) { + throw new Error("Worker workspace generation must be a non-negative safe integer"); + } +} + +function parseRemoteWorkspaceDirectory(stdout: string): string { + const lines = stdout.split(/\r?\n/u).filter(Boolean); + const directory = lines.length === 1 ? lines[0] : undefined; + if ( + !directory || + !path.posix.isAbsolute(directory) || + path.posix.normalize(directory) !== directory || + directory === "/" + ) { + throw new Error("Worker workspace setup returned an invalid remote directory"); + } + return directory; +} + +function parseManifestRef(stdout: string): string { + const lines = stdout.split(/\r?\n/u).filter(Boolean); + const manifestRef = lines.length === 1 ? lines[0] : undefined; + if (!manifestRef || !MANIFEST_REF_PATTERN.test(manifestRef)) { + throw new Error("Worker workspace sync returned an invalid manifest reference"); + } + return manifestRef; +} + +/** Binds workspace commands and synchronization to one connected tunnel owner. */ +export function createWorkerWorkspaceActions( + options: WorkerWorkspaceActionsOptions, +): Pick { + const track = (task: Promise): Promise => { + options.tasks.add(task); + void task.then( + () => options.tasks.delete(task), + () => options.tasks.delete(task), + ); + return task; + }; + + const requirePrepared = (): PreparedWorkerSsh => { + const prepared = options.getPrepared(); + if (!options.isConnected() || !prepared) { + throw new Error("Worker tunnel owner is no longer connected"); + } + return prepared; + }; + + const runTask = (argv: string[], commandOptions: CommandOptions): Promise => + track(options.runner.run(argv, commandOptions)); + + const runWorkspaceCommand = async (command: WorkerWorkspaceCommand): Promise => { + const prepared = requirePrepared(); + return await runTask( + [ + "ssh", + ...workerSshOptions(prepared, { forwarding: "disabled" }), + "-a", + "-x", + "-T", + "-p", + String(prepared.port), + "--", + prepared.sshTarget, + workerSshRemoteCommand(command.argv), + ], + workerSshCommandOptions({ + input: command.input, + timeoutMs: command.timeoutMs ?? WORKSPACE_TIMEOUT_MS, + signal: command.signal + ? AbortSignal.any([options.ownerSignal, command.signal]) + : options.ownerSignal, + }), + ); + }; + + const syncWorkspaceImpl = async ( + request: WorkerWorkspaceSyncRequest, + ): Promise => { + validateWorkspaceSyncRequest(request); + const prepared = requirePrepared(); + const environmentKey = stableWorkerPathComponent(options.environmentId, 16); + const sessionKey = stableWorkerPathComponent(request.sessionId, 32); + const remoteRelative = [ + REMOTE_WORKSPACE_ROOT, + environmentKey, + sessionKey, + String(request.generation), + ].join("/"); + const setup = await runWorkspaceCommand({ + argv: ["sh", "-s", "--", remoteRelative], + input: REMOTE_WORKSPACE_SETUP_SCRIPT, + }); + if (!success(setup)) { + throw workspaceSyncError(setup); + } + const remoteWorkspaceDir = parseRemoteWorkspaceDirectory(setup.stdout.trim()); + + const gitRootResult = await runTask( + ["git", "-C", request.localPath, "rev-parse", "--show-toplevel"], + workerSshCommandOptions({ + timeoutMs: REMOTE_SETUP_TIMEOUT_MS, + signal: options.ownerSignal, + }), + ); + const mode = success(gitRootResult) ? "git" : "plain"; + let baseCommit = ""; + let gitRoot = request.localPath; + const temporaryDirectory = await fs.mkdtemp( + path.join(os.tmpdir(), "openclaw-worker-workspace-sync-"), + ); + const rsyncSsh = workerSshRemoteCommand([ + "ssh", + ...workerSshOptions(prepared, { forwarding: "disabled" }), + "-a", + "-x", + "-T", + "-p", + String(prepared.port), + ]); + try { + let fileListPath: string | undefined; + if (mode === "git") { + gitRoot = gitRootResult.stdout.trim(); + const [canonicalRequestPath, canonicalGitRoot] = await Promise.all([ + fs.realpath(request.localPath), + fs.realpath(gitRoot), + ]); + if (canonicalRequestPath !== canonicalGitRoot) { + throw new Error("Worker git workspace sync requires the managed worktree root"); + } + const gitBase = await runTask( + ["git", "-C", gitRoot, "rev-parse", "--verify", "HEAD"], + workerSshCommandOptions({ + timeoutMs: REMOTE_SETUP_TIMEOUT_MS, + signal: options.ownerSignal, + }), + ); + if (!success(gitBase)) { + throw new Error("Worker git workspace has no base commit"); + } + baseCommit = gitBase.stdout.trim(); + if (!GIT_COMMIT_PATTERN.test(baseCommit)) { + throw new Error("Worker workspace git base is not a commit id"); + } + + const eligiblePath = path.join(temporaryDirectory, "eligible"); + const ignoredPath = path.join(temporaryDirectory, "ignored"); + const selectedPath = path.join(temporaryDirectory, "selected"); + fileListPath = path.join(temporaryDirectory, "transfer-list"); + await runLocalCommandToFile({ + argv: [ + "git", + "-C", + gitRoot, + "ls-files", + "--full-name", + "--cached", + "--others", + "--exclude-standard", + "-z", + ], + outputPath: eligiblePath, + signal: options.ownerSignal, + timeoutMs: WORKSPACE_TIMEOUT_MS, + }); + const worktreeIncludePath = path.join(gitRoot, ".worktreeinclude"); + const worktreeInclude = await fs.lstat(worktreeIncludePath).catch(() => undefined); + if (worktreeInclude?.isFile()) { + await runLocalCommandToFile({ + argv: [ + "git", + "-C", + gitRoot, + "ls-files", + "--full-name", + "--others", + "--ignored", + "--exclude-standard", + "-z", + ], + outputPath: ignoredPath, + signal: options.ownerSignal, + timeoutMs: WORKSPACE_TIMEOUT_MS, + }); + await runLocalCommandToFile({ + argv: [ + "git", + "-C", + gitRoot, + "ls-files", + "--full-name", + "--others", + "--ignored", + `--exclude-from=${worktreeIncludePath}`, + "-z", + ], + outputPath: selectedPath, + signal: options.ownerSignal, + timeoutMs: WORKSPACE_TIMEOUT_MS, + }); + } else { + await Promise.all([ + fs.writeFile(ignoredPath, "", { mode: 0o600 }), + fs.writeFile(selectedPath, "", { mode: 0o600 }), + ]); + } + await writeEligibleGitFiles({ + gitRoot, + eligiblePath, + ignoredPath, + selectedPath, + outputPath: fileListPath, + }); + + const objectListPath = path.join(temporaryDirectory, "base-objects"); + const packPath = path.join(temporaryDirectory, "base.pack"); + await runLocalCommandToFile({ + argv: [ + "git", + "-C", + gitRoot, + "rev-list", + "--objects", + "--no-object-names", + `${baseCommit}^{tree}`, + ], + outputPath: objectListPath, + signal: options.ownerSignal, + timeoutMs: WORKSPACE_TIMEOUT_MS, + }); + await fs.appendFile(objectListPath, `${baseCommit}\n`); + await runLocalCommandToFile({ + argv: ["git", "-C", gitRoot, "pack-objects", "--stdout"], + inputPath: objectListPath, + outputPath: packPath, + signal: options.ownerSignal, + timeoutMs: WORKSPACE_TIMEOUT_MS, + }); + const packTransfer = await runTask( + [ + "rsync", + "--archive", + "--checksum", + "-e", + rsyncSsh, + "--", + packPath, + `${prepared.scpTarget}:${remoteWorkspaceDir}/${REMOTE_GIT_PACK_NAME}`, + ], + workerSshCommandOptions({ + timeoutMs: WORKSPACE_TIMEOUT_MS, + signal: options.ownerSignal, + }), + ); + if (!success(packTransfer)) { + throw workspaceSyncError(packTransfer); + } + const [authorName, authorEmail] = await Promise.all( + ["user.name", "user.email"].map(async (key) => { + const result = await runTask( + ["git", "-C", gitRoot, "config", "--get", key], + workerSshCommandOptions({ + timeoutMs: REMOTE_SETUP_TIMEOUT_MS, + signal: options.ownerSignal, + }), + ); + return success(result) ? result.stdout.trim() : ""; + }), + ); + const seeded = await runWorkspaceCommand({ + argv: [ + "sh", + "-s", + "--", + remoteWorkspaceDir, + path.posix.join(remoteWorkspaceDir, REMOTE_GIT_PACK_NAME), + baseCommit, + authorName ?? "", + authorEmail ?? "", + ], + input: REMOTE_GIT_WORKSPACE_SETUP_SCRIPT, + }); + if (!success(seeded)) { + throw workspaceSyncError(seeded); + } + } + + const localSource = gitRoot.endsWith(path.sep) ? gitRoot : `${gitRoot}${path.sep}`; + const transfer = await runTask( + [ + "rsync", + "--archive", + "--checksum", + "--exclude=.git", + ...(fileListPath ? ["--recursive", "--from0", `--files-from=${fileListPath}`] : []), + "-e", + rsyncSsh, + "--", + localSource, + `${prepared.scpTarget}:${remoteWorkspaceDir}/`, + ], + workerSshCommandOptions({ + timeoutMs: WORKSPACE_TIMEOUT_MS, + signal: options.ownerSignal, + }), + ); + if (!success(transfer)) { + throw workspaceSyncError(transfer); + } + + const manifest = await runWorkspaceCommand({ + argv: ["node", "-e", REMOTE_WORKSPACE_MANIFEST_JS, remoteWorkspaceDir, baseCommit], + }); + if (!success(manifest)) { + throw workspaceSyncError(manifest); + } + return { + mode, + remoteWorkspaceDir, + manifestRef: parseManifestRef(manifest.stdout.trim()), + }; + } finally { + await fs.rm(temporaryDirectory, { recursive: true, force: true }); + } + }; + + return { + runWorkspaceCommand, + syncWorkspace(request) { + // Keep the outer task registered across local-file phases so tunnel stop drains all owner work. + return track(syncWorkspaceImpl(request)); + }, + }; +} diff --git a/src/infra/agent-events.test.ts b/src/infra/agent-events.test.ts index 4d822c2a7beb..1a3624632e12 100644 --- a/src/infra/agent-events.test.ts +++ b/src/infra/agent-events.test.ts @@ -223,6 +223,45 @@ describe("agent-events sequencing", () => { expect(seen).toEqual(["worker"]); }); + test("explicitly adopts only an unowned same-generation context", () => { + const lifecycleGeneration = getAgentEventLifecycleGeneration(); + registerAgentRunContext("adopted-run", { + agentId: "main", + isControlUiVisible: false, + lifecycleGeneration, + sessionId: "session-adopted", + sessionKey: "agent:main:adopted", + }); + + const claimId = claimAgentRunContext( + "adopted-run", + { + agentId: "main", + isControlUiVisible: false, + lifecycleGeneration, + sessionId: "session-adopted", + sessionKey: "agent:main:adopted", + }, + { + adoptExistingUnowned: true, + exclusive: true, + ownsContext: true, + trackOwner: true, + }, + ); + expect(claimId).toBeDefined(); + expect( + claimAgentRunContext( + "adopted-run", + { lifecycleGeneration, sessionKey: "agent:main:adopted" }, + { adoptExistingUnowned: true, exclusive: true, trackOwner: true }, + ), + ).toBeUndefined(); + + releaseAgentRunContext("adopted-run", claimId); + expect(getAgentRunContext("adopted-run")).toBeUndefined(); + }); + test("full event reset clears tracked ownership", () => { const lifecycleGeneration = getAgentEventLifecycleGeneration(); claimAgentRunContext( diff --git a/src/infra/agent-events.ts b/src/infra/agent-events.ts index 15fd0132ef53..0de244c8ca4e 100644 --- a/src/infra/agent-events.ts +++ b/src/infra/agent-events.ts @@ -300,6 +300,8 @@ export function claimAgentRunContext( runId: string, context: AgentRunContext, options: { + /** Adopt a same-generation context only when no tracked execution owns it. */ + adoptExistingUnowned?: boolean; trackOwner?: boolean; ownsContext?: boolean; exclusive?: boolean; @@ -316,11 +318,16 @@ export function claimAgentRunContext( const existingOwners = ownersById.get(runId); const currentOwners = existingOwners?.lifecycleGeneration === lifecycleGeneration ? existingOwners : undefined; + const adoptsExistingUnowned = + options.exclusive === true && + options.adoptExistingUnowned === true && + existing?.lifecycleGeneration === lifecycleGeneration && + currentOwners === undefined; if ( currentOwners?.exclusiveClaimId || (options.exclusive && - (existing?.lifecycleGeneration === lifecycleGeneration || - (currentOwners?.claimIds.size ?? 0) > 0)) + ((existing?.lifecycleGeneration === lifecycleGeneration && !adoptsExistingUnowned) || + currentOwners !== undefined)) ) { return undefined; } diff --git a/src/state/openclaw-state-db.generated.d.ts b/src/state/openclaw-state-db.generated.d.ts index 1c40bd15861d..4c0265415d5d 100644 --- a/src/state/openclaw-state-db.generated.d.ts +++ b/src/state/openclaw-state-db.generated.d.ts @@ -1149,6 +1149,30 @@ export interface WorkerInferenceTurns { updated_at_ms: number; } +export interface WorkerSessionPlacements { + active_owner_epoch: number | null; + agent_id: string; + created_at_ms: number; + environment_id: string | null; + last_live_event_ack_cursor: number | null; + last_transcript_ack_cursor: number | null; + recovery_error: string | null; + remote_workspace_dir: string | null; + session_id: string; + session_key: string; + state: string; + state_changed_at_ms: number; + transition_generation: Generated; + turn_claim_generation: number | null; + turn_claim_id: string | null; + turn_claim_owner: string | null; + turn_claim_owner_epoch: number | null; + turn_claim_run_id: string | null; + updated_at_ms: number; + worker_bundle_hash: string | null; + workspace_base_manifest_ref: string | null; +} + export interface WorkerTranscriptCommitHeads { environment_id: string; next_seq: number; @@ -1272,6 +1296,7 @@ export interface DB { worker_environment_credentials: WorkerEnvironmentCredentials; worker_environments: WorkerEnvironments; worker_inference_turns: WorkerInferenceTurns; + worker_session_placements: WorkerSessionPlacements; worker_transcript_commit_heads: WorkerTranscriptCommitHeads; worker_transcript_commits: WorkerTranscriptCommits; workspace_setup_state: WorkspaceSetupState; diff --git a/src/state/openclaw-state-db.test.ts b/src/state/openclaw-state-db.test.ts index 846c232765ce..3417bd38e130 100644 --- a/src/state/openclaw-state-db.test.ts +++ b/src/state/openclaw-state-db.test.ts @@ -46,6 +46,71 @@ function createTempStateDir(): string { return makeTempDir(stateDbTempDirs, "openclaw-state-db-"); } +type PlacementConstraintProbe = { + sessionId: string; + state: string; + environmentId: string | null; + activeOwnerEpoch: number | null; + workerBundleHash: string | null; + recoveryError: string | null; + workspaceBaseManifestRef?: string; + remoteWorkspaceDir?: string; + lastTranscriptAckCursor?: number; + lastLiveEventAckCursor?: number; + turnClaimOwner?: "local" | "worker"; + turnClaimOwnerEpoch?: number; +}; + +function insertPlacementConstraintProbe( + database: DatabaseSync, + input: PlacementConstraintProbe, +): void { + const hasClaim = input.turnClaimOwner !== undefined; + database + .prepare( + `INSERT INTO worker_session_placements ( + session_id, + agent_id, + session_key, + state, + environment_id, + active_owner_epoch, + workspace_base_manifest_ref, + remote_workspace_dir, + worker_bundle_hash, + last_transcript_ack_cursor, + last_live_event_ack_cursor, + recovery_error, + turn_claim_owner, + turn_claim_id, + turn_claim_run_id, + turn_claim_generation, + turn_claim_owner_epoch, + created_at_ms, + updated_at_ms, + state_changed_at_ms + ) VALUES (?, 'main', ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 1, 1, 1)`, + ) + .run( + input.sessionId, + `agent:main:${input.sessionId}`, + input.state, + input.environmentId, + input.activeOwnerEpoch, + input.workspaceBaseManifestRef ?? null, + input.remoteWorkspaceDir ?? null, + input.workerBundleHash, + input.lastTranscriptAckCursor ?? null, + input.lastLiveEventAckCursor ?? null, + input.recoveryError, + input.turnClaimOwner ?? null, + hasClaim ? `${input.sessionId}-claim` : null, + hasClaim ? `${input.sessionId}-run` : null, + hasClaim ? 0 : null, + input.turnClaimOwnerEpoch ?? null, + ); +} + function statfsFixture(type: number): ReturnType { return { type, @@ -656,6 +721,375 @@ describe("openclaw state database", () => { expect(database.path).toBe(path.join(stateDir, "state", "openclaw.sqlite")); }); + it("rejects a placement turn claim tuple without an owner", () => { + const database = openOpenClawStateDatabase({ + env: { OPENCLAW_STATE_DIR: createTempStateDir() }, + }); + + expect(() => + database.db + .prepare( + `INSERT INTO worker_session_placements ( + session_id, + agent_id, + session_key, + state, + turn_claim_id, + turn_claim_run_id, + turn_claim_generation, + created_at_ms, + updated_at_ms, + state_changed_at_ms + ) VALUES (?, 'main', 'agent:main:placement-claim', 'local', ?, ?, 0, 1, 1, 1)`, + ) + .run("session-placement-claim", "claim-without-owner", "run-without-owner"), + ).toThrow(); + }); + + const validPlacementShapes = [ + { + name: "local placement", + sessionId: "session-local-valid", + state: "local", + environmentId: null, + activeOwnerEpoch: null, + workerBundleHash: null, + recoveryError: null, + }, + { + name: "requested placement", + sessionId: "session-requested-valid", + state: "requested", + environmentId: null, + activeOwnerEpoch: null, + workerBundleHash: null, + recoveryError: null, + }, + { + name: "provisioning placement before environment allocation", + sessionId: "session-provisioning-pending-valid", + state: "provisioning", + environmentId: null, + activeOwnerEpoch: null, + workerBundleHash: null, + recoveryError: null, + }, + { + name: "provisioning placement after environment allocation", + sessionId: "session-provisioning-allocated-valid", + state: "provisioning", + environmentId: "environment-provisioning", + activeOwnerEpoch: null, + workerBundleHash: null, + recoveryError: null, + }, + { + name: "syncing placement", + sessionId: "session-syncing-valid", + state: "syncing", + environmentId: "environment-syncing", + activeOwnerEpoch: null, + workerBundleHash: "bundle-syncing", + recoveryError: null, + }, + { + name: "starting placement", + sessionId: "session-starting-valid", + state: "starting", + environmentId: "environment-starting", + activeOwnerEpoch: null, + workspaceBaseManifestRef: "manifest-starting", + remoteWorkspaceDir: "/workspace/starting", + workerBundleHash: "bundle-starting", + recoveryError: null, + }, + { + name: "active placement", + sessionId: "session-active-valid", + state: "active", + environmentId: "environment-active", + activeOwnerEpoch: 7, + workspaceBaseManifestRef: "manifest-active", + remoteWorkspaceDir: "/workspace/active", + workerBundleHash: "bundle-active", + lastTranscriptAckCursor: 3, + lastLiveEventAckCursor: 4, + recoveryError: null, + }, + { + name: "draining placement", + sessionId: "session-draining-valid", + state: "draining", + environmentId: "environment-draining", + activeOwnerEpoch: 7, + workspaceBaseManifestRef: "manifest-draining", + remoteWorkspaceDir: "/workspace/draining", + workerBundleHash: "bundle-draining", + recoveryError: null, + }, + { + name: "reconciling placement", + sessionId: "session-reconciling-valid", + state: "reconciling", + environmentId: "environment-reconciling", + activeOwnerEpoch: 7, + workspaceBaseManifestRef: "manifest-reconciling", + remoteWorkspaceDir: "/workspace/reconciling", + workerBundleHash: "bundle-reconciling", + recoveryError: null, + }, + { + name: "reclaimed placement with full provenance", + sessionId: "session-reclaimed-valid", + state: "reclaimed", + environmentId: "environment-reclaimed", + activeOwnerEpoch: 7, + workspaceBaseManifestRef: "manifest-reclaimed", + remoteWorkspaceDir: "/workspace/reclaimed", + workerBundleHash: "bundle-reclaimed", + recoveryError: null, + }, + { + name: "failed placement with recovery detail", + sessionId: "session-failed-valid", + state: "failed", + environmentId: "environment-failed", + activeOwnerEpoch: null, + workerBundleHash: null, + recoveryError: "worker placement failed", + }, + ] satisfies Array; + + it.each(validPlacementShapes)("allows a valid $name", (input) => { + const database = openOpenClawStateDatabase({ + env: { OPENCLAW_STATE_DIR: createTempStateDir() }, + }); + + expect(() => insertPlacementConstraintProbe(database.db, input)).not.toThrow(); + }); + + const invalidPlacementShapes = [ + { + name: "local environment", + sessionId: "session-local-environment", + state: "local", + environmentId: "environment-local", + activeOwnerEpoch: null, + workerBundleHash: null, + recoveryError: null, + }, + { + name: "syncing without environment", + sessionId: "session-syncing-environment", + state: "syncing", + environmentId: null, + activeOwnerEpoch: null, + workerBundleHash: "bundle-hash", + recoveryError: null, + }, + { + name: "syncing workspace metadata", + sessionId: "session-syncing-workspace", + state: "syncing", + environmentId: "environment-syncing", + activeOwnerEpoch: null, + workspaceBaseManifestRef: "manifest-syncing", + remoteWorkspaceDir: "/workspace/syncing", + workerBundleHash: "bundle-hash", + recoveryError: null, + }, + { + name: "active without owner epoch", + sessionId: "session-active-epoch", + state: "active", + environmentId: "environment-active", + activeOwnerEpoch: null, + workerBundleHash: "bundle-hash", + recoveryError: null, + workspaceBaseManifestRef: "manifest-active", + remoteWorkspaceDir: "/workspace/active", + }, + { + name: "active without worker bundle", + sessionId: "session-active-bundle", + state: "active", + environmentId: "environment-active", + activeOwnerEpoch: 7, + workerBundleHash: null, + recoveryError: null, + workspaceBaseManifestRef: "manifest-active", + remoteWorkspaceDir: "/workspace/active", + }, + { + name: "starting without manifest", + sessionId: "session-starting-manifest", + state: "starting", + environmentId: "environment-starting", + activeOwnerEpoch: null, + workerBundleHash: "bundle-hash", + recoveryError: null, + remoteWorkspaceDir: "/workspace/starting", + }, + { + name: "starting owner epoch", + sessionId: "session-starting-epoch", + state: "starting", + environmentId: "environment-starting", + activeOwnerEpoch: 7, + workspaceBaseManifestRef: "manifest-starting", + remoteWorkspaceDir: "/workspace/starting", + workerBundleHash: "bundle-hash", + recoveryError: null, + }, + { + name: "requested worker metadata", + sessionId: "session-requested-metadata", + state: "requested", + environmentId: null, + activeOwnerEpoch: null, + workerBundleHash: "bundle-hash", + recoveryError: null, + }, + { + name: "provisioning worker bundle", + sessionId: "session-provisioning-bundle", + state: "provisioning", + environmentId: "environment-provisioning", + activeOwnerEpoch: null, + workerBundleHash: "bundle-hash", + recoveryError: null, + }, + { + name: "active recovery error", + sessionId: "session-active-recovery", + state: "active", + environmentId: "environment-active", + activeOwnerEpoch: 7, + workspaceBaseManifestRef: "manifest-active", + remoteWorkspaceDir: "/workspace/active", + workerBundleHash: "bundle-hash", + recoveryError: "unexpected active recovery detail", + }, + { + name: "reclaimed placement without full provenance", + sessionId: "session-reclaimed-provenance", + state: "reclaimed", + environmentId: "environment-reclaimed", + activeOwnerEpoch: null, + workerBundleHash: "bundle-hash", + recoveryError: null, + }, + { + name: "reclaimed recovery error", + sessionId: "session-reclaimed-recovery", + state: "reclaimed", + environmentId: "environment-reclaimed", + activeOwnerEpoch: 7, + workspaceBaseManifestRef: "manifest-reclaimed", + remoteWorkspaceDir: "/workspace/reclaimed", + workerBundleHash: "bundle-hash", + recoveryError: "unexpected reclaimed recovery detail", + }, + { + name: "failed without recovery error", + sessionId: "session-failed-recovery", + state: "failed", + environmentId: null, + activeOwnerEpoch: null, + workerBundleHash: null, + recoveryError: null, + }, + ] satisfies Array; + + it.each(invalidPlacementShapes)("rejects a placement with $name", (input) => { + const database = openOpenClawStateDatabase({ + env: { OPENCLAW_STATE_DIR: createTempStateDir() }, + }); + + expect(() => insertPlacementConstraintProbe(database.db, input)).toThrow(); + }); + + const invalidPlacementClaimOwners = [ + { + name: "local claim on active placement", + state: "active", + activeOwnerEpoch: 7, + turnClaimOwner: "local", + turnClaimOwnerEpoch: undefined, + }, + { + name: "worker claim on reconciling placement", + state: "reconciling", + activeOwnerEpoch: 7, + turnClaimOwner: "worker", + turnClaimOwnerEpoch: 7, + }, + { + name: "stale worker owner epoch", + state: "active", + activeOwnerEpoch: 7, + turnClaimOwner: "worker", + turnClaimOwnerEpoch: 8, + }, + { + name: "worker claim on reclaimed placement", + state: "reclaimed", + activeOwnerEpoch: 7, + turnClaimOwner: "worker", + turnClaimOwnerEpoch: 7, + }, + ] satisfies Array<{ + name: string; + state: string; + activeOwnerEpoch: number; + turnClaimOwner: "local" | "worker"; + turnClaimOwnerEpoch: number | undefined; + }>; + + it.each(invalidPlacementClaimOwners)("rejects a placement with $name", (input) => { + const database = openOpenClawStateDatabase({ + env: { OPENCLAW_STATE_DIR: createTempStateDir() }, + }); + + expect(() => + insertPlacementConstraintProbe(database.db, { + sessionId: `session-${input.state}-${input.turnClaimOwner}`, + state: input.state, + environmentId: `environment-${input.state}`, + activeOwnerEpoch: input.activeOwnerEpoch, + workspaceBaseManifestRef: `manifest-${input.state}`, + remoteWorkspaceDir: `/workspace/${input.state}`, + workerBundleHash: "bundle-hash", + recoveryError: null, + turnClaimOwner: input.turnClaimOwner, + ...(input.turnClaimOwnerEpoch === undefined + ? {} + : { turnClaimOwnerEpoch: input.turnClaimOwnerEpoch }), + }), + ).toThrow(); + }); + + it("allows an exact worker claim while placement drains", () => { + const database = openOpenClawStateDatabase({ + env: { OPENCLAW_STATE_DIR: createTempStateDir() }, + }); + + expect(() => + insertPlacementConstraintProbe(database.db, { + sessionId: "session-draining-worker", + state: "draining", + environmentId: "environment-draining", + activeOwnerEpoch: 7, + workspaceBaseManifestRef: "manifest-draining", + remoteWorkspaceDir: "/workspace/draining", + workerBundleHash: "bundle-hash", + recoveryError: null, + turnClaimOwner: "worker", + turnClaimOwnerEpoch: 7, + }), + ).not.toThrow(); + }); + it("repairs a same-name shared-state uniqueness index", () => { const stateDir = createTempStateDir(); const env = { OPENCLAW_STATE_DIR: stateDir }; diff --git a/src/state/openclaw-state-schema.generated.ts b/src/state/openclaw-state-schema.generated.ts index e9cbb578c9b8..8d70918e1f69 100644 --- a/src/state/openclaw-state-schema.generated.ts +++ b/src/state/openclaw-state-schema.generated.ts @@ -1586,6 +1586,122 @@ CREATE UNIQUE INDEX IF NOT EXISTS idx_worker_environments_provider_lease ON worker_environments(provider_id, lease_id) WHERE lease_id IS NOT NULL; +-- Session placement lives in the shared state database so local admission, +-- worker admission, and environment attachment use one durable authority. +CREATE TABLE IF NOT EXISTS worker_session_placements ( + session_id TEXT NOT NULL PRIMARY KEY, + agent_id TEXT NOT NULL, + session_key TEXT NOT NULL, + state TEXT NOT NULL CHECK ( + state IN ( + 'local', + 'requested', + 'provisioning', + 'syncing', + 'starting', + 'active', + 'draining', + 'reconciling', + 'reclaimed', + 'failed' + ) + ), + environment_id TEXT, + transition_generation INTEGER NOT NULL DEFAULT 0 CHECK (transition_generation >= 0), + active_owner_epoch INTEGER CHECK (active_owner_epoch IS NULL OR active_owner_epoch >= 1), + workspace_base_manifest_ref TEXT, + remote_workspace_dir TEXT, + worker_bundle_hash TEXT, + last_transcript_ack_cursor INTEGER CHECK ( + last_transcript_ack_cursor IS NULL OR last_transcript_ack_cursor >= 0 + ), + last_live_event_ack_cursor INTEGER CHECK ( + last_live_event_ack_cursor IS NULL OR last_live_event_ack_cursor >= 0 + ), + recovery_error TEXT, + turn_claim_owner TEXT CHECK (turn_claim_owner IN ('local', 'worker')), + turn_claim_id TEXT, + turn_claim_run_id TEXT, + turn_claim_generation INTEGER CHECK ( + turn_claim_generation IS NULL OR turn_claim_generation >= 0 + ), + turn_claim_owner_epoch INTEGER CHECK ( + turn_claim_owner_epoch IS NULL OR turn_claim_owner_epoch >= 1 + ), + created_at_ms INTEGER NOT NULL, + updated_at_ms INTEGER NOT NULL, + state_changed_at_ms INTEGER NOT NULL, + CHECK ( + (state IN ('local', 'requested') + AND environment_id IS NULL AND active_owner_epoch IS NULL + AND workspace_base_manifest_ref IS NULL AND remote_workspace_dir IS NULL + AND worker_bundle_hash IS NULL + AND last_transcript_ack_cursor IS NULL AND last_live_event_ack_cursor IS NULL + AND recovery_error IS NULL) + OR + (state IS 'provisioning' + AND active_owner_epoch IS NULL + AND workspace_base_manifest_ref IS NULL AND remote_workspace_dir IS NULL + AND worker_bundle_hash IS NULL + AND last_transcript_ack_cursor IS NULL AND last_live_event_ack_cursor IS NULL + AND recovery_error IS NULL) + OR + (state IS 'syncing' + AND environment_id IS NOT NULL AND active_owner_epoch IS NULL + AND workspace_base_manifest_ref IS NULL AND remote_workspace_dir IS NULL + AND worker_bundle_hash IS NOT NULL + AND last_transcript_ack_cursor IS NULL AND last_live_event_ack_cursor IS NULL + AND recovery_error IS NULL) + OR + (state IS 'starting' + AND environment_id IS NOT NULL AND active_owner_epoch IS NULL + AND workspace_base_manifest_ref IS NOT NULL AND remote_workspace_dir IS NOT NULL + AND worker_bundle_hash IS NOT NULL + AND last_transcript_ack_cursor IS NULL AND last_live_event_ack_cursor IS NULL + AND recovery_error IS NULL) + OR + (state IN ('active', 'draining', 'reconciling') + AND environment_id IS NOT NULL AND active_owner_epoch IS NOT NULL + AND workspace_base_manifest_ref IS NOT NULL AND remote_workspace_dir IS NOT NULL + AND worker_bundle_hash IS NOT NULL AND recovery_error IS NULL) + OR + (state IS 'reclaimed' + AND environment_id IS NOT NULL AND active_owner_epoch IS NOT NULL + AND workspace_base_manifest_ref IS NOT NULL AND remote_workspace_dir IS NOT NULL + AND worker_bundle_hash IS NOT NULL AND recovery_error IS NULL + AND turn_claim_owner IS NULL AND turn_claim_id IS NULL AND turn_claim_run_id IS NULL + AND turn_claim_generation IS NULL AND turn_claim_owner_epoch IS NULL) + OR + (state IS 'failed' AND recovery_error IS NOT NULL) + ), + CHECK ( + (turn_claim_owner IS NULL AND turn_claim_id IS NULL AND turn_claim_run_id IS NULL + AND turn_claim_generation IS NULL AND turn_claim_owner_epoch IS NULL) + OR + (turn_claim_owner IS 'local' AND turn_claim_id IS NOT NULL + AND turn_claim_run_id IS NOT NULL AND turn_claim_generation IS NOT NULL + AND turn_claim_owner_epoch IS NULL) + OR + (turn_claim_owner IS 'worker' AND turn_claim_id IS NOT NULL + AND turn_claim_run_id IS NOT NULL AND turn_claim_generation IS NOT NULL + AND turn_claim_owner_epoch IS NOT NULL) + ), + CHECK ( + turn_claim_owner IS NULL + OR + (turn_claim_owner IS 'local' AND state IN ('local', 'requested', 'failed')) + OR + (turn_claim_owner IS 'worker' AND state IN ('active', 'draining') + AND turn_claim_owner_epoch IS active_owner_epoch) + ) +); + +CREATE INDEX IF NOT EXISTS idx_worker_session_placements_session_key + ON worker_session_placements(agent_id, session_key); + +CREATE INDEX IF NOT EXISTS idx_worker_session_placements_reconcile + ON worker_session_placements(updated_at_ms, session_id); + -- One active, opaque admission credential per worker environment. Plaintext -- may be retried until delivery acknowledgement but never enters durable state. CREATE TABLE IF NOT EXISTS worker_environment_credentials ( diff --git a/src/state/openclaw-state-schema.sql b/src/state/openclaw-state-schema.sql index e46d3a5d436d..ca0fa73854cd 100644 --- a/src/state/openclaw-state-schema.sql +++ b/src/state/openclaw-state-schema.sql @@ -1581,6 +1581,122 @@ CREATE UNIQUE INDEX IF NOT EXISTS idx_worker_environments_provider_lease ON worker_environments(provider_id, lease_id) WHERE lease_id IS NOT NULL; +-- Session placement lives in the shared state database so local admission, +-- worker admission, and environment attachment use one durable authority. +CREATE TABLE IF NOT EXISTS worker_session_placements ( + session_id TEXT NOT NULL PRIMARY KEY, + agent_id TEXT NOT NULL, + session_key TEXT NOT NULL, + state TEXT NOT NULL CHECK ( + state IN ( + 'local', + 'requested', + 'provisioning', + 'syncing', + 'starting', + 'active', + 'draining', + 'reconciling', + 'reclaimed', + 'failed' + ) + ), + environment_id TEXT, + transition_generation INTEGER NOT NULL DEFAULT 0 CHECK (transition_generation >= 0), + active_owner_epoch INTEGER CHECK (active_owner_epoch IS NULL OR active_owner_epoch >= 1), + workspace_base_manifest_ref TEXT, + remote_workspace_dir TEXT, + worker_bundle_hash TEXT, + last_transcript_ack_cursor INTEGER CHECK ( + last_transcript_ack_cursor IS NULL OR last_transcript_ack_cursor >= 0 + ), + last_live_event_ack_cursor INTEGER CHECK ( + last_live_event_ack_cursor IS NULL OR last_live_event_ack_cursor >= 0 + ), + recovery_error TEXT, + turn_claim_owner TEXT CHECK (turn_claim_owner IN ('local', 'worker')), + turn_claim_id TEXT, + turn_claim_run_id TEXT, + turn_claim_generation INTEGER CHECK ( + turn_claim_generation IS NULL OR turn_claim_generation >= 0 + ), + turn_claim_owner_epoch INTEGER CHECK ( + turn_claim_owner_epoch IS NULL OR turn_claim_owner_epoch >= 1 + ), + created_at_ms INTEGER NOT NULL, + updated_at_ms INTEGER NOT NULL, + state_changed_at_ms INTEGER NOT NULL, + CHECK ( + (state IN ('local', 'requested') + AND environment_id IS NULL AND active_owner_epoch IS NULL + AND workspace_base_manifest_ref IS NULL AND remote_workspace_dir IS NULL + AND worker_bundle_hash IS NULL + AND last_transcript_ack_cursor IS NULL AND last_live_event_ack_cursor IS NULL + AND recovery_error IS NULL) + OR + (state IS 'provisioning' + AND active_owner_epoch IS NULL + AND workspace_base_manifest_ref IS NULL AND remote_workspace_dir IS NULL + AND worker_bundle_hash IS NULL + AND last_transcript_ack_cursor IS NULL AND last_live_event_ack_cursor IS NULL + AND recovery_error IS NULL) + OR + (state IS 'syncing' + AND environment_id IS NOT NULL AND active_owner_epoch IS NULL + AND workspace_base_manifest_ref IS NULL AND remote_workspace_dir IS NULL + AND worker_bundle_hash IS NOT NULL + AND last_transcript_ack_cursor IS NULL AND last_live_event_ack_cursor IS NULL + AND recovery_error IS NULL) + OR + (state IS 'starting' + AND environment_id IS NOT NULL AND active_owner_epoch IS NULL + AND workspace_base_manifest_ref IS NOT NULL AND remote_workspace_dir IS NOT NULL + AND worker_bundle_hash IS NOT NULL + AND last_transcript_ack_cursor IS NULL AND last_live_event_ack_cursor IS NULL + AND recovery_error IS NULL) + OR + (state IN ('active', 'draining', 'reconciling') + AND environment_id IS NOT NULL AND active_owner_epoch IS NOT NULL + AND workspace_base_manifest_ref IS NOT NULL AND remote_workspace_dir IS NOT NULL + AND worker_bundle_hash IS NOT NULL AND recovery_error IS NULL) + OR + (state IS 'reclaimed' + AND environment_id IS NOT NULL AND active_owner_epoch IS NOT NULL + AND workspace_base_manifest_ref IS NOT NULL AND remote_workspace_dir IS NOT NULL + AND worker_bundle_hash IS NOT NULL AND recovery_error IS NULL + AND turn_claim_owner IS NULL AND turn_claim_id IS NULL AND turn_claim_run_id IS NULL + AND turn_claim_generation IS NULL AND turn_claim_owner_epoch IS NULL) + OR + (state IS 'failed' AND recovery_error IS NOT NULL) + ), + CHECK ( + (turn_claim_owner IS NULL AND turn_claim_id IS NULL AND turn_claim_run_id IS NULL + AND turn_claim_generation IS NULL AND turn_claim_owner_epoch IS NULL) + OR + (turn_claim_owner IS 'local' AND turn_claim_id IS NOT NULL + AND turn_claim_run_id IS NOT NULL AND turn_claim_generation IS NOT NULL + AND turn_claim_owner_epoch IS NULL) + OR + (turn_claim_owner IS 'worker' AND turn_claim_id IS NOT NULL + AND turn_claim_run_id IS NOT NULL AND turn_claim_generation IS NOT NULL + AND turn_claim_owner_epoch IS NOT NULL) + ), + CHECK ( + turn_claim_owner IS NULL + OR + (turn_claim_owner IS 'local' AND state IN ('local', 'requested', 'failed')) + OR + (turn_claim_owner IS 'worker' AND state IN ('active', 'draining') + AND turn_claim_owner_epoch IS active_owner_epoch) + ) +); + +CREATE INDEX IF NOT EXISTS idx_worker_session_placements_session_key + ON worker_session_placements(agent_id, session_key); + +CREATE INDEX IF NOT EXISTS idx_worker_session_placements_reconcile + ON worker_session_placements(updated_at_ms, session_id); + -- One active, opaque admission credential per worker environment. Plaintext -- may be retried until delivery acknowledgement but never enters durable state. CREATE TABLE IF NOT EXISTS worker_environment_credentials ( diff --git a/src/worker/embedded-agent-live.runtime.ts b/src/worker/embedded-agent-live.runtime.ts index afff3012c1c9..cd79abafdab9 100644 --- a/src/worker/embedded-agent-live.runtime.ts +++ b/src/worker/embedded-agent-live.runtime.ts @@ -166,6 +166,7 @@ type WorkerLiveRuntime = { handleSessionEvent: (event: AgentSessionEvent) => void; enqueueRunFailure: (failure: { aborted: boolean; error: Error }) => void; flush: () => Promise; + emitTerminal: () => Promise; }; export function createWorkerLiveRuntime(client: WorkerLiveClient): WorkerLiveRuntime { @@ -219,6 +220,9 @@ export function createWorkerLiveRuntime(client: WorkerLiveClient): WorkerLiveRun }; const startedAt = Date.now(); let lifecycleFinished = false; + // Terminal lifecycle events are deferred past the final transcript flush so the + // gateway never sees an end/error before the authoritative transcript commit. + let terminalLiveEvent: WorkerLiveEvent | undefined; let streamedText = ""; let streamedThinking = ""; const handleSessionEvent = (event: AgentSessionEvent) => { @@ -310,17 +314,18 @@ export function createWorkerLiveRuntime(client: WorkerLiveClient): WorkerLiveRun .toReversed() .find((message): message is AssistantMessage => message.role === "assistant"); if (lastAssistant?.stopReason === "error") { - enqueueLive({ + terminalLiveEvent = { kind: "lifecycle", payload: { phase: "error", startedAt, endedAt: Date.now(), error: lastAssistant.errorMessage ?? "Worker inference failed.", + fallbackExhaustedFailure: true, }, - }); + }; } else if (lastAssistant?.stopReason === "aborted") { - enqueueLive({ + terminalLiveEvent = { kind: "lifecycle", payload: { phase: "end", @@ -329,12 +334,12 @@ export function createWorkerLiveRuntime(client: WorkerLiveClient): WorkerLiveRun stopReason: "aborted", aborted: true, }, - }); + }; } else { - enqueueLive({ + terminalLiveEvent = { kind: "lifecycle", payload: { phase: "end", startedAt, endedAt: Date.now() }, - }); + }; } } }; @@ -343,7 +348,7 @@ export function createWorkerLiveRuntime(client: WorkerLiveClient): WorkerLiveRun return; } if (failure.aborted) { - enqueueLive({ + terminalLiveEvent = { kind: "lifecycle", payload: { phase: "end", @@ -352,18 +357,27 @@ export function createWorkerLiveRuntime(client: WorkerLiveClient): WorkerLiveRun stopReason: "aborted", aborted: true, }, - }); + }; } else { - enqueueLive({ + terminalLiveEvent = { kind: "lifecycle", payload: { phase: "error", startedAt, endedAt: Date.now(), error: failure.error.message, + fallbackExhaustedFailure: true, }, - }); + }; } }; - return { handleSessionEvent, enqueueRunFailure, flush }; + // Emits directly (not via the degradable preview queue): the terminal event drives + // gateway turn settlement and must survive a degraded live stream. + const emitTerminal = async () => { + if (!terminalLiveEvent) { + return; + } + await client.emit(boundLiveEvent(terminalLiveEvent)); + }; + return { handleSessionEvent, enqueueRunFailure, flush, emitTerminal }; } diff --git a/src/worker/embedded-agent-transcript.runtime.ts b/src/worker/embedded-agent-transcript.runtime.ts index 7b027ea1c24c..87edbd83c900 100644 --- a/src/worker/embedded-agent-transcript.runtime.ts +++ b/src/worker/embedded-agent-transcript.runtime.ts @@ -4,126 +4,14 @@ import type { WorkerInferenceContext } from "../../packages/gateway-protocol/src import { WORKER_INFERENCE_MAX_CONTEXT_MESSAGES } from "../../packages/gateway-protocol/src/schema/worker-inference.js"; import type { AgentMessage } from "../agents/runtime/index.js"; import type { AgentSessionWriteLockRunner } from "../agents/sessions/agent-session.js"; -import type { AssistantMessage, Context, Message } from "../llm/types.js"; -import { isWorkerTranscriptMessageFrameSafe } from "./transcript-message.js"; - -function cloneTextContent(part: { type: "text"; text: string; textSignature?: string }) { - return { - type: "text" as const, - text: part.text, - ...(part.textSignature ? { textSignature: part.textSignature } : {}), - }; -} - -function cloneImageContent(part: { type: "image"; data: string; mimeType: string }) { - return { type: "image" as const, data: part.data, mimeType: part.mimeType }; -} - -function cloneUsage(message: AssistantMessage): WorkerTranscriptMessage & { role: "assistant" } { - return { - role: "assistant", - content: message.content.map((part) => { - if (part.type === "text") { - return cloneTextContent(part); - } - if (part.type === "thinking") { - return { - type: "thinking" as const, - thinking: part.thinking, - ...(part.thinkingSignature ? { thinkingSignature: part.thinkingSignature } : {}), - ...(part.redacted === undefined ? {} : { redacted: part.redacted }), - }; - } - return { - type: "toolCall" as const, - id: part.id, - name: part.name, - arguments: structuredClone(part.arguments), - ...(part.thoughtSignature ? { thoughtSignature: part.thoughtSignature } : {}), - ...(part.executionMode ? { executionMode: part.executionMode } : {}), - }; - }), - api: message.api, - provider: message.provider, - model: message.model, - ...(message.responseModel ? { responseModel: message.responseModel } : {}), - ...(message.responseId ? { responseId: message.responseId } : {}), - ...(message.diagnostics - ? { - diagnostics: message.diagnostics.map((diagnostic) => ({ - type: diagnostic.type, - timestamp: diagnostic.timestamp, - ...(diagnostic.error - ? { - error: { - ...(diagnostic.error.name ? { name: diagnostic.error.name } : {}), - message: diagnostic.error.message, - ...(diagnostic.error.stack ? { stack: diagnostic.error.stack } : {}), - ...(diagnostic.error.code === undefined ? {} : { code: diagnostic.error.code }), - }, - } - : {}), - ...(diagnostic.details ? { details: structuredClone(diagnostic.details) } : {}), - })), - } - : {}), - usage: { - input: message.usage.input, - output: message.usage.output, - cacheRead: message.usage.cacheRead, - cacheWrite: message.usage.cacheWrite, - ...(message.usage.contextUsage - ? { contextUsage: structuredClone(message.usage.contextUsage) } - : {}), - totalTokens: message.usage.totalTokens, - cost: { - input: message.usage.cost.input, - output: message.usage.cost.output, - cacheRead: message.usage.cost.cacheRead, - cacheWrite: message.usage.cost.cacheWrite, - total: message.usage.cost.total, - ...(message.usage.cost.totalOrigin ? { totalOrigin: message.usage.cost.totalOrigin } : {}), - }, - }, - stopReason: message.stopReason, - ...(message.errorMessage ? { errorMessage: message.errorMessage } : {}), - ...(message.errorCode ? { errorCode: message.errorCode } : {}), - ...(message.errorType ? { errorType: message.errorType } : {}), - ...(message.errorBody ? { errorBody: message.errorBody } : {}), - timestamp: message.timestamp, - }; -} - -export function toWorkerTranscriptMessage( - message: AgentMessage, -): WorkerTranscriptMessage | undefined { - if (message.role === "user") { - const content = - typeof message.content === "string" - ? [{ type: "text" as const, text: message.content }] - : message.content.map((part) => - part.type === "text" ? cloneTextContent(part) : cloneImageContent(part), - ); - return { role: "user", content, timestamp: message.timestamp }; - } - if (message.role === "assistant") { - return cloneUsage(message); - } - if (message.role === "toolResult") { - return { - role: "toolResult", - toolCallId: message.toolCallId, - toolName: message.toolName, - content: message.content.map((part) => - part.type === "text" ? cloneTextContent(part) : cloneImageContent(part), - ), - ...(message.details === undefined ? {} : { details: structuredClone(message.details) }), - isError: message.isError, - timestamp: message.timestamp, - }; - } - return undefined; -} +import type { Context, Message } from "../llm/types.js"; +import { + cloneImageContent, + cloneTextContent, + cloneUsage, + isWorkerTranscriptMessageFrameSafe, + toWorkerTranscriptMessage, +} from "./transcript-message.js"; export function toAgentMessage(message: WorkerTranscriptMessage): Message { if (message.role === "user") { diff --git a/src/worker/embedded-agent.runtime.ts b/src/worker/embedded-agent.runtime.ts index 3b83e2e71991..9a750433fdb1 100644 --- a/src/worker/embedded-agent.runtime.ts +++ b/src/worker/embedded-agent.runtime.ts @@ -26,8 +26,8 @@ import { createWorkerTranscriptRuntime, toAgentMessage, toWorkerInferenceContext, - toWorkerTranscriptMessage, } from "./embedded-agent-transcript.runtime.js"; +import { toWorkerTranscriptMessage } from "./transcript-message.js"; const LOCAL_WORKER_TOOL_NAMES = [ "read", @@ -75,6 +75,7 @@ type RunWorkerEmbeddedTurnParams = { transcript: WorkerEmbeddedTranscriptClient; live: WorkerEmbeddedLiveClient; initialMessages?: WorkerTranscriptMessage[]; + suppressPromptTranscript?: boolean; systemPrompt?: string; inferenceOptions?: WorkerInferenceOptions; signal?: AbortSignal; @@ -110,7 +111,7 @@ export async function runWorkerEmbeddedTurn( noPromptTemplates: true, noThemes: true, noContextFiles: true, - ...(params.systemPrompt === undefined ? {} : { systemPrompt: params.systemPrompt }), + ...(params.systemPrompt === undefined ? {} : { appendSystemPrompt: [params.systemPrompt] }), agentsFilesOverride: () => ({ agentsFiles: contextFiles }), }); await resourceLoader.reload(); @@ -122,6 +123,7 @@ export async function runWorkerEmbeddedTurn( const transcriptRuntime = createWorkerTranscriptRuntime(params.transcript); const sessionManager = guardSessionManager(baseSessionManager, { + suppressNextUserMessagePersistence: params.suppressPromptTranscript, onMessagePersisted: transcriptRuntime.onMessagePersisted, }); @@ -221,13 +223,14 @@ export async function runWorkerEmbeddedTurn( let finalTranscriptFailure: Error | undefined; try { - if (!params.signal?.aborted) { - try { - await transcriptRuntime.withSessionWriteLock(() => undefined); - } catch (error) { - finalTranscriptFailure = toError(error, "Worker transcript flush failed."); - } - await liveRuntime.flush(); + try { + await transcriptRuntime.withSessionWriteLock(() => undefined); + } catch (error) { + finalTranscriptFailure = toError(error, "Worker transcript flush failed."); + } + await liveRuntime.flush(); + if (finalTranscriptFailure === undefined) { + await liveRuntime.emitTerminal(); } } finally { params.signal?.removeEventListener("abort", abortTurn); diff --git a/src/worker/launch-descriptor.test.ts b/src/worker/launch-descriptor.test.ts index 09dddd04ab48..35c0f4a09927 100644 --- a/src/worker/launch-descriptor.test.ts +++ b/src/worker/launch-descriptor.test.ts @@ -28,6 +28,7 @@ function launchDescriptor(): WorkerLaunchDescriptor { runId: "run-1", turnId: "turn-1", prompt: "Inspect the workspace.", + suppressPromptTranscript: false, workspaceDir: "/tmp/openclaw-worker/workspace", modelRef: { provider: "provider-1", model: "model-1" }, inferenceOptions: { reasoning: "medium", maxTokens: 512 }, @@ -52,7 +53,7 @@ describe("worker launch descriptor", () => { expect(buildWorkerConnectParams(descriptor)).toMatchObject({ role: "worker", client: { id: "openclaw-worker", mode: "worker", version: "2026.7.12" }, - admission: descriptor.admission, + admission: { ...descriptor.admission, runId: descriptor.assignment.runId }, }); }); diff --git a/src/worker/launch-descriptor.ts b/src/worker/launch-descriptor.ts index 5cf9bae6c0fd..ee3eb50728fb 100644 --- a/src/worker/launch-descriptor.ts +++ b/src/worker/launch-descriptor.ts @@ -31,6 +31,7 @@ type WorkerLaunchAssignment = { runId: string; turnId: string; prompt: string; + suppressPromptTranscript: boolean; workspaceDir: string; modelRef: WorkerInferenceModelRef; inferenceOptions: WorkerInferenceOptions; @@ -46,7 +47,7 @@ type WorkerLaunchAssignment = { }; }; -type WorkerLaunchAdmission = WorkerConnectParams["admission"] & { +type WorkerLaunchAdmission = Omit & { sessionId: string; }; @@ -94,6 +95,7 @@ function parseAssignment(value: unknown): WorkerLaunchAssignment | undefined { "runId", "turnId", "prompt", + "suppressPromptTranscript", "workspaceDir", "modelRef", "inferenceOptions", @@ -110,6 +112,7 @@ function parseAssignment(value: unknown): WorkerLaunchAssignment | undefined { !isIdentifier(value.runId) || !isIdentifier(value.turnId) || typeof value.prompt !== "string" || + typeof value.suppressPromptTranscript !== "boolean" || !isIdentifier(value.workspaceDir) || !path.isAbsolute(value.workspaceDir) || (value.systemPrompt !== undefined && typeof value.systemPrompt !== "string") || @@ -146,7 +149,7 @@ function parseAssignment(value: unknown): WorkerLaunchAssignment | undefined { } export function buildWorkerConnectParams( - descriptor: Pick, + descriptor: Pick, ): WorkerConnectParams { return { minProtocol: PROTOCOL_VERSION, @@ -158,7 +161,10 @@ export function buildWorkerConnectParams( mode: GATEWAY_CLIENT_MODES.WORKER, }, role: "worker", - admission: descriptor.admission, + admission: { + ...descriptor.admission, + runId: descriptor.assignment.runId, + }, }; } diff --git a/src/worker/transcript-message.ts b/src/worker/transcript-message.ts index ffd4c8fdf083..6f4e7167ad19 100644 --- a/src/worker/transcript-message.ts +++ b/src/worker/transcript-message.ts @@ -6,9 +6,131 @@ import { WORKER_PROTOCOL_MAX_IDENTIFIER_LENGTH, WORKER_PROTOCOL_MAX_PAYLOAD_BYTES, } from "../../packages/gateway-protocol/src/schema/worker-admission.js"; +import type { AgentMessage } from "../agents/runtime/index.js"; +import type { AssistantMessage } from "../llm/types.js"; const SIZE_FRAME_ID = "00000000-0000-4000-8000-000000000000"; +export function cloneTextContent(part: { type: "text"; text: string; textSignature?: string }) { + return { + type: "text" as const, + text: part.text, + ...(part.textSignature ? { textSignature: part.textSignature } : {}), + }; +} + +export function cloneImageContent(part: { type: "image"; data: string; mimeType: string }) { + return { type: "image" as const, data: part.data, mimeType: part.mimeType }; +} + +export function cloneUsage( + message: AssistantMessage, +): WorkerTranscriptMessage & { role: "assistant" } { + return { + role: "assistant", + content: message.content.map((part) => { + if (part.type === "text") { + return cloneTextContent(part); + } + if (part.type === "thinking") { + return { + type: "thinking" as const, + thinking: part.thinking, + ...(part.thinkingSignature ? { thinkingSignature: part.thinkingSignature } : {}), + ...(part.redacted === undefined ? {} : { redacted: part.redacted }), + }; + } + return { + type: "toolCall" as const, + id: part.id, + name: part.name, + arguments: structuredClone(part.arguments), + ...(part.thoughtSignature ? { thoughtSignature: part.thoughtSignature } : {}), + ...(part.executionMode ? { executionMode: part.executionMode } : {}), + }; + }), + api: message.api, + provider: message.provider, + model: message.model, + ...(message.responseModel ? { responseModel: message.responseModel } : {}), + ...(message.responseId ? { responseId: message.responseId } : {}), + ...(message.diagnostics + ? { + diagnostics: message.diagnostics.map((diagnostic) => ({ + type: diagnostic.type, + timestamp: diagnostic.timestamp, + ...(diagnostic.error + ? { + error: { + ...(diagnostic.error.name ? { name: diagnostic.error.name } : {}), + message: diagnostic.error.message, + ...(diagnostic.error.stack ? { stack: diagnostic.error.stack } : {}), + ...(diagnostic.error.code === undefined ? {} : { code: diagnostic.error.code }), + }, + } + : {}), + ...(diagnostic.details ? { details: structuredClone(diagnostic.details) } : {}), + })), + } + : {}), + usage: { + input: message.usage.input, + output: message.usage.output, + cacheRead: message.usage.cacheRead, + cacheWrite: message.usage.cacheWrite, + ...(message.usage.contextUsage + ? { contextUsage: structuredClone(message.usage.contextUsage) } + : {}), + totalTokens: message.usage.totalTokens, + cost: { + input: message.usage.cost.input, + output: message.usage.cost.output, + cacheRead: message.usage.cost.cacheRead, + cacheWrite: message.usage.cost.cacheWrite, + total: message.usage.cost.total, + ...(message.usage.cost.totalOrigin ? { totalOrigin: message.usage.cost.totalOrigin } : {}), + }, + }, + stopReason: message.stopReason, + ...(message.errorMessage ? { errorMessage: message.errorMessage } : {}), + ...(message.errorCode ? { errorCode: message.errorCode } : {}), + ...(message.errorType ? { errorType: message.errorType } : {}), + ...(message.errorBody ? { errorBody: message.errorBody } : {}), + timestamp: message.timestamp, + }; +} + +export function toWorkerTranscriptMessage( + message: AgentMessage, +): WorkerTranscriptMessage | undefined { + if (message.role === "user") { + const content = + typeof message.content === "string" + ? [{ type: "text" as const, text: message.content }] + : message.content.map((part) => + part.type === "text" ? cloneTextContent(part) : cloneImageContent(part), + ); + return { role: "user", content, timestamp: message.timestamp }; + } + if (message.role === "assistant") { + return cloneUsage(message); + } + if (message.role === "toolResult") { + return { + role: "toolResult", + toolCallId: message.toolCallId, + toolName: message.toolName, + content: message.content.map((part) => + part.type === "text" ? cloneTextContent(part) : cloneImageContent(part), + ), + ...(message.details === undefined ? {} : { details: structuredClone(message.details) }), + isError: message.isError, + timestamp: message.timestamp, + }; + } + return undefined; +} + export function isWorkerTranscriptMessageFrameSafe(message: WorkerTranscriptMessage): boolean { const frame: WorkerTranscriptCommitRequestFrame = { type: "req", diff --git a/src/worker/worker.fault-injection.test.ts b/src/worker/worker.fault-injection.test.ts index 4af879340a02..d039ea4729bf 100644 --- a/src/worker/worker.fault-injection.test.ts +++ b/src/worker/worker.fault-injection.test.ts @@ -320,6 +320,7 @@ class ComposedGatewayHarness { baseLeafId?: string | null; initialSeq?: number; initialAckedSeq?: number; + runId?: string; } = {}, ): WorkerClients { const epoch = params.epoch ?? this.epoch; @@ -336,12 +337,13 @@ class ComposedGatewayHarness { handshake: HANDSHAKE, }, assignment: { - runId: RUN_ID, + runId: params.runId ?? RUN_ID, turnId: "fault-turn", prompt: "fault injection", workspaceDir: this.root, modelRef: MODEL_REF, inferenceOptions: {}, + suppressPromptTranscript: false, initialMessages: [], transcript: { baseLeafId: params.baseLeafId ?? null, nextSeq: params.initialSeq ?? 1 }, liveEvents: { @@ -857,10 +859,13 @@ describe("cloud worker milestone 2 fault injection", () => { await oldInferenceRejected; harness.providerPlan = { kind: "immediate", text: "new owner reply" }; + // Milestone-3 admission binds the worker to a single run; the fresh owner + // must be admitted for the run it executes. const fresh = harness.createClients({ admissionProof: REPLACEMENT_CREDENTIAL, epoch: newEpoch, baseLeafId: oldCommit.newLeafId, + runId: "fresh-run", }); clients.push(fresh); await fresh.connection.start(); diff --git a/src/worker/worker.runtime.test.ts b/src/worker/worker.runtime.test.ts index a32f19eb7a2e..6ca2737ff8ef 100644 --- a/src/worker/worker.runtime.test.ts +++ b/src/worker/worker.runtime.test.ts @@ -128,6 +128,7 @@ class FakeWorkerGateway { readonly acceptedTranscriptRequests: WorkerTranscriptCommitParams[] = []; readonly liveEventRequests: WorkerLiveEventParams[] = []; readonly inferenceRequests: WorkerInferenceStartParams[] = []; + readonly applicationOrder: string[] = []; constructor(private readonly options: FakeGatewayOptions = {}) { this.httpServer = createServer(); @@ -315,6 +316,7 @@ class FakeWorkerGateway { return; } this.acceptedTranscriptRequests.push(structuredClone(frame.params)); + this.applicationOrder.push(`transcript:${frame.params.seq}`); this.send(socket, { type: "res", id: frame.id, @@ -331,6 +333,11 @@ class FakeWorkerGateway { private handleLiveEvent(socket: WebSocket, frame: WorkerLiveEventRequestFrame): void { this.methods.push(frame.method); this.liveEventRequests.push(structuredClone(frame.params)); + this.applicationOrder.push( + frame.params.event.kind === "lifecycle" + ? `live:lifecycle:${frame.params.event.payload.phase}` + : `live:${frame.params.event.kind}`, + ); if (this.options.silenceFirstLiveEvent && !this.droppedLiveEvent) { this.droppedLiveEvent = true; return; @@ -706,6 +713,7 @@ function descriptor(socketPath: string, workspaceDir: string): WorkerLaunchDescr runId: RUN_ID, turnId: "worker-turn", prompt: "Complete the worker turn.", + suppressPromptTranscript: false, workspaceDir, modelRef: MODEL_REF, inferenceOptions: { reasoning: "off" }, @@ -754,6 +762,14 @@ describe("worker runtime", () => { expect(gateway.inferenceRequests[0]?.context.systemPrompt).toContain("worker-bootstrap-marker"); const toolNames = gateway.inferenceRequests[0]?.context.tools?.map((tool) => tool.name) ?? []; expect(toolNames).toHaveLength(6); + const terminalIndex = gateway.applicationOrder.findIndex( + (entry) => entry === "live:lifecycle:end", + ); + const finalTranscriptIndex = gateway.applicationOrder.findLastIndex((entry) => + entry.startsWith("transcript:"), + ); + expect(finalTranscriptIndex).toBeGreaterThanOrEqual(0); + expect(terminalIndex).toBeGreaterThan(finalTranscriptIndex); expect(toolNames).toEqual( expect.arrayContaining(["read", "write", "edit", "apply_patch", "exec", "process"]), ); @@ -806,7 +822,7 @@ describe("worker runtime", () => { gateway.liveEventRequests.some( (request) => request.event.kind === "lifecycle" && request.event.payload.phase === "error", ), - ).toBe(true); + ).toBe(false); }); it("renumbers live events after a gateway cursor reset without aborting the run", async () => { @@ -824,10 +840,10 @@ describe("worker runtime", () => { expect(gateway.liveEventRequests[1]?.event).toEqual(gateway.liveEventRequests[0]?.event); }); - it("degrades hard live-event failures without affecting inference or transcript commits", async () => { + it("requires authoritative terminal delivery after degrading preview live events", async () => { const { gateway, launch } = await setup({ liveFailure: "capacity-exceeded" }); - await expect(runWorkerDescriptor(launch)).resolves.toMatchObject({ status: "completed" }); + await expect(runWorkerDescriptor(launch)).rejects.toThrow("worker live event rejected"); expect(gateway.inferenceRequests).toHaveLength(1); expect( @@ -835,7 +851,11 @@ describe("worker runtime", () => { .flatMap((request) => request.messages) .map((message) => message.role), ).toEqual(["user", "assistant"]); - expect(gateway.liveEventRequests).toHaveLength(1); + expect(gateway.liveEventRequests.length).toBeGreaterThanOrEqual(2); + expect(gateway.liveEventRequests.at(-1)?.event).toMatchObject({ + kind: "lifecycle", + payload: { phase: "end" }, + }); }); it("degrades a repeated no-progress live resync without hanging the run", async () => { @@ -849,7 +869,11 @@ describe("worker runtime", () => { expect(gateway.inferenceRequests).toHaveLength(1); expect(gateway.acceptedTranscriptRequests).toHaveLength(2); - expect(gateway.liveEventRequests).toHaveLength(2); + expect(gateway.liveEventRequests).toHaveLength(3); + expect(gateway.liveEventRequests.at(-1)?.event).toMatchObject({ + kind: "lifecycle", + payload: { phase: "end" }, + }); }); it("fails closed when worker admission is rejected", async () => { @@ -887,6 +911,10 @@ describe("worker runtime", () => { await expect(result).rejects.toThrow("operator stopped worker"); expect(gateway.methods).toContain("worker.inference.cancel"); + expect(gateway.liveEventRequests.at(-1)?.event).toMatchObject({ + kind: "lifecycle", + payload: { phase: "end", aborted: true }, + }); }); it("bounds shutdown when remote inference cancellation cannot settle", async () => { @@ -905,14 +933,17 @@ describe("worker runtime", () => { }); it.each([ - ["error", "fixture provider failed", "error", "error"], - ["cancelled", "fixture inference cancelled", "aborted", "end"], + ["error", "error", "error"], + ["cancelled", "aborted", "end"], ] as const)( "reports remote inference %s terminals as failed turns", - async (plan, message, stopReason, lifecyclePhase) => { + async (plan, stopReason, lifecyclePhase) => { const { gateway, launch } = await setup({ inferencePlans: [plan] }); - await expect(runWorkerDescriptor(launch)).rejects.toThrow(message); + await expect(runWorkerDescriptor(launch)).resolves.toEqual({ + status: "failed", + reason: "turn-failed", + }); const assistant = gateway.transcriptRequests .flatMap((request) => request.messages) .toReversed() @@ -926,6 +957,19 @@ describe("worker runtime", () => { }, ); + it("keeps an unacknowledged failed-turn terminal as an infrastructure failure", async () => { + const { gateway, launch } = await setup({ + inferencePlans: ["error"], + liveFailure: "capacity-exceeded", + }); + + await expect(runWorkerDescriptor(launch)).rejects.toThrow("worker live event rejected"); + expect(gateway.liveEventRequests.at(-1)?.event).toMatchObject({ + kind: "lifecycle", + payload: { phase: "error" }, + }); + }); + it("fails closed when a heartbeat is rejected without fencing", async () => { const { launch } = await setup({ inferencePlans: ["hold"], @@ -958,9 +1002,10 @@ describe("worker runtime", () => { async (plan) => { const { gateway, launch } = await setup({ inferencePlans: [plan] }); - await expect(runWorkerDescriptor(launch)).rejects.toThrow( - "Worker inference result exceeds the transcript message limit.", - ); + await expect(runWorkerDescriptor(launch)).resolves.toEqual({ + status: "failed", + reason: "turn-failed", + }); const assistant = gateway.transcriptRequests .flatMap((request) => request.messages) .toReversed() diff --git a/src/worker/worker.runtime.ts b/src/worker/worker.runtime.ts index e9a9694f9b1d..7dd3ae538292 100644 --- a/src/worker/worker.runtime.ts +++ b/src/worker/worker.runtime.ts @@ -9,8 +9,11 @@ import { WorkerTranscriptCommitClient, } from "./worker-rpc-clients.js"; -type WorkerRuntimeResult = +// Cross-process contract: serialized to stdout by runWorkerCommand and parsed by the +// gateway worker turn launcher. +export type WorkerRuntimeResult = | { status: "completed"; transcriptLeafId: string | null; transcriptNextSeq: number } + | { status: "failed"; reason: "turn-failed" } | { status: "fenced"; reason: "credential-replaced" | "owner-epoch-mismatch" }; const WORKER_REMOTE_CANCEL_GRACE_MS = 1_000; @@ -52,6 +55,7 @@ export async function runWorkerDescriptor( const abortController = new AbortController(); let turnStarted = false; + let terminalLiveAcked = false; let forcedStopTimer: NodeJS.Timeout | undefined; const connection = createWorkerConnection({ socketPath: descriptor.socketPath, @@ -121,6 +125,7 @@ export async function runWorkerDescriptor( sessionKey: `worker:${descriptor.admission.sessionId}`, runId: descriptor.assignment.runId, prompt: descriptor.assignment.prompt, + suppressPromptTranscript: descriptor.assignment.suppressPromptTranscript, modelRef: descriptor.assignment.modelRef, initialMessages: descriptor.assignment.initialMessages, ...(descriptor.assignment.systemPrompt === undefined @@ -136,6 +141,12 @@ export async function runWorkerDescriptor( live: { emit: async (event) => { await live.emit(descriptor.assignment.runId, event); + if ( + event.kind === "lifecycle" && + (event.payload.phase === "end" || event.payload.phase === "error") + ) { + terminalLiveAcked = true; + } }, }, signal: abortController.signal, @@ -148,6 +159,12 @@ export async function runWorkerDescriptor( if (fenced) { return fenced; } + if (options.signal?.aborted) { + throw toError(options.signal.reason, "worker interrupted"); + } + if (terminalLiveAcked && connection.state.kind === "ready") { + return { status: "failed", reason: "turn-failed" }; + } throw toError(error, "worker session failed"); } const fenced = fencedResult(connection.state);