From 07d9f9c92ca24db6b04f93be0dfa2dac3e55b700 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 11 Aug 2026 23:58:06 -0700 Subject: [PATCH] refactor(gateway): extract source-agnostic desktop relay core (#122455) * refactor(gateway): extract source-agnostic desktop relay core Split the cloud-worker desktop observer into a reusable core under src/gateway/desktop/ so upcoming desktop sources (gateway host, nodes) plug into one relay pipeline: - attachment.ts: RfbAttachment union (unix socket | loopback TCP) - session-registry.ts: generic observer lifecycle (8-observer cap, controller eviction, 60s linger, owner-epoch fencing) with typed stale-owner/stopped errors - observe-bridge.ts: single-use observe tokens + WS relay, path moves to /desktop/observe (clients consume wsPath verbatim) - rfb-view-only-filter.ts: pure move desktop-tunnel.ts keeps only worker SSH acquisition and app launch, plugged into the shared registry; external API unchanged. Behavior is identical; worker desktop tests pass unmodified apart from import and path renames. Drops the never-used now() option. * fix(gateway): fence only superseded desktop sessions The desktop core extraction replaced the owner fence's "stop strictly older owners" check with an unconditional session stop. A launcher that claims an owner epoch first, then reaches its async fencing pass after a same-epoch observe has already created the session, tore that session down and failed the observer with "stopped before connecting". Restore the original invariant in the registry that owns it: stopSuperseded() retires an entry only when its epoch is strictly lower than the claimant's, so peers sharing a generation keep the session. The regression test drives launch-then-observe at one epoch and fails on the pre-fix code inside fenceReplacedOwners. * refactor(gateway): drop unused WorkerDesktopTunnels type export --- src/gateway/desktop/attachment.test.ts | 48 ++ src/gateway/desktop/attachment.ts | 11 + .../observe-bridge.test.ts} | 32 +- .../observe-bridge.ts} | 46 +- .../rfb-view-only-filter.test.ts | 0 .../rfb-view-only-filter.ts | 0 src/gateway/desktop/session-registry.ts | 254 +++++++++++ src/gateway/server-http.ts | 15 +- .../server-methods/environments.test.ts | 6 +- src/gateway/server-runtime-state-prepare.ts | 9 +- src/gateway/server-runtime-state.ts | 6 +- .../server-worker-environment-startup.ts | 6 +- ...server.plugin-node-capability-auth.test.ts | 33 +- .../desktop-tunnel.test.ts | 22 + .../worker-environments/desktop-tunnel.ts | 413 +++++++----------- .../environment-access.test.ts | 4 +- .../worker-environments/environment-access.ts | 14 +- src/gateway/worker-environments/tunnel.ts | 7 +- .../components/desktop/desktop-client.test.ts | 20 +- ui/src/e2e/desktop-panel.e2e.test.ts | 6 +- 20 files changed, 594 insertions(+), 358 deletions(-) create mode 100644 src/gateway/desktop/attachment.test.ts create mode 100644 src/gateway/desktop/attachment.ts rename src/gateway/{worker-environments/desktop-observe.test.ts => desktop/observe-bridge.test.ts} (91%) rename src/gateway/{worker-environments/desktop-observe.ts => desktop/observe-bridge.ts} (80%) rename src/gateway/{worker-environments => desktop}/rfb-view-only-filter.test.ts (100%) rename src/gateway/{worker-environments => desktop}/rfb-view-only-filter.ts (100%) create mode 100644 src/gateway/desktop/session-registry.ts diff --git a/src/gateway/desktop/attachment.test.ts b/src/gateway/desktop/attachment.test.ts new file mode 100644 index 000000000000..3bf7e639b8ee --- /dev/null +++ b/src/gateway/desktop/attachment.test.ts @@ -0,0 +1,48 @@ +import net from "node:net"; +import { afterEach, describe, expect, it } from "vitest"; +import { connectRfbAttachment } from "./attachment.js"; + +const servers: net.Server[] = []; +const sockets: net.Socket[] = []; + +afterEach(async () => { + for (const socket of sockets.splice(0)) { + socket.destroy(); + } + await Promise.all( + servers.splice(0).map( + (server) => + new Promise((resolve) => { + server.close(() => resolve()); + }), + ), + ); +}); + +describe("RFB attachments", () => { + it("connects a loopback TCP attachment", async () => { + const accepted = new Promise((resolve) => { + const server = net.createServer((socket) => { + sockets.push(socket); + resolve(); + }); + servers.push(server); + server.listen(0, "127.0.0.1"); + }); + const server = servers[0]; + if (!server) { + throw new Error("expected TCP test server"); + } + await new Promise((resolve) => { + server.once("listening", resolve); + }); + const address = server.address(); + if (!address || typeof address === "string") { + throw new Error("expected TCP test server address"); + } + + sockets.push(connectRfbAttachment({ kind: "tcp", host: "127.0.0.1", port: address.port })); + + await expect(accepted).resolves.toBeUndefined(); + }); +}); diff --git a/src/gateway/desktop/attachment.ts b/src/gateway/desktop/attachment.ts new file mode 100644 index 000000000000..bde7cdcb7623 --- /dev/null +++ b/src/gateway/desktop/attachment.ts @@ -0,0 +1,11 @@ +import net from "node:net"; + +export type RfbAttachment = + | { kind: "unix-socket"; socketPath: string } + | { kind: "tcp"; host: "127.0.0.1"; port: number }; + +export function connectRfbAttachment(attachment: RfbAttachment): net.Socket { + return attachment.kind === "unix-socket" + ? net.connect(attachment.socketPath) + : net.connect(attachment.port, attachment.host); +} diff --git a/src/gateway/worker-environments/desktop-observe.test.ts b/src/gateway/desktop/observe-bridge.test.ts similarity index 91% rename from src/gateway/worker-environments/desktop-observe.test.ts rename to src/gateway/desktop/observe-bridge.test.ts index 4ae72dcc28e2..a20ddc1a3818 100644 --- a/src/gateway/worker-environments/desktop-observe.test.ts +++ b/src/gateway/desktop/observe-bridge.test.ts @@ -6,10 +6,10 @@ import path from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; import { WebSocket } from "ws"; import { - handleWorkerDesktopUpgrade, - mintWorkerDesktopObserverToken, - WORKER_DESKTOP_OBSERVE_PATH, -} from "./desktop-observe.js"; + DESKTOP_OBSERVE_PATH, + handleDesktopObserveUpgrade, + mintDesktopObserverToken, +} from "./observe-bridge.js"; const cleanup: Array<() => Promise> = []; @@ -20,11 +20,11 @@ afterEach(async () => { describe("worker desktop observer tokens", () => { it("mints opaque tokens that expire after 60 seconds", () => { - const minted = mintWorkerDesktopObserverToken({ - environmentId: "worker:one", + const minted = mintDesktopObserverToken({ + sourceKey: "worker:one", ownerEpoch: 3, control: true, - localSocketPath: "/tmp/desktop.sock", + attachment: { kind: "unix-socket", socketPath: "/tmp/desktop.sock" }, nowMs: 1_000, }); expect(minted.token).toMatch(/^[a-f0-9]{48}$/u); @@ -55,8 +55,8 @@ async function createProxyHarness( const closeObserver = vi.fn(); const httpServer = http.createServer(); httpServer.on("upgrade", (req, socket, head) => { - handleWorkerDesktopUpgrade(req, socket, head, { - tunnels: { + handleDesktopObserveUpgrade(req, socket, head, { + registry: { attachObserver: (_environmentId, observer) => { closeObserver.mockImplementation((code: number, reason: string) => { observer.close(code, reason); @@ -81,14 +81,14 @@ async function createProxyHarness( }); await fs.rm(root, { recursive: true, force: true }); }); - const minted = mintWorkerDesktopObserverToken({ - environmentId: "worker:pump", + const minted = mintDesktopObserverToken({ + sourceKey: "worker:pump", ownerEpoch: 2, control: params.control ?? false, - localSocketPath, + attachment: { kind: "unix-socket", socketPath: localSocketPath }, }); const ws = new WebSocket( - `ws://127.0.0.1:${address.port}${WORKER_DESKTOP_OBSERVE_PATH}?token=${minted.token}`, + `ws://127.0.0.1:${address.port}${DESKTOP_OBSERVE_PATH}?token=${minted.token}`, ); cleanup.push(async () => ws.terminate()); await new Promise((resolve, reject) => { @@ -133,11 +133,11 @@ describe("worker desktop observer proxy", () => { const harness = await createProxyHarness(); await expectUnauthorizedObserver(harness.observerUrl); - const expired = mintWorkerDesktopObserverToken({ - environmentId: "worker:expired", + const expired = mintDesktopObserverToken({ + sourceKey: "worker:expired", ownerEpoch: 1, control: false, - localSocketPath: "/tmp/expired.sock", + attachment: { kind: "unix-socket", socketPath: "/tmp/expired.sock" }, nowMs: 0, }); const observerUrl = new URL(harness.observerUrl); diff --git a/src/gateway/worker-environments/desktop-observe.ts b/src/gateway/desktop/observe-bridge.ts similarity index 80% rename from src/gateway/worker-environments/desktop-observe.ts rename to src/gateway/desktop/observe-bridge.ts index 75640c27eaf3..5d65c9e676a7 100644 --- a/src/gateway/worker-environments/desktop-observe.ts +++ b/src/gateway/desktop/observe-bridge.ts @@ -1,30 +1,30 @@ import crypto from "node:crypto"; import type { IncomingMessage } from "node:http"; -import net from "node:net"; import type { Duplex } from "node:stream"; import { WebSocket, WebSocketServer, type RawData } from "ws"; -import type { WorkerDesktopTunnels } from "./desktop-tunnel.js"; +import { connectRfbAttachment, type RfbAttachment } from "./attachment.js"; import { createRfbClientMessageFilter } from "./rfb-view-only-filter.js"; +import type { DesktopSessionRegistry } from "./session-registry.js"; -export const WORKER_DESKTOP_OBSERVE_PATH = "/worker-desktop/observe"; +export const DESKTOP_OBSERVE_PATH = "/desktop/observe"; const TOKEN_TTL_MS = 60_000; const TOKEN_PATTERN = /^[a-f0-9]{48}$/u; const MAX_PAYLOAD_BYTES = 1024 * 1024; const PAUSE_BUFFERED_BYTES = 4 * 1024 * 1024; const RESUME_CHECK_MS = 25; -type WorkerDesktopObserverTokenEntry = { - environmentId: string; +type DesktopObserverTokenEntry = { + sourceKey: string; ownerEpoch: number; control: boolean; - localSocketPath: string; + attachment: RfbAttachment; expiresAt: number; }; -const observerTokens = new Map(); +const observerTokens = new Map(); const desktopObserverWss = new WebSocketServer({ noServer: true, maxPayload: MAX_PAYLOAD_BYTES }); -function pruneWorkerDesktopObserverTokens(nowMs: number): void { +function pruneDesktopObserverTokens(nowMs: number): void { for (const [token, entry] of observerTokens) { if (entry.expiresAt <= nowMs) { observerTokens.delete(token); @@ -32,32 +32,32 @@ function pruneWorkerDesktopObserverTokens(nowMs: number): void { } } -export function mintWorkerDesktopObserverToken(params: { - environmentId: string; +export function mintDesktopObserverToken(params: { + sourceKey: string; ownerEpoch: number; control: boolean; - localSocketPath: string; + attachment: RfbAttachment; nowMs?: number; }): { token: string; expiresAtMs: number } { const nowMs = params.nowMs ?? Date.now(); - pruneWorkerDesktopObserverTokens(nowMs); + pruneDesktopObserverTokens(nowMs); const token = crypto.randomBytes(24).toString("hex"); const expiresAtMs = nowMs + TOKEN_TTL_MS; observerTokens.set(token, { - environmentId: params.environmentId, + sourceKey: params.sourceKey, ownerEpoch: params.ownerEpoch, control: params.control, - localSocketPath: params.localSocketPath, + attachment: params.attachment, expiresAt: expiresAtMs, }); return { token, expiresAtMs }; } -function consumeWorkerDesktopObserverToken( +function consumeDesktopObserverToken( token: string, nowMs = Date.now(), -): WorkerDesktopObserverTokenEntry | undefined { - pruneWorkerDesktopObserverTokens(nowMs); +): DesktopObserverTokenEntry | undefined { + pruneDesktopObserverTokens(nowMs); const normalized = token.trim(); if (!TOKEN_PATTERN.test(normalized)) { return undefined; @@ -86,28 +86,28 @@ function rawDataBuffer(data: RawData): Buffer { } /** Upgrades one authenticated observer token into a raw bidirectional RFB stream. */ -export function handleWorkerDesktopUpgrade( +export function handleDesktopObserveUpgrade( req: IncomingMessage, socket: Duplex, head: Buffer, deps: { - tunnels: Pick; + registry: Pick; getBufferedAmount?: (ws: WebSocket) => number; }, ): boolean { const resource = new URL(req.url ?? "/", "http://127.0.0.1"); - if (resource.pathname !== WORKER_DESKTOP_OBSERVE_PATH) { + if (resource.pathname !== DESKTOP_OBSERVE_PATH) { return false; } const token = resource.searchParams.get("token") ?? ""; - const entry = consumeWorkerDesktopObserverToken(token); + const entry = consumeDesktopObserverToken(token); if (!entry) { writeUnauthorized(socket); return true; } desktopObserverWss.handleUpgrade(req, socket, head, (ws) => { // View-only is enforced here at the RFB message boundary; the UI setting is only UX. - const observer = deps.tunnels.attachObserver(entry.environmentId, { + const observer = deps.registry.attachObserver(entry.sourceKey, { control: entry.control, ownerEpoch: entry.ownerEpoch, close: (code, reason) => ws.close(code, reason), @@ -116,7 +116,7 @@ export function handleWorkerDesktopUpgrade( ws.close(1013, "desktop observer limit"); return; } - const desktopSocket = net.connect(entry.localSocketPath); + const desktopSocket = connectRfbAttachment(entry.attachment); const clientMessageFilter = entry.control ? undefined : createRfbClientMessageFilter(); let closed = false; let resumeTimer: ReturnType | undefined; diff --git a/src/gateway/worker-environments/rfb-view-only-filter.test.ts b/src/gateway/desktop/rfb-view-only-filter.test.ts similarity index 100% rename from src/gateway/worker-environments/rfb-view-only-filter.test.ts rename to src/gateway/desktop/rfb-view-only-filter.test.ts diff --git a/src/gateway/worker-environments/rfb-view-only-filter.ts b/src/gateway/desktop/rfb-view-only-filter.ts similarity index 100% rename from src/gateway/worker-environments/rfb-view-only-filter.ts rename to src/gateway/desktop/rfb-view-only-filter.ts diff --git a/src/gateway/desktop/session-registry.ts b/src/gateway/desktop/session-registry.ts new file mode 100644 index 000000000000..ff8154a366aa --- /dev/null +++ b/src/gateway/desktop/session-registry.ts @@ -0,0 +1,254 @@ +import type { RfbAttachment } from "./attachment.js"; + +const DEFAULT_LINGER_MS = 60_000; +const MAX_OBSERVERS = 8; + +export class DesktopSessionStaleOwnerError extends Error { + constructor() { + super("Desktop session owner epoch is stale"); + this.name = "DesktopSessionStaleOwnerError"; + } +} + +export class DesktopSessionStoppedError extends Error { + constructor() { + super("Desktop session stopped before connecting"); + this.name = "DesktopSessionStoppedError"; + } +} + +type DesktopSessionObserver = { + control: boolean; + /** Epoch the observer token was minted against; a stale token must not reach a newer entry. */ + ownerEpoch: number; + close(code: number, reason: string): void; +}; + +type DesktopSessionAcquireResult = { + attachment: RfbAttachment; + vncPassword?: string; +}; + +type DesktopSessionAcquireRequest = { + sourceKey: string; + ownerEpoch: number; + start: (isCurrent: () => boolean) => Promise; + teardown?: () => Promise; +}; + +type ObserverEntry = DesktopSessionObserver & { released: boolean }; +type DesktopSessionEntry = { + sourceKey: string; + ownerEpoch: number; + initialization?: Promise; + stopPromise?: Promise; + ready: Promise; + resolveReady: (result: DesktopSessionAcquireResult) => void; + rejectReady: (error: Error) => void; + readySettled: boolean; + observers: Set; + controller?: ObserverEntry; + lingerTimer?: ReturnType; + stopped: boolean; + start: DesktopSessionAcquireRequest["start"]; + teardown?: DesktopSessionAcquireRequest["teardown"]; +}; + +/** Owns per-source desktop sessions and their connected observer lifetimes. */ +export function createDesktopSessionRegistry( + deps: { + lingerMs?: number; + } = {}, +) { + const lingerMs = deps.lingerMs ?? DEFAULT_LINGER_MS; + const entries = new Map(); + const claimedOwnerEpochs = new Map(); + + const claimOwnerEpoch = (sourceKey: string, ownerEpoch: number): boolean => { + const claimedEpoch = claimedOwnerEpochs.get(sourceKey); + if (claimedEpoch !== undefined && ownerEpoch < claimedEpoch) { + throw new DesktopSessionStaleOwnerError(); + } + if (claimedEpoch === undefined || ownerEpoch > claimedEpoch) { + claimedOwnerEpochs.set(sourceKey, ownerEpoch); + return true; + } + return false; + }; + + const isCurrent = (entry: DesktopSessionEntry) => + entries.get(entry.sourceKey) === entry && !entry.stopped; + + const closeObserver = (observer: ObserverEntry, code: number, reason: string) => { + try { + observer.close(code, reason); + } catch { + // Observer cleanup remains authoritative when the transport close callback fails. + } + }; + + const stopEntry = (entry: DesktopSessionEntry): Promise => { + if (entry.stopPromise) { + return entry.stopPromise; + } + entry.stopPromise = (async () => { + entry.stopped = true; + if (entries.get(entry.sourceKey) === entry) { + entries.delete(entry.sourceKey); + } + clearTimeout(entry.lingerTimer); + entry.lingerTimer = undefined; + for (const observer of entry.observers) { + observer.released = true; + closeObserver(observer, 1012, "desktop tunnel closed"); + } + entry.observers.clear(); + entry.controller = undefined; + if (!entry.readySettled) { + entry.readySettled = true; + entry.rejectReady(new DesktopSessionStoppedError()); + } + // Teardown brackets initialization so a source can stop the currently published + // transport, then dispose anything initialization publishes before it settles. + await entry.teardown?.().catch(() => undefined); + await entry.initialization?.catch(() => undefined); + await entry.teardown?.().catch(() => undefined); + })(); + return entry.stopPromise; + }; + + async function acquire( + request: DesktopSessionAcquireRequest, + ): Promise { + claimOwnerEpoch(request.sourceKey, request.ownerEpoch); + const current = entries.get(request.sourceKey); + if (current) { + if (request.ownerEpoch < current.ownerEpoch) { + throw new DesktopSessionStaleOwnerError(); + } + if (request.ownerEpoch === current.ownerEpoch) { + return await current.ready; + } + } + + let resolveReady!: (result: DesktopSessionAcquireResult) => void; + let rejectReady!: (error: Error) => void; + const ready = new Promise((resolve, reject) => { + resolveReady = resolve; + rejectReady = reject; + }); + void ready.catch(() => undefined); + const entry: DesktopSessionEntry = { + sourceKey: request.sourceKey, + ownerEpoch: request.ownerEpoch, + ready, + resolveReady, + rejectReady, + readySettled: false, + observers: new Set(), + stopped: false, + start: request.start, + ...(request.teardown ? { teardown: request.teardown } : {}), + }; + entries.set(request.sourceKey, entry); + entry.initialization = (async () => { + if (current) { + await stopEntry(current); + } + if (!isCurrent(entry)) { + return; + } + const result = await entry.start(() => isCurrent(entry)); + if (!isCurrent(entry)) { + return; + } + entry.readySettled = true; + entry.resolveReady(result); + })(); + void entry.initialization.catch((error: unknown) => { + if (!entry.readySettled) { + entry.readySettled = true; + entry.rejectReady(error instanceof Error ? error : new Error("Desktop session failed")); + } + void stopEntry(entry); + }); + return await ready; + } + + function attachObserver(sourceKey: string, observer: DesktopSessionObserver) { + const entry = entries.get(sourceKey); + if (!entry || !entry.readySettled || entry.stopped || entry.observers.size >= MAX_OBSERVERS) { + return undefined; + } + // A token minted against a replaced entry must not reach this one; otherwise a stale + // control token would evict the current controller of a desktop it never observed. + if (observer.ownerEpoch !== entry.ownerEpoch) { + return undefined; + } + clearTimeout(entry.lingerTimer); + entry.lingerTimer = undefined; + if (observer.control && entry.controller) { + const previous = entry.controller; + previous.released = true; + entry.observers.delete(previous); + entry.controller = undefined; + closeObserver(previous, 4000, "control-taken"); + } + const attached: ObserverEntry = { ...observer, released: false }; + entry.observers.add(attached); + if (attached.control) { + entry.controller = attached; + } + return { + release() { + if (attached.released) { + return; + } + attached.released = true; + entry.observers.delete(attached); + if (entry.controller === attached) { + entry.controller = undefined; + } + if (entry.observers.size === 0 && isCurrent(entry)) { + entry.lingerTimer = setTimeout(() => void stopEntry(entry), lingerMs); + entry.lingerTimer.unref?.(); + } + }, + }; + } + + async function stop(sourceKey: string, ownerEpoch?: number): Promise { + const entry = entries.get(sourceKey); + if (entry && (ownerEpoch === undefined || ownerEpoch === entry.ownerEpoch)) { + await stopEntry(entry); + } + } + + /** + * Retires only owners strictly older than the claimant. An equal epoch shares the + * session, so fencing must not tear down a peer that claimed the same generation. + */ + async function stopSuperseded(sourceKey: string, ownerEpoch: number): Promise { + const entry = entries.get(sourceKey); + if (entry && entry.ownerEpoch < ownerEpoch) { + await stopEntry(entry); + } + } + + async function stopAll(): Promise { + await Promise.all([...entries.values()].map(stopEntry)); + } + + return { + acquire, + attachObserver, + claimOwnerEpoch, + isOwnerEpochCurrent: (sourceKey: string, ownerEpoch: number) => + claimedOwnerEpochs.get(sourceKey) === ownerEpoch, + stop, + stopSuperseded, + stopAll, + }; +} + +export type DesktopSessionRegistry = ReturnType; diff --git a/src/gateway/server-http.ts b/src/gateway/server-http.ts index 4a40852e4ef0..8c4e47e6529f 100644 --- a/src/gateway/server-http.ts +++ b/src/gateway/server-http.ts @@ -44,6 +44,7 @@ import { isControlUiPluginManagerRequest, } from "./control-ui-routing.js"; import type { ControlUiRootState } from "./control-ui.js"; +import type { DesktopSessionRegistry } from "./desktop/session-registry.js"; import { classifyGatewayProbePath, classifyMcpAppStandalonePath, @@ -79,7 +80,6 @@ import { } from "./server/ws-types.js"; import { isTerminalConfigEnabled } from "./terminal/enabled.js"; import { canonicalizeUserProfileAvatarPath } from "./user-profiles-http-path.js"; -import type { WorkerDesktopTunnels } from "./worker-environments/desktop-tunnel.js"; type PluginGatewayDispatchContext = { gatewayAuthSatisfied?: boolean; @@ -906,7 +906,7 @@ export function attachGatewayUpgradeHandler(opts: { rateLimiter?: AuthRateLimiter; /** Optional logger for error diagnostics. */ log?: { warn: (msg: string) => void }; - workerDesktopTunnels?: WorkerDesktopTunnels; + desktopSessionRegistry?: DesktopSessionRegistry; }) { const { httpServer, @@ -1006,8 +1006,8 @@ export function attachGatewayUpgradeHandler(opts: { return; } } - if (requestPath === "/worker-desktop/observe") { - if (!opts.workerDesktopTunnels) { + if (requestPath === "/desktop/observe") { + if (!opts.desktopSessionRegistry) { writeGatewayUpgradeServiceUnavailable(socket, "desktop observe unavailable"); socket.destroy(); return; @@ -1020,10 +1020,9 @@ export function attachGatewayUpgradeHandler(opts: { socket.destroy(); return; } - const { handleWorkerDesktopUpgrade } = - await import("./worker-environments/desktop-observe.js"); - handleWorkerDesktopUpgrade(req, socket, head, { - tunnels: opts.workerDesktopTunnels, + const { handleDesktopObserveUpgrade } = await import("./desktop/observe-bridge.js"); + handleDesktopObserveUpgrade(req, socket, head, { + registry: opts.desktopSessionRegistry, }); return; } diff --git a/src/gateway/server-methods/environments.test.ts b/src/gateway/server-methods/environments.test.ts index fd53bf1366ec..8283489e461f 100644 --- a/src/gateway/server-methods/environments.test.ts +++ b/src/gateway/server-methods/environments.test.ts @@ -130,7 +130,7 @@ function workerService(overrides: Partial = {}) { destroyUnattached: vi.fn(async () => workerRecord({ state: "destroyed" })), observeDesktop: vi.fn(async ({ control }) => ({ transport: "rfb" as const, - wsPath: "/worker-desktop/observe?token=abc", + wsPath: "/desktop/observe?token=abc", expiresAtMs: 70_000, control, })), @@ -474,7 +474,7 @@ describe("environment gateway methods", () => { it("starts desktop observation with explicit and default control modes", async () => { const observeDesktop = vi.fn(async ({ control }: { control: boolean }) => ({ transport: "rfb" as const, - wsPath: "/worker-desktop/observe?token=abc", + wsPath: "/desktop/observe?token=abc", expiresAtMs: 70_000, control, })); @@ -493,7 +493,7 @@ describe("environment gateway methods", () => { true, { transport: "rfb", - wsPath: "/worker-desktop/observe?token=abc", + wsPath: "/desktop/observe?token=abc", expiresAtMs: 70_000, control: true, }, diff --git a/src/gateway/server-runtime-state-prepare.ts b/src/gateway/server-runtime-state-prepare.ts index 097e35ea5d2d..cd07dfb04381 100644 --- a/src/gateway/server-runtime-state-prepare.ts +++ b/src/gateway/server-runtime-state-prepare.ts @@ -14,6 +14,7 @@ import type { RuntimeEnv } from "../runtime.js"; import { getActiveSecretsRuntimeConfigSnapshot } from "../secrets/runtime-state.js"; import { createAuthRateLimiter, type AuthRateLimiter } from "./auth-rate-limit.js"; import { resolveGatewayAuth } from "./auth.js"; +import { createDesktopSessionRegistry } from "./desktop/session-registry.js"; import { isLoopbackHost } from "./net.js"; import { createNodeReapprovalCoordinator } from "./node-reapproval-coordinator.js"; import { resolveGatewayPluginConfig } from "./runtime-plugin-config.js"; @@ -117,13 +118,17 @@ export async function prepareGatewayKernelState(params: { const workerGatewayEndpoint = { resolve: (() => undefined) as () => { host: "127.0.0.1" | "::1"; port: number } | undefined, }; + const desktopSessionRegistry = shouldStartWorkerEnvironmentService + ? createDesktopSessionRegistry() + : undefined; const workerEnvironmentRuntime = - workerEnvironmentStartup && shouldStartWorkerEnvironmentService + workerEnvironmentStartup && desktopSessionRegistry ? await startupTrace.measure("worker-environments.runtime-imports", async () => { const workerModule = await loadWorkerEnvironmentStartupModule(); return await workerModule.createGatewayWorkerEnvironmentRuntime({ getPluginRegistry: () => pluginRuntime.registry, resolveWorkerGateway: () => workerGatewayEndpoint.resolve(), + desktopSessionRegistry, startup: workerEnvironmentStartup, log, }); @@ -415,7 +420,7 @@ export async function prepareGatewayKernelState(params: { handleWatchNodeRequest: async (req: IncomingMessage, res: ServerResponse) => (await watchNodeRequestHandler.current?.(req, res)) ?? false, workerIngressEnabled: Boolean(workerEnvironmentService), - workerDesktopTunnels: workerTunnelManager?.desktop, + desktopSessionRegistry: workerTunnelManager ? desktopSessionRegistry : undefined, clients: connectionState.clients, }); const { diff --git a/src/gateway/server-runtime-state.ts b/src/gateway/server-runtime-state.ts index 03f95d017929..8cb2865c29b2 100644 --- a/src/gateway/server-runtime-state.ts +++ b/src/gateway/server-runtime-state.ts @@ -19,6 +19,7 @@ import type { PluginRegistry } from "../plugins/registry.js"; import type { AuthRateLimiter } from "./auth-rate-limit.js"; import type { ResolvedGatewayAuth } from "./auth.js"; import type { ControlUiRootState } from "./control-ui.js"; +import type { DesktopSessionRegistry } from "./desktop/session-registry.js"; import type { HooksConfigResolved } from "./hooks.js"; import type { AuthorizedGatewayHttpRequest } from "./http-auth-utils.js"; import { createSandboxHostHttpServer } from "./mcp-app-sandbox-http.js"; @@ -43,7 +44,6 @@ import { } from "./server/preauth-connection-budget.js"; import type { ReadinessChecker, StartupChecker } from "./server/readiness.js"; import type { GatewayWsClient } from "./server/ws-types.js"; -import type { WorkerDesktopTunnels } from "./worker-environments/desktop-tunnel.js"; type GatewayPluginRequestHandler = ( req: IncomingMessage, @@ -118,7 +118,7 @@ export async function createGatewayHttpTransport(params: { isTerminalEnabled: () => boolean; handleWatchNodeRequest?: (req: IncomingMessage, res: ServerResponse) => Promise; workerIngressEnabled?: boolean; - workerDesktopTunnels?: WorkerDesktopTunnels; + desktopSessionRegistry?: DesktopSessionRegistry; clients: Set; }): Promise<{ httpServer: HttpServer; @@ -302,7 +302,7 @@ export async function createGatewayHttpTransport(params: { getResolvedAuth: params.getResolvedAuth, rateLimiter: params.rateLimiter, log: params.log, - workerDesktopTunnels: params.workerDesktopTunnels, + desktopSessionRegistry: params.desktopSessionRegistry, }); gatewayHttpServers.push(httpServer); httpServers.push(httpServer); diff --git a/src/gateway/server-worker-environment-startup.ts b/src/gateway/server-worker-environment-startup.ts index c1f0c649966f..b3033216c7ed 100644 --- a/src/gateway/server-worker-environment-startup.ts +++ b/src/gateway/server-worker-environment-startup.ts @@ -6,6 +6,7 @@ import { getActiveSecretsRuntimeEnvState, } from "../secrets/runtime-state.js"; import { createLazyRuntimeModule } from "../shared/lazy-runtime.js"; +import type { DesktopSessionRegistry } from "./desktop/session-registry.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"; @@ -77,6 +78,7 @@ export async function loadGatewayWorkerEnvironmentStartupState(): Promise Pick; resolveWorkerGateway: () => WorkerGatewayEndpoint; + desktopSessionRegistry: DesktopSessionRegistry; startup: GatewayWorkerEnvironmentStartupState; log: WorkerEnvironmentLogger; }): Promise { @@ -144,7 +146,9 @@ export async function createGatewayWorkerEnvironmentRuntime(params: { startupBindings.map((binding) => [binding.environmentId, binding.runEpoch] as const), ), }); - const workerTunnelManager = createWorkerTunnelManager(); + const workerTunnelManager = createWorkerTunnelManager({ + desktopSessionRegistry: params.desktopSessionRegistry, + }); let executeSessionTool: ReturnType = async () => { throw new Error("Worker session tools are unavailable"); }; diff --git a/src/gateway/server.plugin-node-capability-auth.test.ts b/src/gateway/server.plugin-node-capability-auth.test.ts index e218716a534a..aaaf83ff28e1 100644 --- a/src/gateway/server.plugin-node-capability-auth.test.ts +++ b/src/gateway/server.plugin-node-capability-auth.test.ts @@ -15,16 +15,13 @@ import { import { withTimeout } from "../utils/with-timeout.js"; import { createAuthRateLimiter } from "./auth-rate-limit.js"; import type { ResolvedGatewayAuth } from "./auth.js"; +import { DESKTOP_OBSERVE_PATH, mintDesktopObserverToken } from "./desktop/observe-bridge.js"; import { PLUGIN_NODE_CAPABILITY_PATH_PREFIX } from "./plugin-node-capability.js"; import { MAX_PREAUTH_PAYLOAD_BYTES } from "./server-constants.js"; import { attachGatewayUpgradeHandler, createGatewayHttpServer } from "./server-http.js"; import { createPreauthConnectionBudget } from "./server/preauth-connection-budget.js"; import type { GatewayWsClient } from "./server/ws-types.js"; import { withTempConfig } from "./test-temp-config.js"; -import { - mintWorkerDesktopObserverToken, - WORKER_DESKTOP_OBSERVE_PATH, -} from "./worker-environments/desktop-observe.js"; const WS_REJECT_TIMEOUT_MS = 2_000; const WS_CONNECT_TIMEOUT_MS = 5_000; @@ -357,7 +354,9 @@ async function withCanvasGatewayHarness(params: { resolvePluginNodeCapabilityRoute?: Parameters< typeof attachGatewayUpgradeHandler >[0]["resolvePluginNodeCapabilityRoute"]; - workerDesktopTunnels?: Parameters[0]["workerDesktopTunnels"]; + desktopSessionRegistry?: Parameters< + typeof attachGatewayUpgradeHandler + >[0]["desktopSessionRegistry"]; run: (ctx: { listener: Awaited>; clients: Set; @@ -426,7 +425,7 @@ async function withCanvasGatewayHarness(params: { resolvedAuth: params.resolvedAuth, getResolvedAuth: params.getResolvedAuth, rateLimiter: params.rateLimiter, - workerDesktopTunnels: params.workerDesktopTunnels, + desktopSessionRegistry: params.desktopSessionRegistry, }); const listener = await listen(httpServer, params.listenHost); @@ -810,25 +809,25 @@ describe("gateway plugin node capability auth", () => { { message: "desktop unix server listen timed out" }, ); const release = vi.fn(); - const workerDesktopTunnels = { + const desktopSessionRegistry = { attachObserver: () => ({ release }), } as unknown as NonNullable< - Parameters[0]["workerDesktopTunnels"] + Parameters[0]["desktopSessionRegistry"] >; try { await withCanvasGatewayHarness({ resolvedAuth: tokenResolvedAuth, handleHttpRequest: async () => false, resolvePluginNodeCapabilityRoute: () => undefined, - workerDesktopTunnels, + desktopSessionRegistry, run: async ({ listener }) => { - const minted = mintWorkerDesktopObserverToken({ - environmentId: "worker:boundary", + const minted = mintDesktopObserverToken({ + sourceKey: "worker:boundary", ownerEpoch: 4, control: false, - localSocketPath, + attachment: { kind: "unix-socket", socketPath: localSocketPath }, }); - const url = `ws://127.0.0.1:${listener.port}${WORKER_DESKTOP_OBSERVE_PATH}?token=${minted.token}`; + const url = `ws://127.0.0.1:${listener.port}${DESKTOP_OBSERVE_PATH}?token=${minted.token}`; const ws = new WebSocket(url); const received = new Promise((resolve, reject) => { ws.once("message", (data) => resolve(Buffer.from(data as Buffer))); @@ -843,16 +842,16 @@ describe("gateway plugin node capability auth", () => { // A draining Gateway must refuse new desktop observers like every other // core upgrade; otherwise restart/suspension leaks long-lived sockets. - const draining = mintWorkerDesktopObserverToken({ - environmentId: "worker:boundary", + const draining = mintDesktopObserverToken({ + sourceKey: "worker:boundary", ownerEpoch: 4, control: false, - localSocketPath, + attachment: { kind: "unix-socket", socketPath: localSocketPath }, }); markGatewayRestartDraining(); try { await expectWsRejected( - `ws://127.0.0.1:${listener.port}${WORKER_DESKTOP_OBSERVE_PATH}?token=${draining.token}`, + `ws://127.0.0.1:${listener.port}${DESKTOP_OBSERVE_PATH}?token=${draining.token}`, {}, 503, ); diff --git a/src/gateway/worker-environments/desktop-tunnel.test.ts b/src/gateway/worker-environments/desktop-tunnel.test.ts index 069f9469c490..d39d5e597cf8 100644 --- a/src/gateway/worker-environments/desktop-tunnel.test.ts +++ b/src/gateway/worker-environments/desktop-tunnel.test.ts @@ -396,4 +396,26 @@ describe("worker desktop tunnels", () => { await expect(launchApp(failed)).rejects.toThrow("launcher failed"); await failed.stopAll(); }); + + it("keeps a same-epoch desktop session alive when an app launch fences replaced owners", async () => { + const fake = fakeRunner(); + const manager = createWorkerDesktopTunnels({ runner: fake.runner }); + // The launcher claims the epoch first, so its fencing pass runs after the + // observer session for that same epoch already exists. Fencing must only + // retire strictly older owners; equal epochs share the session. + const launching = launchApp(manager, "browser", 1); + const starting = acquire(manager, 1); + await waitForStarts(fake.starts, 1); + fake.starts[0]?.process.becomeReady(); + await starting; + await launching; + + const observer = manager.attachObserver("worker:one", { + control: false, + ownerEpoch: 1, + close: vi.fn(), + }); + expect(observer).toBeDefined(); + observer?.release(); + }); }); diff --git a/src/gateway/worker-environments/desktop-tunnel.ts b/src/gateway/worker-environments/desktop-tunnel.ts index b3f1e414be23..60940d4b9990 100644 --- a/src/gateway/worker-environments/desktop-tunnel.ts +++ b/src/gateway/worker-environments/desktop-tunnel.ts @@ -6,6 +6,13 @@ import type { WorkerDesktopEndpoint, WorkerSshEndpoint, } from "../../plugins/types.js"; +import type { RfbAttachment } from "../desktop/attachment.js"; +import { + createDesktopSessionRegistry, + DesktopSessionStaleOwnerError, + DesktopSessionStoppedError, + type DesktopSessionRegistry, +} from "../desktop/session-registry.js"; import { prepareWorkerSsh, type PreparedWorkerSsh, @@ -21,10 +28,8 @@ import { WORKER_TUNNEL_READY_MARKER, } from "./tunnel-ssh-runner.js"; -const DEFAULT_LINGER_MS = 60_000; const PASSWORD_READ_TIMEOUT_MS = 20_000; const APP_LAUNCH_TIMEOUT_MS = 30_000; -const MAX_OBSERVERS = 8; const REMOTE_DESKTOP_READY_SCRIPT = String.raw`set -eu printf '%s\n' '${WORKER_TUNNEL_READY_MARKER}' @@ -32,13 +37,6 @@ trap 'exit 0' HUP INT TERM while :; do sleep 3600; done `; -type WorkerDesktopObserver = { - control: boolean; - /** Epoch the observer token was minted against; a stale token must not reach a newer entry. */ - ownerEpoch: number; - close(code: number, reason: string): void; -}; - type DesktopAcquireRequest = { environmentId: string; ownerEpoch: number; @@ -47,25 +45,7 @@ type DesktopAcquireRequest = { resolveIdentity: WorkerSshIdentityResolver; }; -type DesktopAcquireResult = { localSocketPath: string; vncPassword?: string }; -type ObserverEntry = WorkerDesktopObserver & { released: boolean }; -type DesktopEntry = { - environmentId: string; - ownerEpoch: number; - localSocketPath?: string; - prepared?: PreparedWorkerSsh; - process?: WorkerSshProcess; - initialization?: Promise; - stopPromise?: Promise; - ready: Promise; - resolveReady: (result: DesktopAcquireResult) => void; - rejectReady: (error: Error) => void; - readySettled: boolean; - observers: Set; - controller?: ObserverEntry; - lingerTimer?: ReturnType; - stopped: boolean; -}; +type DesktopAcquireResult = { attachment: RfbAttachment; vncPassword?: string }; type DesktopAppLaunchEntry = { environmentId: string; @@ -88,65 +68,20 @@ function successful(result: Awaited>): boolea return result.termination === "exit" && result.code === 0; } -/** Owns per-environment local desktop forwards and their connected observer lifetimes. */ +/** Owns worker-specific desktop SSH acquisition and app launch processes. */ export function createWorkerDesktopTunnels(deps: { runner: WorkerSshRunner; - now?: () => number; + registry?: DesktopSessionRegistry; lingerMs?: number; platform?: NodeJS.Platform; }) { - const lingerMs = deps.lingerMs ?? DEFAULT_LINGER_MS; const platform = deps.platform ?? process.platform; - const entries = new Map(); + const sessions = deps.registry ?? createDesktopSessionRegistry({ lingerMs: deps.lingerMs }); const appLaunches = new Map(); - const claimedOwnerEpochs = new Map(); const appLaunchKey = (environmentId: string, appId: WorkerDesktopApp["id"]) => `${environmentId}\0${appId}`; - const isCurrent = (entry: DesktopEntry) => - entries.get(entry.environmentId) === entry && !entry.stopped; - - const closeObserver = (observer: ObserverEntry, code: number, reason: string) => { - try { - observer.close(code, reason); - } catch { - // Observer cleanup remains authoritative when the transport close callback fails. - } - }; - - const stopEntry = (entry: DesktopEntry): Promise => { - if (entry.stopPromise) { - return entry.stopPromise; - } - entry.stopPromise = (async () => { - entry.stopped = true; - if (entries.get(entry.environmentId) === entry) { - entries.delete(entry.environmentId); - } - clearTimeout(entry.lingerTimer); - entry.lingerTimer = undefined; - for (const observer of entry.observers) { - observer.released = true; - closeObserver(observer, 1012, "desktop tunnel closed"); - } - entry.observers.clear(); - entry.controller = undefined; - if (!entry.readySettled) { - entry.readySettled = true; - entry.rejectReady(new Error("Worker desktop tunnel stopped before connecting")); - } - const processBeforeInitialization = entry.process; - await processBeforeInitialization?.stop().catch(() => undefined); - await entry.initialization?.catch(() => undefined); - if (entry.process !== processBeforeInitialization) { - await entry.process?.stop().catch(() => undefined); - } - await entry.prepared?.dispose().catch(() => undefined); - })(); - return entry.stopPromise; - }; - const stopAppLaunches = async (environmentId: string, ownerEpoch?: number): Promise => { const matching = [...appLaunches.values()].filter( (entry) => @@ -159,23 +94,19 @@ export function createWorkerDesktopTunnels(deps: { await Promise.allSettled(matching.map((entry) => entry.operation)); }; - const reserveOwnerEpoch = (environmentId: string, ownerEpoch: number): boolean => { - const claimedEpoch = claimedOwnerEpochs.get(environmentId); - if (claimedEpoch !== undefined && ownerEpoch < claimedEpoch) { - throw new Error("Worker desktop owner epoch is stale"); + const claimOwnerEpoch = (environmentId: string, ownerEpoch: number): boolean => { + try { + return sessions.claimOwnerEpoch(environmentId, ownerEpoch); + } catch (error) { + if (error instanceof DesktopSessionStaleOwnerError) { + throw new Error("Worker desktop owner epoch is stale", { cause: error }); + } + throw error; } - if (claimedEpoch === undefined || ownerEpoch > claimedEpoch) { - claimedOwnerEpochs.set(environmentId, ownerEpoch); - return true; - } - return false; }; const fenceReplacedOwners = async (environmentId: string, ownerEpoch: number): Promise => { - const current = entries.get(environmentId); - if (current && current.ownerEpoch < ownerEpoch) { - await stopEntry(current); - } + await sessions.stopSuperseded(environmentId, ownerEpoch); const staleLaunches = [...appLaunches.values()].filter( (entry) => entry.environmentId === environmentId && entry.ownerEpoch < ownerEpoch, ); @@ -185,139 +116,144 @@ export function createWorkerDesktopTunnels(deps: { await Promise.allSettled(staleLaunches.map((entry) => entry.operation)); }; - const startEntry = async (entry: DesktopEntry, request: DesktopAcquireRequest) => { - const prepared = await prepareWorkerSsh({ - ssh: request.ssh, - pinnedHostKey: request.ssh.hostKey, - resolveIdentity: request.resolveIdentity, - temporaryDirectoryPrefix: "openclaw-worker-desktop-", - }); - entry.prepared = prepared; - if (!isCurrent(entry)) { - await prepared.dispose(); - entry.prepared = undefined; - return; - } - const localSocketPath = path.join(path.dirname(prepared.knownHostsPath), "desktop.sock"); - entry.localSocketPath = localSocketPath; - const child = deps.runner.start( - [ - "ssh", - ...workerSshOptions(prepared, { forwarding: "explicit" }), - "-a", - "-x", - "-T", - "-o", - "ServerAliveInterval=15", - "-o", - "ServerAliveCountMax=3", - "-o", - "StreamLocalBindMask=0177", - "-L", - `${localSocketPath}:127.0.0.1:${request.desktop.port}`, - "-p", - String(prepared.port), - "--", - prepared.sshTarget, - workerSshRemoteCommand(["sh", "-s"]), - ], - workerSshCommandOptions({ - input: REMOTE_DESKTOP_READY_SCRIPT, - timeoutMs: Number.MAX_SAFE_INTEGER, - }), - ); - entry.process = child; - void child.exited.then(() => { - if (isCurrent(entry)) { - void stopEntry(entry); + const createSessionHooks = (request: DesktopAcquireRequest) => { + let prepared: PreparedWorkerSsh | undefined; + let child: WorkerSshProcess | undefined; + let stoppedChild: WorkerSshProcess | undefined; + let startSettled = false; + + const start = async (isCurrent: () => boolean): Promise => { + try { + prepared = await prepareWorkerSsh({ + ssh: request.ssh, + pinnedHostKey: request.ssh.hostKey, + resolveIdentity: request.resolveIdentity, + temporaryDirectoryPrefix: "openclaw-worker-desktop-", + }); + if (!isCurrent()) { + await prepared.dispose(); + prepared = undefined; + throw new Error("Worker desktop tunnel stopped before connecting"); + } + const localSocketPath = path.join(path.dirname(prepared.knownHostsPath), "desktop.sock"); + child = deps.runner.start( + [ + "ssh", + ...workerSshOptions(prepared, { forwarding: "explicit" }), + "-a", + "-x", + "-T", + "-o", + "ServerAliveInterval=15", + "-o", + "ServerAliveCountMax=3", + "-o", + "StreamLocalBindMask=0177", + "-L", + `${localSocketPath}:127.0.0.1:${request.desktop.port}`, + "-p", + String(prepared.port), + "--", + prepared.sshTarget, + workerSshRemoteCommand(["sh", "-s"]), + ], + workerSshCommandOptions({ + input: REMOTE_DESKTOP_READY_SCRIPT, + timeoutMs: Number.MAX_SAFE_INTEGER, + }), + ); + const startedChild = child; + void startedChild.exited.then(() => { + if (isCurrent()) { + void sessions.stop(request.environmentId, request.ownerEpoch); + } + }); + await startedChild.ready; + if (!isCurrent()) { + await startedChild.stop(); + throw new Error("Worker desktop tunnel stopped before connecting"); + } + let vncPassword: string | undefined; + if (request.desktop.passwordFilePath) { + const result = await deps.runner.run( + [ + "ssh", + ...workerSshOptions(prepared, { forwarding: "disabled" }), + "-a", + "-x", + "-T", + "-p", + String(prepared.port), + "--", + prepared.sshTarget, + workerSshRemoteCommand(["cat", request.desktop.passwordFilePath]), + ], + workerSshCommandOptions({ timeoutMs: PASSWORD_READ_TIMEOUT_MS }), + ); + if (!successful(result)) { + throw workerSshProcessError(result.stderr || result.stdout); + } + vncPassword = result.stdout.replace(/(?:\r?\n)+$/u, ""); + if (!vncPassword) { + throw new Error("Worker desktop password file is empty"); + } + registerSecretValueForRedaction(vncPassword); + } + return { + attachment: { kind: "unix-socket", socketPath: localSocketPath }, + ...(vncPassword ? { vncPassword } : {}), + }; + } finally { + startSettled = true; } - }); - await child.ready; - if (!isCurrent(entry)) { - await child.stop(); - return; - } - let vncPassword: string | undefined; - if (request.desktop.passwordFilePath) { - const result = await deps.runner.run( - [ - "ssh", - ...workerSshOptions(prepared, { forwarding: "disabled" }), - "-a", - "-x", - "-T", - "-p", - String(prepared.port), - "--", - prepared.sshTarget, - workerSshRemoteCommand(["cat", request.desktop.passwordFilePath]), - ], - workerSshCommandOptions({ timeoutMs: PASSWORD_READ_TIMEOUT_MS }), - ); - if (!successful(result)) { - throw workerSshProcessError(result.stderr || result.stdout); + }; + + const teardown = async (): Promise => { + if (child && child !== stoppedChild) { + stoppedChild = child; + await child.stop().catch(() => undefined); } - vncPassword = result.stdout.replace(/(?:\r?\n)+$/u, ""); - if (!vncPassword) { - throw new Error("Worker desktop password file is empty"); + if (!startSettled) { + return; } - registerSecretValueForRedaction(vncPassword); - } - entry.readySettled = true; - entry.resolveReady({ localSocketPath, ...(vncPassword ? { vncPassword } : {}) }); + if (child && child !== stoppedChild) { + stoppedChild = child; + await child?.stop().catch(() => undefined); + } + await prepared?.dispose().catch(() => undefined); + prepared = undefined; + }; + + return { start, teardown }; }; async function acquire(request: DesktopAcquireRequest): Promise { if (platform === "win32") { throw new WorkerDesktopUnsupportedError(); } - const ownerAdvanced = reserveOwnerEpoch(request.environmentId, request.ownerEpoch); + const ownerAdvanced = claimOwnerEpoch(request.environmentId, request.ownerEpoch); if (ownerAdvanced) { await fenceReplacedOwners(request.environmentId, request.ownerEpoch); } - if (claimedOwnerEpochs.get(request.environmentId) !== request.ownerEpoch) { + if (!sessions.isOwnerEpochCurrent(request.environmentId, request.ownerEpoch)) { throw new Error("Worker desktop owner epoch is stale"); } - const current = entries.get(request.environmentId); - if (current?.ownerEpoch === request.ownerEpoch) { - return await current.ready; + const hooks = createSessionHooks(request); + try { + return await sessions.acquire({ + sourceKey: request.environmentId, + ownerEpoch: request.ownerEpoch, + ...hooks, + }); + } catch (error) { + if (error instanceof DesktopSessionStaleOwnerError) { + throw new Error("Worker desktop owner epoch is stale", { cause: error }); + } + if (error instanceof DesktopSessionStoppedError) { + throw new Error("Worker desktop tunnel stopped before connecting", { cause: error }); + } + throw error; } - let resolveReady!: (result: DesktopAcquireResult) => void; - let rejectReady!: (error: Error) => void; - const ready = new Promise((resolve, reject) => { - resolveReady = resolve; - rejectReady = reject; - }); - void ready.catch(() => undefined); - const entry: DesktopEntry = { - environmentId: request.environmentId, - ownerEpoch: request.ownerEpoch, - ready, - resolveReady, - rejectReady, - readySettled: false, - observers: new Set(), - stopped: false, - }; - entries.set(request.environmentId, entry); - entry.initialization = (async () => { - if (current) { - await stopEntry(current); - } - if (isCurrent(entry)) { - await startEntry(entry, request); - } - })(); - void entry.initialization.catch((error: unknown) => { - if (!entry.readySettled) { - entry.readySettled = true; - entry.rejectReady( - error instanceof Error ? error : new Error("Worker desktop tunnel failed"), - ); - } - void stopEntry(entry); - }); - return await ready; } function launchApp(request: { @@ -332,7 +268,7 @@ export function createWorkerDesktopTunnels(deps: { } let ownerAdvanced: boolean; try { - ownerAdvanced = reserveOwnerEpoch(request.environmentId, request.ownerEpoch); + ownerAdvanced = claimOwnerEpoch(request.environmentId, request.ownerEpoch); } catch (error) { return Promise.reject( error instanceof Error @@ -354,7 +290,7 @@ export function createWorkerDesktopTunnels(deps: { const execution = (async () => { await startGate; abortController.signal.throwIfAborted(); - if (claimedOwnerEpochs.get(request.environmentId) !== request.ownerEpoch) { + if (!sessions.isOwnerEpochCurrent(request.environmentId, request.ownerEpoch)) { throw new Error("Worker desktop app launch owner was replaced"); } if (current) { @@ -430,54 +366,9 @@ export function createWorkerDesktopTunnels(deps: { return operation; } - function attachObserver(environmentId: string, observer: WorkerDesktopObserver) { - const entry = entries.get(environmentId); - if (!entry || !entry.readySettled || entry.stopped || entry.observers.size >= MAX_OBSERVERS) { - return undefined; - } - // A token minted against a replaced entry must not reach this one; otherwise a stale - // control token would evict the current controller of a desktop it never observed. - if (observer.ownerEpoch !== entry.ownerEpoch) { - return undefined; - } - clearTimeout(entry.lingerTimer); - entry.lingerTimer = undefined; - if (observer.control && entry.controller) { - const previous = entry.controller; - previous.released = true; - entry.observers.delete(previous); - entry.controller = undefined; - closeObserver(previous, 4000, "control-taken"); - } - const attached: ObserverEntry = { ...observer, released: false }; - entry.observers.add(attached); - if (attached.control) { - entry.controller = attached; - } - return { - release() { - if (attached.released) { - return; - } - attached.released = true; - entry.observers.delete(attached); - if (entry.controller === attached) { - entry.controller = undefined; - } - if (entry.observers.size === 0 && isCurrent(entry)) { - entry.lingerTimer = setTimeout(() => void stopEntry(entry), lingerMs); - entry.lingerTimer.unref?.(); - } - }, - }; - } - async function stop(environmentId: string, ownerEpoch?: number): Promise { - const entry = entries.get(environmentId); await Promise.all([ - entry && (ownerEpoch === undefined || ownerEpoch === entry.ownerEpoch) - ? stopEntry(entry) - : Promise.resolve(), + sessions.stop(environmentId, ownerEpoch), stopAppLaunches(environmentId, ownerEpoch), ]); } @@ -487,12 +378,16 @@ export function createWorkerDesktopTunnels(deps: { entry.abortController.abort(new Error("Worker desktop app launcher stopped")); } await Promise.all([ - ...[...entries.values()].map(stopEntry), + sessions.stopAll(), ...[...appLaunches.values()].map((entry) => entry.operation.catch(() => undefined)), ]); } - return { acquire, attachObserver, launchApp, stop, stopAll }; + return { + acquire, + attachObserver: sessions.attachObserver, + launchApp, + stop, + stopAll, + }; } - -export type WorkerDesktopTunnels = ReturnType; diff --git a/src/gateway/worker-environments/environment-access.test.ts b/src/gateway/worker-environments/environment-access.test.ts index 505380be7216..ad2ffedc0439 100644 --- a/src/gateway/worker-environments/environment-access.test.ts +++ b/src/gateway/worker-environments/environment-access.test.ts @@ -241,7 +241,7 @@ describe("worker environment service", () => { const record = support.seedReadyDesktop("worker-desktop-observe"); const desktopPassword = ["desktop", String.fromCharCode(45), "secret"].join(""); const acquire = vi.fn(async () => ({ - localSocketPath: "/tmp/worker-desktop.sock", + attachment: { kind: "unix-socket" as const, socketPath: "/tmp/worker-desktop.sock" }, vncPassword: desktopPassword, })); const tunnelManager = { @@ -262,7 +262,7 @@ describe("worker environment service", () => { workerService.observeDesktop({ environmentId: record.environmentId, control: true }), ).resolves.toMatchObject({ transport: "rfb", - wsPath: expect.stringMatching(/^\/worker-desktop\/observe\?token=[a-f0-9]{48}$/u), + wsPath: expect.stringMatching(/^\/desktop\/observe\?token=[a-f0-9]{48}$/u), expiresAtMs: support.testState.nowMs + 60_000, control: true, vncPassword: desktopPassword, diff --git a/src/gateway/worker-environments/environment-access.ts b/src/gateway/worker-environments/environment-access.ts index 5f4a8febf4f2..e317b3b68f5f 100644 --- a/src/gateway/worker-environments/environment-access.ts +++ b/src/gateway/worker-environments/environment-access.ts @@ -163,7 +163,7 @@ export function createWorkerEnvironmentAccess(options: WorkerEnvironmentAccessOp if (!tunnels) { throw serviceError("invalid_state", "Worker tunnel runtime is unavailable"); } - let startup: Promise<{ localSocketPath: string; vncPassword?: string }> | undefined; + let startup: ReturnType | undefined; let ownerEpoch: number | undefined; await withLock(request.environmentId, async () => { stopping = options.isStopping(); @@ -203,18 +203,18 @@ export function createWorkerEnvironmentAccess(options: WorkerEnvironmentAccessOp throw serviceError("invalid_state", "Worker desktop tunnel failed to start"); } const acquired = await startup; - const { WORKER_DESKTOP_OBSERVE_PATH, mintWorkerDesktopObserverToken } = - await import("./desktop-observe.js"); - const minted = mintWorkerDesktopObserverToken({ - environmentId: request.environmentId, + const { DESKTOP_OBSERVE_PATH, mintDesktopObserverToken } = + await import("../desktop/observe-bridge.js"); + const minted = mintDesktopObserverToken({ + sourceKey: request.environmentId, ownerEpoch, control: request.control, - localSocketPath: acquired.localSocketPath, + attachment: acquired.attachment, nowMs: now(), }); return { transport: "rfb", - wsPath: `${WORKER_DESKTOP_OBSERVE_PATH}?token=${minted.token}`, + wsPath: `${DESKTOP_OBSERVE_PATH}?token=${minted.token}`, expiresAtMs: minted.expiresAtMs, control: request.control, ...(acquired.vncPassword ? { vncPassword: acquired.vncPassword } : {}), diff --git a/src/gateway/worker-environments/tunnel.ts b/src/gateway/worker-environments/tunnel.ts index f6f76acf2a45..55ded6840b43 100644 --- a/src/gateway/worker-environments/tunnel.ts +++ b/src/gateway/worker-environments/tunnel.ts @@ -5,6 +5,7 @@ import { createSubsystemLogger } from "../../logging/subsystem.js"; import type { WorkerSshEndpoint } from "../../plugins/types.js"; import type { SpawnResult } from "../../process/exec.js"; import { createDeferredCore, type Deferred } from "../../shared/deferred.js"; +import type { DesktopSessionRegistry } from "../desktop/session-registry.js"; import { createWorkerDesktopTunnels } from "./desktop-tunnel.js"; import { advanceWorkerSshAfterTransportExit, @@ -108,6 +109,7 @@ type TunnelEntry = { type WorkerTunnelManagerOptions = { runner?: WorkerSshRunner; + desktopSessionRegistry?: DesktopSessionRegistry; sleep?: (ms: number, signal?: AbortSignal) => Promise; backoff?: BackoffPolicy; now?: () => number; @@ -145,7 +147,10 @@ export function createWorkerTunnelManager(options: WorkerTunnelManagerOptions = const backoff = options.backoff ?? DEFAULT_BACKOFF; const now = options.now ?? Date.now; const stableConnectionMs = options.stableConnectionMs ?? DEFAULT_STABLE_CONNECTION_MS; - const desktop = createWorkerDesktopTunnels({ runner, now }); + const desktop = createWorkerDesktopTunnels({ + runner, + ...(options.desktopSessionRegistry ? { registry: options.desktopSessionRegistry } : {}), + }); const entries = new Map(); const claimedOwnerEpochs = new Map(); diff --git a/ui/src/components/desktop/desktop-client.test.ts b/ui/src/components/desktop/desktop-client.test.ts index 41e53595af51..535dac9a444b 100644 --- a/ui/src/components/desktop/desktop-client.test.ts +++ b/ui/src/components/desktop/desktop-client.test.ts @@ -37,14 +37,8 @@ function createFakeRfb() { describe("DesktopClient", () => { it.each([ - [ - "http://control.example.test/chat", - "ws://control.example.test/worker-desktop/observe?token=abc", - ], - [ - "https://control.example.test/chat", - "wss://control.example.test/worker-desktop/observe?token=abc", - ], + ["http://control.example.test/chat", "ws://control.example.test/desktop/observe?token=abc"], + ["https://control.example.test/chat", "wss://control.example.test/desktop/observe?token=abc"], ])("resolves relative observer URLs against %s", async (gatewayUrl, expectedUrl) => { const { Rfb, instances } = createFakeRfb(); const sockets: FakeSocket[] = []; @@ -57,7 +51,7 @@ describe("DesktopClient", () => { await client.connect({ gatewayUrl, - wsUrl: "/worker-desktop/observe?token=abc", + wsUrl: "/desktop/observe?token=abc", password: "secret", viewOnly: true, target, @@ -70,12 +64,12 @@ describe("DesktopClient", () => { it("propagates RFB options and disconnects through the returned handle", async () => { const { Rfb, instances } = createFakeRfb(); - const socket = new FakeSocket("ws://control.example.test/worker-desktop/observe"); + const socket = new FakeSocket("ws://control.example.test/desktop/observe"); const client = new DesktopClient(Rfb, () => socket as unknown as WebSocket); const handle = await client.connect({ gatewayUrl: "ws://control.example.test", - wsUrl: "/worker-desktop/observe", + wsUrl: "/desktop/observe", password: "secret", background: "rgb(8, 8, 8)", viewOnly: false, @@ -93,12 +87,12 @@ describe("DesktopClient", () => { it("forwards socket close metadata through the RFB disconnect callback", async () => { const { Rfb, instances } = createFakeRfb(); - const socket = new FakeSocket("ws://control.example.test/worker-desktop/observe"); + const socket = new FakeSocket("ws://control.example.test/desktop/observe"); const onDisconnect = vi.fn(); const client = new DesktopClient(Rfb, () => socket as unknown as WebSocket); await client.connect({ - wsUrl: "ws://control.example.test/worker-desktop/observe", + wsUrl: "ws://control.example.test/desktop/observe", viewOnly: true, target: document.createElement("div"), onDisconnect, diff --git a/ui/src/e2e/desktop-panel.e2e.test.ts b/ui/src/e2e/desktop-panel.e2e.test.ts index 2a286501b253..d23c872b5fb4 100644 --- a/ui/src/e2e/desktop-panel.e2e.test.ts +++ b/ui/src/e2e/desktop-panel.e2e.test.ts @@ -141,7 +141,7 @@ suite.define(() => { match: { environmentId: "worker-desktop-1", control: false }, response: { transport: "rfb", - wsPath: "/worker-desktop/observe?token=view", + wsPath: "/desktop/observe?token=view", expiresAtMs: 60_000, control: false, }, @@ -150,7 +150,7 @@ suite.define(() => { match: { environmentId: "worker-desktop-1", control: true }, response: { transport: "rfb", - wsPath: "/worker-desktop/observe?token=control", + wsPath: "/desktop/observe?token=control", expiresAtMs: 60_000, control: true, }, @@ -274,7 +274,7 @@ suite.define(() => { }, "worker.desktop.observe": { transport: "rfb", - wsPath: "/worker-desktop/observe?token=view", + wsPath: "/desktop/observe?token=view", expiresAtMs: 60_000, control: false, },