diff --git a/mcp/packages/server/src/PluginBridge.test.ts b/mcp/packages/server/src/PenpotConnection.test.ts similarity index 98% rename from mcp/packages/server/src/PluginBridge.test.ts rename to mcp/packages/server/src/PenpotConnection.test.ts index b0d15936be..28f2464f88 100644 --- a/mcp/packages/server/src/PluginBridge.test.ts +++ b/mcp/packages/server/src/PenpotConnection.test.ts @@ -1,6 +1,6 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { assertPluginResponsive, HEARTBEAT_STALE_THRESHOLD_MS } from "./PluginBridge"; +import { assertPluginResponsive, HEARTBEAT_STALE_THRESHOLD_MS } from "./PenpotConnection"; test("passes for a responsive connection with a recent heartbeat", () => { const now = 1_000_000; diff --git a/mcp/packages/server/src/PenpotConnection.ts b/mcp/packages/server/src/PenpotConnection.ts index 267c1b159b..8635f30f14 100644 --- a/mcp/packages/server/src/PenpotConnection.ts +++ b/mcp/packages/server/src/PenpotConnection.ts @@ -11,6 +11,42 @@ export interface PluginLivenessState { frozen: boolean; } +/** + * Maximum plugin heartbeat age before a connection is stale. + * + * This uses plugin heartbeats rather than WebSocket pongs because the browser can answer + * protocol pings while the tab's JavaScript event loop is frozen. + */ +export const HEARTBEAT_STALE_THRESHOLD_MS = 30000; + +/** + * Throws if the plugin tab cannot currently run tasks. + * + * A socket can stay open while the page event loop is paused, so task dispatch must check + * plugin-level liveness before sending work. + */ +export function assertPluginResponsive( + state: PluginLivenessState, + now: number, + staleThresholdMs: number = HEARTBEAT_STALE_THRESHOLD_MS +): void { + if (state.frozen) { + throw new Error( + `The Penpot plugin tab has been frozen by the browser and cannot run tasks. ` + + `Please click/focus the Penpot tab to wake it, then retry.` + ); + } + + const heartbeatAge = now - state.lastHeartbeat; + if (heartbeatAge > staleThresholdMs) { + throw new Error( + `The Penpot plugin tab appears to be suspended by the browser (no heartbeat for ` + + `${Math.round(heartbeatAge / 1000)}s). Please click/focus the Penpot tab to wake it, ` + + `then retry.` + ); + } +} + /** One local Penpot WebSocket connection and its session metadata. */ export interface PenpotConnection extends PluginLivenessState { readonly socket: WebSocket; diff --git a/mcp/packages/server/src/PluginBridge.sessions.test.ts b/mcp/packages/server/src/PluginBridge.sessions.test.ts index 1046fdb062..20874c60fc 100644 --- a/mcp/packages/server/src/PluginBridge.sessions.test.ts +++ b/mcp/packages/server/src/PluginBridge.sessions.test.ts @@ -37,14 +37,17 @@ async function connect(sessionId: string, token = "alice", fileName = "Design", const socket = new WebSocket(`ws://127.0.0.1:${port}?userToken=${token}`); sockets.push(socket); await once(socket, "open"); + // answer tasks from the start, since one may arrive right behind the initialization + socket.on("message", (raw) => { + const request = JSON.parse(raw.toString()); + if (request.type !== "initialized") { + socket.send(JSON.stringify({ id: request.id, success: true, data: reply })); + } + }); const initialized = once(socket, "message", { signal: AbortSignal.timeout(1500) }); socket.send(JSON.stringify({ type: "initialize", session: { sessionId, fileId: "file-1", fileName } })); const [message] = await initialized; assert.equal(JSON.parse(message.toString()).type, "initialized"); - socket.on("message", (raw) => { - const request = JSON.parse(raw.toString()); - socket.send(JSON.stringify({ id: request.id, success: true, data: reply })); - }); return socket; } @@ -234,3 +237,141 @@ test("a wake agent connection without a token is rejected in multi-user mode", a const [code] = await once(socket, "close"); assert.equal(code, 1008); }); + +test("a plugin message resolves a liveness wait started before it", async () => { + const plugin = await connect("tab"); + const connection = bridge.getUserConnections("alice")!.getRegistered("tab")!; + plugin.send(JSON.stringify({ type: "freeze" })); + while (!connection.frozen) { + await new Promise((resolve) => setImmediate(resolve)); + } + const { signal } = bridge.liveness.wait(connection, 1000); + plugin.send(JSON.stringify({ type: "heartbeat" })); + await signal; + assert.equal(connection.frozen, false); +}); + +test("a plugin reconnecting for the same session resolves a liveness wait with its new connection", async () => { + const plugin = await connect("tab"); + const connection = bridge.getUserConnections("alice")!.getRegistered("tab")!; + const { signal } = bridge.liveness.wait(connection, 1000); + const closed = once(plugin, "close"); + plugin.close(); + await closed; + await connect("tab"); + const replacement = await signal; + assert.notEqual(replacement, connection); + assert.equal(replacement, bridge.getUserConnections("alice")!.get("tab")); +}); + +/** Reports a browser freeze from a plugin socket and waits until the bridge has recorded it. */ +async function freeze(plugin: WebSocket, sessionId = "tab") { + const connection = bridge.getUserConnections("alice")!.getRegistered(sessionId)!; + plugin.send(JSON.stringify({ type: "freeze" })); + while (!connection.frozen) { + await new Promise((resolve) => setImmediate(resolve)); + } +} + +/** Answers wake requests like a browser extension whose activation resumes the plugin. */ +function wakeOnRequest(agent: WebSocket, plugin: WebSocket) { + agent.on("message", (raw) => { + const request = JSON.parse(raw.toString()); + agent.send(JSON.stringify({ type: "wake-result", requestId: request.requestId, outcome: "activated" })); + plugin.send(JSON.stringify({ type: "heartbeat" })); + }); +} + +test("a task for a frozen session runs once a wake agent has woken it", async () => { + const plugin = await connect("tab"); + const agent = await connectAgent(["tab"]); + await freeze(plugin); + wakeOnRequest(agent, plugin); + assert.equal((await bridge.executePluginTask(new PluginTask("test", {}), "tab")).data, "tab"); +}); + +test("a task for a frozen session goes to the new connection when the woken page reconnects", async () => { + const plugin = await connect("tab", "alice", "Design", "previous"); + const agent = await connectAgent(["tab"]); + await freeze(plugin); + // the resumed page drops its socket and opens a new one for the same session, as Chrome does + agent.on("message", (raw) => { + const request = JSON.parse(raw.toString()); + agent.send(JSON.stringify({ type: "wake-result", requestId: request.requestId, outcome: "activated" })); + plugin.close(); + void connect("tab", "alice", "Design", "current"); + }); + assert.equal((await bridge.executePluginTask(new PluginTask("test", {}), "tab")).data, "current"); +}); + +test("a task for a frozen session waits for the new connection when the woken page answers, then closes", async () => { + const plugin = await connect("tab", "alice", "Design", "previous"); + const agent = await connectAgent(["tab"]); + await freeze(plugin); + // Chrome runs the page's queued heartbeat, then closes its socket with 1001 + agent.on("message", (raw) => { + const request = JSON.parse(raw.toString()); + agent.send(JSON.stringify({ type: "wake-result", requestId: request.requestId, outcome: "activated" })); + plugin.send(JSON.stringify({ type: "heartbeat" })); + plugin.close(1001); + setTimeout(() => void connect("tab", "alice", "Design", "current"), 50); + }); + assert.equal((await bridge.executePluginTask(new PluginTask("test", {}), "tab")).data, "current"); +}); + +test("a task for a frozen session fails as before when the agent cannot wake it", async () => { + const plugin = await connect("tab"); + const agent = await connectAgent(["tab"]); + await freeze(plugin); + answerWake(agent, "tab-closed"); + await assert.rejects(bridge.executePluginTask(new PluginTask("test", {}), "tab"), /has been frozen/); +}); + +test("a task for a frozen session fails as before when no wake agent is connected", async () => { + const plugin = await connect("tab"); + await freeze(plugin); + await assert.rejects(bridge.executePluginTask(new PluginTask("test", {}), "tab"), /has been frozen/); +}); + +test("waking a session never replays a task that timed out", async () => { + const plugin = new WebSocket(`ws://127.0.0.1:${port}?userToken=alice`); + sockets.push(plugin); + await once(plugin, "open"); + const initialized = once(plugin, "message"); + plugin.send(JSON.stringify({ type: "initialize", session: { sessionId: "tab", fileId: null, fileName: null } })); + await initialized; + // the plugin stops answering after the first request, as a tab frozen mid-task would + const requests: string[] = []; + plugin.on("message", (raw) => { + const request = JSON.parse(raw.toString()); + requests.push(request.id); + if (requests.length > 1) { + plugin.send(JSON.stringify({ id: request.id, success: true, data: "woken" })); + } + }); + const lost = new PluginTask("test", {}); + await assert.rejects(bridge.executePluginTask(lost, "tab"), /timed out/); + await freeze(plugin); + wakeOnRequest(await connectAgent(["tab"]), plugin); + const task = new PluginTask("test", {}); + assert.equal((await bridge.executePluginTask(task, "tab")).data, "woken"); + assert.deepEqual(requests, [lost.id, task.id]); +}); + +test("a task fails at once when its plugin connection closes before answering", async () => { + const plugin = new WebSocket(`ws://127.0.0.1:${port}?userToken=alice`); + sockets.push(plugin); + await once(plugin, "open"); + const initialized = once(plugin, "message"); + plugin.send(JSON.stringify({ type: "initialize", session: { sessionId: "tab", fileId: null, fileName: null } })); + await initialized; + // the socket closes once the task reaches it, as an old socket does right after a page resumes + plugin.on("message", () => plugin.terminate()); + const started = Date.now(); + await assert.rejects( + bridge.executePluginTask(new PluginTask("test", {}), "tab"), + /connection closed before the task finished/ + ); + // well within the 1s task timeout of these tests + assert.ok(Date.now() - started < 500); +}); diff --git a/mcp/packages/server/src/PluginBridge.ts b/mcp/packages/server/src/PluginBridge.ts index b4b8e084b3..4c6125ebf0 100644 --- a/mcp/packages/server/src/PluginBridge.ts +++ b/mcp/packages/server/src/PluginBridge.ts @@ -6,25 +6,19 @@ import { PluginTaskResponse, PluginTaskResult } from "@penpot/mcp-common"; import { createLogger } from "./logger"; import type { PenpotMcpServer } from "./PenpotMcpServer"; import type { RedisBridge } from "./RedisBridge"; -import type { PenpotConnection, PluginLivenessState } from "./PenpotConnection"; +import type { PenpotConnection } from "./PenpotConnection"; import { UserPenpotConnections } from "./UserPenpotConnections"; import type { TaskDispatcher, TaskDispatchHost } from "./dispatch/TaskDispatcher"; import { SingleInstanceTaskDispatcher } from "./dispatch/SingleInstanceTaskDispatcher"; import { MultiInstanceTaskDispatcher } from "./dispatch/MultiInstanceTaskDispatcher"; import { WakeAgentRegistry } from "./wake/WakeAgentRegistry"; import { wakeAgentInitSchema } from "./wake/WakeAgentProtocol"; +import { LivenessWaiter } from "./wake/LivenessWaiter"; +import { LocalSessionWaker, type SessionWaker } from "./wake/SessionWaker"; const KEEP_ALIVE_TIME = 30000; // 30 seconds const INITIALIZATION_TIMEOUT_MS = 10_000; -/** - * Maximum plugin heartbeat age before a connection is stale. - * - * This uses plugin heartbeats rather than WebSocket pongs because the browser can answer - * protocol pings while the tab's JavaScript event loop is frozen. - */ -export const HEARTBEAT_STALE_THRESHOLD_MS = 30000; - const connectionInitSchema = z.object({ type: z.literal("initialize"), session: z.object({ @@ -35,34 +29,6 @@ const connectionInitSchema = z.object({ }), }); -/** - * Throws if the plugin tab cannot currently run tasks. - * - * A socket can stay open while the page event loop is paused, so task dispatch must check - * plugin-level liveness before sending work. - */ -export function assertPluginResponsive( - state: PluginLivenessState, - now: number, - staleThresholdMs: number = HEARTBEAT_STALE_THRESHOLD_MS -): void { - if (state.frozen) { - throw new Error( - `The Penpot plugin tab has been frozen by the browser and cannot run tasks. ` + - `Please click/focus the Penpot tab to wake it, then retry.` - ); - } - - const heartbeatAge = now - state.lastHeartbeat; - if (heartbeatAge > staleThresholdMs) { - throw new Error( - `The Penpot plugin tab appears to be suspended by the browser (no heartbeat for ` + - `${Math.round(heartbeatAge / 1000)}s). Please click/focus the Penpot tab to wake it, ` + - `then retry.` - ); - } -} - /** * Manages WebSocket connections to Penpot plugin instances and handles plugin tasks * over these connections. @@ -82,6 +48,9 @@ export class PluginBridge implements TaskDispatchHost { private readonly dispatcher: TaskDispatcher; /** the browser extensions that can wake the tabs of frozen plugin sessions. */ public readonly wakeAgents = new WakeAgentRegistry(KEEP_ALIVE_TIME); + /** the waits for fresh liveness signals from the plugins of sessions. */ + public readonly liveness = new LivenessWaiter(); + private readonly sessionWaker: SessionWaker = new LocalSessionWaker(this.wakeAgents, this.liveness, this); /** * Creates the plugin bridge and starts its WebSocket server. @@ -190,6 +159,7 @@ export class PluginBridge implements TaskDispatchHost { return; } connection.frozen = false; + this.liveness.notify(connection); if (message?.type === "context-update") { const context = connectionInitSchema.shape.session .omit({ sessionId: true }) @@ -206,8 +176,8 @@ export class PluginBridge implements TaskDispatchHost { } }); - ws.on("close", () => { - this.logger.info("WebSocket connection closed"); + ws.on("close", (code: number, reason: Buffer) => { + this.logger.info("WebSocket connection closed (code %d%s)", code, reason.length ? `: ${reason}` : ""); clearTimeout(initializationTimeout); this.wakeAgents.remove(ws); this.removeConnection(ws); @@ -261,6 +231,8 @@ export class PluginBridge implements TaskDispatchHost { if (ws.readyState === WebSocket.OPEN && this.connectionsBySocket.get(ws) === connection) { connection.ready = true; ws.send(JSON.stringify({ type: "initialized" })); + // a page resuming from a browser freeze may replace its connection + this.liveness.notify(connection); } } @@ -283,8 +255,9 @@ export class PluginBridge implements TaskDispatchHost { * * Clears the per-connection keep-alive interval and removes the connection from the * socket-keyed index and the user's session collection. The token-keyed index entry - * is removed only when the user's last local connection closes. Safe to call with a - * socket that is not (or no longer) registered. + * is removed only when the user's last local connection closes. Tasks sent over the + * connection and still awaiting a response fail at once. Safe to call with a socket + * that is not (or no longer) registered. * * @param ws - The WebSocket whose connection state should be removed */ @@ -298,6 +271,13 @@ export class PluginBridge implements TaskDispatchHost { this.connectionsBySocket.delete(ws); const connections = this.getUserConnections(connection.userToken)!; connections.remove(connection); + // a task sent over the closed socket can no longer be answered, since responses only count + // from the connection the task went to + for (const [taskId, taskConnection] of [...this.taskConnections]) { + if (taskConnection === connection) { + this.rejectPendingTask(taskId, new Error("The plugin connection closed before the task finished.")); + } + } if (connection.userToken !== null && connections.size === 0) { this.userPenpotConnectionsByToken.delete(connection.userToken); } @@ -454,16 +434,20 @@ export class PluginBridge implements TaskDispatchHost { this.startTaskTimeout(task, () => void redisBridge.unsubscribeFromResponse(task.id)); } - /** Sends a task directly to a ready local connection, with execution tracking. */ - public sendLocalTask(task: AbstractPluginTask, connection: PenpotConnection): void { - const target = connection; - if (!target.ready || target.socket.readyState !== 1) { - // WebSocket is not open - throw new Error(`Plugin instance is disconnected. Task could not be sent.`); - } + /** + * Sends a task directly to a ready local connection, with execution tracking. + * + * A plugin whose tab the browser has frozen is woken first, if a wake agent can do so; the + * task then goes to the session's new connection if the woken page replaced its WebSocket. + * The task's timeout starts once the task is sent. + */ + public async sendLocalTask(task: AbstractPluginTask, connection: PenpotConnection): Promise { + this.assertConnected(connection); // the socket can be open while browser-throttled plugin JS cannot run tasks - assertPluginResponsive(target, Date.now()); + const target = await this.sessionWaker.ensureResponsive(connection); + // the connection may have closed after its liveness signal + this.assertConnected(target); // register the task for result correlation, then send over the socket this.pendingTasks.set(task.id, task); @@ -472,6 +456,13 @@ export class PluginBridge implements TaskDispatchHost { this.startTaskTimeout(task); } + /** Throws if the connection cannot receive tasks because its socket is not open. */ + private assertConnected(connection: PenpotConnection): void { + if (!connection.ready || connection.socket.readyState !== WebSocket.OPEN) { + throw new Error(`Plugin instance is disconnected. Task could not be sent.`); + } + } + /** Arms the shared execution timeout for a locally or remotely dispatched task. */ private startTaskTimeout(task: AbstractPluginTask, onTimeout?: () => void): void { // Set up a timeout to reject the task if no response is received diff --git a/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.test.ts b/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.test.ts index 4c1b153916..5d7c518bb9 100644 --- a/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.test.ts +++ b/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.test.ts @@ -150,7 +150,7 @@ class Instance implements TaskDispatchHost { return this.users.get(token!); } - sendLocalTask(task: AbstractPluginTask, connection: PenpotConnection) { + async sendLocalTask(task: AbstractPluginTask, connection: PenpotConnection) { if (!connection.ready) throw new Error("Plugin disconnected"); this.executions.push(connection.session.sessionId); task.resolveWithResult({ data: connection.session.sessionId }); @@ -369,3 +369,27 @@ test("disconnecting the last tab removes its session claim subscription", async await owner.disconnect(connection); assert.equal(network.claims.get("alice")?.size, 0); }); + +test("a forwarded task waits while the owning instance wakes the session", async () => { + const network = new RedisNetwork(); + const owner = new Instance(network); + await owner.connect("alice", "tab"); + const sendLocalTask = owner.sendLocalTask.bind(owner); + owner.sendLocalTask = async (task, connection) => { + await new Promise((resolve) => setTimeout(resolve, 10)); + return sendLocalTask(task, connection); + }; + assert.equal((await new Instance(network).execute("alice", "tab")).data, "tab"); +}); + +test("a forwarded task fails with the owner's error when its session cannot be woken", async () => { + const network = new RedisNetwork(); + const owner = new Instance(network); + await owner.connect("alice", "tab"); + owner.sendLocalTask = async () => { + await new Promise((resolve) => setImmediate(resolve)); + throw new Error("The Penpot plugin tab has been frozen by the browser and cannot run tasks."); + }; + await assert.rejects(new Instance(network).execute("alice", "tab"), /has been frozen/); + assert.deepEqual(owner.executions, []); +}); diff --git a/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.ts b/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.ts index 10c895e659..8acce23110 100644 --- a/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.ts +++ b/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.ts @@ -190,11 +190,9 @@ export class MultiInstanceTaskDispatcher extends TaskDispatcher { const task = new RemotePluginTask(request.task, request.params, this.redisBridge, request.id); this.logger.debug("Dispatching remote task %s as %s to Penpot via WebSocket", request.id, task.id); - try { - this.host.sendLocalTask(task, connection); - } catch (error) { + void this.host.sendLocalTask(task, connection).catch((error) => { task.rejectWithError(error instanceof Error ? error : new Error(String(error))); - } + }); } /** Serializes subscription changes for a user across connects and disconnects. */ diff --git a/mcp/packages/server/src/dispatch/SingleInstanceTaskDispatcher.ts b/mcp/packages/server/src/dispatch/SingleInstanceTaskDispatcher.ts index 60ed052105..40e677308e 100644 --- a/mcp/packages/server/src/dispatch/SingleInstanceTaskDispatcher.ts +++ b/mcp/packages/server/src/dispatch/SingleInstanceTaskDispatcher.ts @@ -18,7 +18,7 @@ export class SingleInstanceTaskDispatcher extends TaskDispatcher { if (!connection) { throw new Error(`Penpot session ${JSON.stringify(targetId)} is not connected for this user.`); } - this.host.sendLocalTask(task, connection); + await this.host.sendLocalTask(task, connection); } async onNewConnection(): Promise {} diff --git a/mcp/packages/server/src/dispatch/TaskDispatcher.ts b/mcp/packages/server/src/dispatch/TaskDispatcher.ts index 062af6b80c..b941b57747 100644 --- a/mcp/packages/server/src/dispatch/TaskDispatcher.ts +++ b/mcp/packages/server/src/dispatch/TaskDispatcher.ts @@ -6,7 +6,7 @@ import type { UserPenpotConnections } from "../UserPenpotConnections"; /** Local connection access and tracked task execution provided by PluginBridge. */ export interface TaskDispatchHost { getUserConnections(userToken: string | null): UserPenpotConnections | undefined; - sendLocalTask(task: AbstractPluginTask, connection: PenpotConnection): void; + sendLocalTask(task: AbstractPluginTask, connection: PenpotConnection): Promise; sendRemoteTask(task: AbstractPluginTask, userToken: string, sessionId: string): void; displaceConnection(connection: PenpotConnection): void; } diff --git a/mcp/packages/server/src/wake/LivenessWaiter.test.ts b/mcp/packages/server/src/wake/LivenessWaiter.test.ts new file mode 100644 index 0000000000..4d69f5ccf5 --- /dev/null +++ b/mcp/packages/server/src/wake/LivenessWaiter.test.ts @@ -0,0 +1,88 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import type { PenpotConnection } from "../PenpotConnection"; +import { LivenessWaiter } from "./LivenessWaiter"; + +function frozenConnection(sessionId = "tab", userToken: string | null = "alice"): PenpotConnection { + return { + userToken, + session: { sessionId, fileId: null, fileName: null }, + lastHeartbeat: 0, + frozen: true, + } as PenpotConnection; +} + +/** Records a plugin message as the plugin bridge does. */ +function receive(waiter: LivenessWaiter, connection: PenpotConnection, frozen = false): void { + connection.lastHeartbeat = Date.now(); + connection.frozen = frozen; + waiter.notify(connection); +} + +test("resolves with the connection of the first plugin message after the wait starts", async () => { + const waiter = new LivenessWaiter(); + const connection = frozenConnection(); + const { signal } = waiter.wait(connection, 1000); + receive(waiter, connection); + assert.equal(await signal, connection); +}); + +test("a message received before the wait starts does not satisfy it", async () => { + const waiter = new LivenessWaiter(); + const connection = frozenConnection(); + receive(waiter, connection); + const { signal } = waiter.wait(connection, 20); + await assert.rejects(signal, /No plugin signal arrived within 20ms/); +}); + +test("a freeze report does not satisfy the wait", async () => { + const waiter = new LivenessWaiter(); + const connection = frozenConnection(); + const { signal } = waiter.wait(connection, 1000); + receive(waiter, connection, true); + let settled = false; + void signal.then(() => (settled = true)); + await new Promise((resolve) => setImmediate(resolve)); + assert.equal(settled, false); + receive(waiter, connection); + await signal; +}); + +test("a new connection for the same session satisfies the wait with that connection", async () => { + const waiter = new LivenessWaiter(); + const previous = frozenConnection(); + const { signal } = waiter.wait(previous, 1000); + const replacement = frozenConnection(); + replacement.frozen = false; + waiter.notify(replacement); + assert.equal(await signal, replacement); +}); + +test("resolves every wait of the session and none of other sessions or users", async () => { + const waiter = new LivenessWaiter(); + const connection = frozenConnection(); + const first = waiter.wait(connection, 1000); + const second = waiter.wait(connection, 1000); + const otherSession = waiter.wait(frozenConnection("other"), 20); + const otherUser = waiter.wait(frozenConnection("tab", "bob"), 20); + receive(waiter, connection); + await Promise.all([first.signal, second.signal]); + await assert.rejects(otherSession.signal, /No plugin signal/); + await assert.rejects(otherUser.signal, /No plugin signal/); +}); + +test("a cancelled wait neither settles nor keeps its timeout", async () => { + const waiter = new LivenessWaiter(); + const connection = frozenConnection(); + const wait = waiter.wait(connection, 10); + wait.cancel(); + let settled = false; + void wait.signal.then( + () => (settled = true), + () => (settled = true) + ); + await new Promise((resolve) => setTimeout(resolve, 30)); + receive(waiter, connection); + await new Promise((resolve) => setImmediate(resolve)); + assert.equal(settled, false); +}); diff --git a/mcp/packages/server/src/wake/LivenessWaiter.ts b/mcp/packages/server/src/wake/LivenessWaiter.ts new file mode 100644 index 0000000000..37a1f3a5ea --- /dev/null +++ b/mcp/packages/server/src/wake/LivenessWaiter.ts @@ -0,0 +1,85 @@ +import type { PenpotConnection } from "../PenpotConnection"; + +/** A pending wait for a fresh liveness signal from the plugin of a session. */ +export interface LivenessWait { + /** resolves with the connection that showed the plugin can run, or rejects when the wait times out. */ + readonly signal: Promise; + + /** Stops waiting, leaving the signal unsettled. */ + cancel(): void; +} + +/** One registered wait and the means to settle it. */ +interface Waiter { + readonly resolve: (connection: PenpotConnection) => void; + readonly reject: (error: Error) => void; + readonly timeout: NodeJS.Timeout; +} + +/** + * Tracks waits for fresh liveness signals from the plugins of sessions. + * + * A wait is satisfied by the first signal from the session after it starts: a plugin message + * that leaves its connection unfrozen, or a new connection for the session completing its + * initialization, since a page resuming from a browser freeze may replace its WebSocket. The + * owner of the connections reports both via {@link notify}. + */ +export class LivenessWaiter { + private readonly waiters = new Map>(); + + /** + * Starts waiting for the next liveness signal from the session of a connection. + * + * @param connection - A connection of the session to observe + * @param timeoutMs - Maximum time, in ms, to wait for the signal + */ + wait(connection: PenpotConnection, timeoutMs: number): LivenessWait { + const key = LivenessWaiter.sessionKey(connection); + let waiter!: Waiter; + const signal = new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + this.release(key, waiter); + reject(new Error(`No plugin signal arrived within ${timeoutMs}ms.`)); + }, timeoutMs); + waiter = { resolve, reject, timeout }; + }); + const waiters = this.waiters.get(key) ?? new Set(); + waiters.add(waiter); + this.waiters.set(key, waiters); + return { signal, cancel: () => this.release(key, waiter) }; + } + + /** + * Reports a plugin message, after its connection's liveness state has been updated, or a + * connection that has completed its initialization. + * + * @param connection - The connection that sent the message or was initialized + */ + notify(connection: PenpotConnection): void { + if (connection.frozen) { + return; + } + const key = LivenessWaiter.sessionKey(connection); + const waiters = [...(this.waiters.get(key) ?? [])]; + this.waiters.delete(key); + for (const waiter of waiters) { + clearTimeout(waiter.timeout); + waiter.resolve(connection); + } + } + + /** Identifies the session of a connection across its user's reconnections. */ + private static sessionKey(connection: PenpotConnection): string { + return JSON.stringify([connection.userToken, connection.session.sessionId]); + } + + /** Removes a single wait, clearing its timeout. */ + private release(key: string, waiter: Waiter): void { + clearTimeout(waiter.timeout); + const waiters = this.waiters.get(key); + waiters?.delete(waiter); + if (waiters?.size === 0) { + this.waiters.delete(key); + } + } +} diff --git a/mcp/packages/server/src/wake/SessionWaker.test.ts b/mcp/packages/server/src/wake/SessionWaker.test.ts new file mode 100644 index 0000000000..e026513c1d --- /dev/null +++ b/mcp/packages/server/src/wake/SessionWaker.test.ts @@ -0,0 +1,206 @@ +import assert from "node:assert/strict"; +import { afterEach, beforeEach, test } from "node:test"; +import type { WakeRequest } from "@penpot/mcp-common"; +import type { PenpotConnection } from "../PenpotConnection"; +import { UserPenpotConnections } from "../UserPenpotConnections"; +import { LivenessWaiter } from "./LivenessWaiter"; +import { LocalSessionWaker } from "./SessionWaker"; +import { WakeAgentRegistry, type WakeAgentSocket } from "./WakeAgentRegistry"; + +/** Agent connection answering each wake request with a fixed outcome. */ +class FakeAgent { + readonly requests: WakeRequest[] = []; + onWake?: () => void; + + constructor(private readonly outcome: string | null) {} + + send(data: string): void { + const request: WakeRequest = JSON.parse(data); + this.requests.push(request); + if (this.outcome !== null) { + // answer asynchronously, as a real agent would + setImmediate(() => { + registry.handleMessage(this.asSocket(), { + type: "wake-result", + requestId: request.requestId, + outcome: this.outcome, + }); + this.onWake?.(); + }); + } + } + + close(): void {} + + ping(): void {} + + asSocket(): WakeAgentSocket { + return this as unknown as WakeAgentSocket; + } +} + +let registry: WakeAgentRegistry; +let liveness: LivenessWaiter; +let open: UserPenpotConnections; +let waker: LocalSessionWaker; + +beforeEach(() => { + registry = new WakeAgentRegistry(); + liveness = new LivenessWaiter(); + open = new UserPenpotConnections(); + waker = new LocalSessionWaker(registry, liveness, { getUserConnections: () => open }, 100, 10); +}); + +afterEach(() => { + registry.close(); +}); + +/** Creates an open connection of the session and registers it as the session's connection. */ +function connection(state: { frozen?: boolean; lastHeartbeat?: number } = {}): PenpotConnection { + // the waker matches connections by user token and session ID + const created = { + userToken: "alice", + session: { sessionId: "tab", fileId: null, fileName: null }, + frozen: state.frozen ?? true, + lastHeartbeat: state.lastHeartbeat ?? Date.now(), + ready: true, + socket: { readyState: 1 }, + } as PenpotConnection; + open.add(created); + return created; +} + +/** Closes a connection as the plugin bridge does when its socket closes. */ +function close(target: PenpotConnection): void { + target.ready = false; + open.remove(target); +} + +function addAgent(outcome: string | null): FakeAgent { + const agent = new FakeAgent(outcome); + registry.add(agent.asSocket(), "alice", { type: "initialize-wake-agent", agentId: "agent", sessionIds: ["tab"] }); + return agent; +} + +/** Records a plugin message as the plugin bridge does. */ +function receive(target: PenpotConnection): void { + target.lastHeartbeat = Date.now(); + target.frozen = false; + liveness.notify(target); +} + +test("returns the connection at once for a responsive plugin without asking the agent", async () => { + const agent = addAgent("activated"); + const target = connection({ frozen: false }); + assert.equal(await waker.ensureResponsive(target), target); + assert.equal(agent.requests.length, 0); +}); + +test("fails with the liveness error when no agent can wake the session", async () => { + await assert.rejects(waker.ensureResponsive(connection()), /has been frozen/); +}); + +test("wakes a frozen session and returns once the plugin signals", async () => { + const agent = addAgent("activated"); + const target = connection(); + agent.onWake = () => receive(target); + assert.equal(await waker.ensureResponsive(target), target); + assert.equal(agent.requests.length, 1); + assert.equal(agent.requests[0].sessionId, "tab"); +}); + +test("returns the new connection when the woken page replaces its connection", async () => { + const agent = addAgent("activated"); + const target = connection(); + agent.onWake = () => { + close(target); + liveness.notify(connection({ frozen: false })); + }; + const live = await waker.ensureResponsive(target); + assert.notEqual(live, target); + assert.equal(live, open.get("tab")); +}); + +test("waits for the new connection when the woken page closes its socket right after its first message", async () => { + const agent = addAgent("activated"); + const target = connection(); + let replacement: PenpotConnection | undefined; + agent.onWake = () => { + receive(target); + close(target); + setTimeout(() => { + replacement = connection({ frozen: false }); + liveness.notify(replacement); + }, 30); + }; + const live = await waker.ensureResponsive(target); + assert.equal(live, replacement); +}); + +test("fails with the liveness error when the woken page closes its socket and does not reconnect", async () => { + const agent = addAgent("activated"); + const target = connection(); + agent.onWake = () => { + receive(target); + close(target); + }; + await assert.rejects(waker.ensureResponsive(target), /has been frozen/); +}); + +test("wakes a session whose heartbeat is stale", async () => { + const agent = addAgent("activated"); + const target = connection({ frozen: false, lastHeartbeat: Date.now() - 60_000 }); + agent.onWake = () => receive(target); + await waker.ensureResponsive(target); +}); + +test("a signal arriving before the agent's answer is not missed", async () => { + const agent = addAgent("activated"); + const target = connection(); + agent.send = (data: string) => { + const request: WakeRequest = JSON.parse(data); + receive(target); + setImmediate(() => + registry.handleMessage(agent.asSocket(), { + type: "wake-result", + requestId: request.requestId, + outcome: "activated", + }) + ); + }; + await waker.ensureResponsive(target); +}); + +test("fails with the liveness error when the agent cannot wake the tab", async () => { + addAgent("tab-closed"); + await assert.rejects(waker.ensureResponsive(connection()), /has been frozen/); +}); + +test("fails with the liveness error when the agent does not answer", async () => { + addAgent(null); + await assert.rejects(waker.ensureResponsive(connection()), /has been frozen/); +}); + +test("fails with the liveness error when no plugin signal follows the activation", async () => { + addAgent("activated"); + const target = connection({ frozen: false, lastHeartbeat: Date.now() - 60_000 }); + await assert.rejects(waker.ensureResponsive(target), /no heartbeat for 60s/); +}); + +test("concurrent tasks for one session share a single wake request", async () => { + const agent = addAgent("activated"); + const target = connection(); + const sibling = connection(); + agent.onWake = () => receive(target); + await Promise.all([waker.ensureResponsive(target), waker.ensureResponsive(sibling)]); + assert.equal(agent.requests.length, 1); +}); + +test("a later task after a failed wake tries again", async () => { + const agent = addAgent("activated"); + const target = connection(); + await assert.rejects(waker.ensureResponsive(target), /has been frozen/); + agent.onWake = () => receive(target); + await waker.ensureResponsive(target); + assert.equal(agent.requests.length, 2); +}); diff --git a/mcp/packages/server/src/wake/SessionWaker.ts b/mcp/packages/server/src/wake/SessionWaker.ts new file mode 100644 index 0000000000..0cf2521dd2 --- /dev/null +++ b/mcp/packages/server/src/wake/SessionWaker.ts @@ -0,0 +1,139 @@ +import { assertPluginResponsive, type PenpotConnection } from "../PenpotConnection"; +import type { UserPenpotConnections } from "../UserPenpotConnections"; +import { createLogger } from "../logger"; +import type { LivenessWaiter } from "./LivenessWaiter"; +import type { WakeAgentRegistry } from "./WakeAgentRegistry"; + +/** the maximum time, in ms, to wait for a woken plugin to show it can run. */ +export const WAKE_WAIT_MS = 10_000; + +/** + * the time, in ms, a woken page gets to settle before its task is sent: a page resuming from a + * browser freeze may close its WebSocket right after its first message. + */ +export const WAKE_SETTLE_MS = 250; + +/** + * Strategy ensuring that a plugin connection can run a task before the task is sent to it. + */ +export interface SessionWaker { + /** + * Returns once the plugin of the connection's session can run tasks. + * + * @param connection - The connection the task is about to be sent to + * @returns The connection to send the task to: the given one, or the connection that replaced + * it for the same session while the plugin was being woken + * @throws Error the error of {@link assertPluginResponsive} if the plugin cannot run tasks + */ + ensureResponsive(connection: PenpotConnection): Promise; +} + +/** Lookup of the ready connections of each user. */ +export interface SessionConnections { + getUserConnections(userToken: string | null): UserPenpotConnections | undefined; +} + +/** + * Session waker that asks a wake agent connected to this instance to wake the tab of an + * unresponsive plugin, then waits for a fresh liveness signal from the plugin. + * + * After the signal, the page gets a moment to settle, and the task goes to the session's open + * connection at that point, waiting for a new one if the page closed its socket on resume. + * Without an agent able to wake the session, it fails at once, as the plain liveness check does. + * Concurrent tasks for the same session share a single wake attempt. + */ +export class LocalSessionWaker implements SessionWaker { + private readonly logger = createLogger("LocalSessionWaker"); + private readonly wakes = new Map>(); + + /** + * @param registry - The wake agents connected to this instance + * @param liveness - The waits for plugin liveness signals, fed by the plugin connections + * @param connections - The lookup of the open connections of each session + * @param waitMs - Maximum time, in ms, for the agent's answer and the plugin's signal together + * @param settleMs - Time, in ms, a woken page gets to settle before its task is sent + */ + constructor( + private readonly registry: WakeAgentRegistry, + private readonly liveness: LivenessWaiter, + private readonly connections: SessionConnections, + private readonly waitMs: number = WAKE_WAIT_MS, + private readonly settleMs: number = WAKE_SETTLE_MS + ) {} + + async ensureResponsive(connection: PenpotConnection): Promise { + try { + assertPluginResponsive(connection, Date.now()); + return connection; + } catch (error) { + if (!this.registry.canWake(connection.userToken, connection.session.sessionId)) { + throw error; + } + const key = JSON.stringify([connection.userToken, connection.session.sessionId]); + let wake = this.wakes.get(key); + if (!wake) { + wake = this.wake(connection).finally(() => this.wakes.delete(key)); + this.wakes.set(key, wake); + } + try { + return await wake; + } catch (wakeError) { + this.logger.info("Could not wake the plugin session: %s", (wakeError as Error).message); + throw error; + } + } + } + + /** + * Asks the agent to wake the session and waits until the session has an open connection whose + * plugin has shown it can run. + * + * @returns The connection to send the task to + */ + private async wake(connection: PenpotConnection): Promise { + const deadline = Date.now() + this.waitMs; + // start waiting before the request, so that a fast signal is not missed + const wait = this.liveness.wait(connection, this.waitMs); + const signal = wait.signal; + // the signal may fail while the agent's answer is awaited + signal.catch(() => {}); + try { + const outcome = await this.registry.wake(connection.userToken, connection.session.sessionId, this.waitMs); + if (outcome !== "activated") { + throw new Error(`The wake agent answered ${JSON.stringify(outcome)}.`); + } + await signal; + } finally { + wait.cancel(); + } + await new Promise((resolve) => setTimeout(resolve, this.settleMs)); + const live = await this.waitForOpenConnection(connection, deadline); + this.logger.info( + live === connection ? "Woke the plugin session" : "Woke the plugin session on a new connection" + ); + return live; + } + + /** Returns the session's open connection, waiting for a new one until the deadline if there is none. */ + private async waitForOpenConnection(connection: PenpotConnection, deadline: number): Promise { + const { userToken } = connection; + const { sessionId } = connection.session; + for (;;) { + const open = this.connections.getUserConnections(userToken)?.get(sessionId); + if (open) { + return open; + } + const remaining = deadline - Date.now(); + if (remaining <= 0) { + throw new Error("The woken plugin closed its connection and did not reconnect in time."); + } + const wait = this.liveness.wait(connection, remaining); + try { + // a connection initialized after the lookup above has notified this wait + await wait.signal; + } finally { + wait.cancel(); + } + } + } +}