✨ Wake frozen Penpot tabs before dispatching MCP tasks

Before a task goes to a plugin whose tab the browser has frozen, ask a
wake agent to wake the tab and wait until the plugin can run again.

A LivenessWaiter tells when the plugin of a session shows it can run:
the first plugin message after the wait starts that leaves its
connection unfrozen, or a new connection for the session completing its
initialization, since a page resuming from a freeze may replace its
WebSocket.

The SessionWaker strategy makes sure a plugin can run a task before the
task is sent. When the plugin is frozen or its heartbeat is stale,
LocalSessionWaker asks the agent holding the session to wake the tab and
waits for the session's next liveness signal. It then gives the page
250 ms to settle and returns the session's open connection, waiting for
a new one when the page has closed its socket on resume, as Chrome does.
The whole wait is bounded by 10 s. Without an agent, or when any step
fails, the task fails with the same error as before. Concurrent tasks
for one session share a single wake. The liveness check moves to
PenpotConnection, so the waker can use it without the plugin bridge.

sendLocalTask becomes asynchronous and waits for the waker. Both
dispatchers await it. The task timeout starts once the task is sent, so
a wake never replays a task that already timed out.

When a plugin connection closes, its tasks still awaiting a response
fail at once instead of at their timeout, since responses only count
from the connection a task went to. A resuming page closes its old
socket right after a last heartbeat, so a task may go out on it.
Connections log their close code and reason, to tell a page closing its
socket from a browser or network failure.

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;

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;

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);
+});

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<void> {
+        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<any, any>, onTimeout?: () => void): void {
         // Set up a timeout to reject the task if no response is received

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, []);
+});

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. */

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<void> {}

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<void>;
     sendRemoteTask(task: AbstractPluginTask, userToken: string, sessionId: string): void;
     displaceConnection(connection: PenpotConnection): void;
 }

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);
+});

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<PenpotConnection>;
+
+    /** 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<string, Set<Waiter>>();
+
+    /**
+     * 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<PenpotConnection>((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);
+        }
+    }
+}

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);
+});

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<PenpotConnection>;
+}
+
+/** 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<string, Promise<PenpotConnection>>();
+
+    /**
+     * @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<PenpotConnection> {
+        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<PenpotConnection> {
+        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<PenpotConnection> {
+        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();
+            }
+        }
+    }
+}