github/penpot
1
0
Fork 0
mirror of https://github.com/penpot/penpot.git synced 2026-10-29 05:50:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

✨ Route MCP wake requests across instances through Redis

Browse Source 
The instance that holds a frozen plugin's connection may not hold the
browser extension that can wake its tab. The session waker now asks wake
agents through a WakeAgents interface. With Redis, an agent of this
instance is asked first; otherwise the request goes out on a wake channel
per user, to which an instance subscribes while it holds an agent of that
user. Each instance answers on a response channel per request: with its
agent's outcome, or with unknown-session at once when none of its agents
holds the session. The first activation wins. A request that reaches no
instance fails at once, so a frozen session without an agent still fails
as before.

When a forwarded task finds its plugin frozen, the instance holding the
plugin connection reports the wake on the task's response channel. A
waking notice lets the requester extend its timeout by the wake budget,
and a sent notice restarts the plain task timeout, so the wake does not
count against the task. Tasks that need no wake send no notices.

A woken page reconnects through the load balancer, perhaps to another
instance. With Redis the waker does not wait for a local reconnect:
when the woken page has closed its socket, the owner sends a resend
notice and the requester publishes the task again until an instance
holds the session, for up to the wake wait. A closing connection also
ends a pending liveness wait, so a page that closes without a last
heartbeat is not waited for in vain.
This commit is contained in:
alonso.torres 2026-10-26 10:25:57 +00:00
parent 57640425b9
commit b132939b71
14 changed files with 1008 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
mcp/packages/server/src/PluginBridge.sessions.test.ts
View File

@ -251,13 +251,25 @@ test("a plugin message resolves a liveness wait started before it", async () =>
assert.equal(connection.frozen, false);
});
test("a plugin reconnecting for the same session resolves a liveness wait with its new connection", async () => {
test("a closing plugin connection ends a liveness wait of its session", async () => {
const plugin = await connect("tab");
const connection = bridge.getUserConnections("alice")!.getRegistered("tab")!;
const { signal } = bridge.liveness.wait(connection, 1000);
plugin.close();
assert.equal(await signal, connection);
assert.equal(connection.ready, 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 closed = once(plugin, "close");
plugin.close();
await closed;
while (bridge.getUserConnections("alice")?.get("tab")) {
await new Promise((resolve) => setImmediate(resolve));
}
const { signal } = bridge.liveness.wait(connection, 1000);
await connect("tab");
const replacement = await signal;
assert.notEqual(replacement, connection);

136
mcp/packages/server/src/PluginBridge.ts
View File

@ -5,7 +5,7 @@ import { AbstractPluginTask, PluginTask } from "./PluginTask";
import { PluginTaskResponse, PluginTaskResult } from "@penpot/mcp-common";
import { createLogger } from "./logger";
import type { PenpotMcpServer } from "./PenpotMcpServer";
import type { RedisBridge } from "./RedisBridge";
import type { RedisBridge, TaskNotice } from "./RedisBridge";
import type { PenpotConnection } from "./PenpotConnection";
import { UserPenpotConnections } from "./UserPenpotConnections";
import type { TaskDispatcher, TaskDispatchHost } from "./dispatch/TaskDispatcher";
@ -14,10 +14,13 @@ import { MultiInstanceTaskDispatcher } from "./dispatch/MultiInstanceTaskDispatc
import { WakeAgentRegistry } from "./wake/WakeAgentRegistry";
import { wakeAgentInitSchema } from "./wake/WakeAgentProtocol";
import { LivenessWaiter } from "./wake/LivenessWaiter";
import { LocalSessionWaker, type SessionWaker } from "./wake/SessionWaker";
import { LocalSessionWaker, WAKE_SETTLE_MS, WAKE_WAIT_MS, type SessionWaker } from "./wake/SessionWaker";
import { RedisWakeAgents } from "./wake/RedisWakeAgents";
const KEEP_ALIVE_TIME = 30000; // 30 seconds
const INITIALIZATION_TIMEOUT_MS = 10_000;
/** the interval, in ms, between publications of a task for a woken page that is reconnecting. */
const RESEND_INTERVAL_MS = 250;
const connectionInitSchema = z.object({
type: z.literal("initialize"),
@ -50,7 +53,9 @@ export class PluginBridge implements TaskDispatchHost {
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);
/** the wake agents of the other instances, when tasks are routed across instances. */
private readonly remoteWakeAgents?: RedisWakeAgents;
private readonly sessionWaker: SessionWaker;
/**
* Creates the plugin bridge and starts its WebSocket server.
@ -60,7 +65,8 @@ export class PluginBridge implements TaskDispatchHost {
* @param redisBridge - Optional Redis bridge enabling multi-instance task routing.
* When provided, tasks handled by this instance are routed to the instance
* holding the relevant plugin's WebSocket connection (which may be this same
* instance) via Redis, rather than dispatched directly over a local socket.
* instance) via Redis, rather than dispatched directly over a local socket, and
* frozen tabs can be woken by wake agents connected to other instances.
* @param taskTimeoutSecs - Timeout, in seconds, for plugin task execution
* (defaults to {@link DEFAULT_TASK_TIMEOUT_SECS})
*/
@ -73,6 +79,21 @@ export class PluginBridge implements TaskDispatchHost {
this.dispatcher = redisBridge
? new MultiInstanceTaskDispatcher(this, redisBridge)
: new SingleInstanceTaskDispatcher(this);
if (redisBridge) {
this.remoteWakeAgents = new RedisWakeAgents(this.wakeAgents, redisBridge);
this.wakeAgents.observe(this.remoteWakeAgents);
// behind a shared endpoint, a woken page may reconnect to any instance
this.sessionWaker = new LocalSessionWaker(
this.remoteWakeAgents,
this.liveness,
this,
WAKE_WAIT_MS,
WAKE_SETTLE_MS,
"any-instance"
);
} else {
this.sessionWaker = new LocalSessionWaker(this.wakeAgents, this.liveness, this);
}
this.wsServer = new WebSocketServer({ port: port, host: mcpServer.host });
this.setupWebSocketHandlers();
}
@ -278,6 +299,8 @@ export class PluginBridge implements TaskDispatchHost {
this.rejectPendingTask(taskId, new Error("The plugin connection closed before the task finished."));
}
}
// a wake waiting for this session stops waiting for the closed connection
this.liveness.notifyClosed(connection);
if (connection.userToken !== null && connections.size === 0) {
this.userPenpotConnectionsByToken.delete(connection.userToken);
}
@ -402,36 +425,96 @@ export class PluginBridge implements TaskDispatchHost {
* if the published request reached no instance, i.e. if no instance holds a plugin
* connection for the user's session, or if publishing fails outright.
*
* The instance holding the plugin connection reports the wake of a frozen plugin's tab
* through task notices: the timeout makes room for the wake, and the task is published
* again when the woken page reconnects, possibly to another instance.
*
* @param task - The task to dispatch
* @param userToken - The user token identifying the target session's owner
* @param sessionId - The target Penpot session
*/
public sendRemoteTask(task: AbstractPluginTask, userToken: string, sessionId: string): void {
const redisBridge = this.redisBridge!;
this.logger.debug("Dispatching task %s via Redis", task.id);
// register the task for result correlation, then publish the request via Redis
this.pendingTasks.set(task.id, task);
void redisBridge
.sendTaskRequest(userToken, sessionId, task.toRequest(), (response) =>
this.handlePluginTaskResponse(response)
)
this.publishRemoteTask(task, userToken, sessionId);
this.startTaskTimeout(task, () => this.releaseRemoteTask(task));
}
/**
* Publishes the request of a remote task via Redis.
*
* Without a resend deadline, the task fails at once if the request reaches no instance.
* With one, the request is published again until it reaches an instance or the deadline
* passes, since a woken page reconnects after a moment.
*/
private publishRemoteTask(
task: AbstractPluginTask,
userToken: string,
sessionId: string,
resendDeadline?: number
): void {
void this.redisBridge!.sendTaskRequest(
userToken,
sessionId,
task.toRequest(),
(response) => this.handlePluginTaskResponse(response),
(notice) => this.handleTaskNotice(task, userToken, sessionId, notice)
)
.then((receiverCount) => {
// fail fast when no instance received the request (no connection with matching user token and session ID in any instance)
if (receiverCount === 0) {
if (receiverCount > 0 || !this.pendingTasks.has(task.id)) {
return;
}
if (resendDeadline === undefined) {
// fail fast when no instance received the request (no connection with matching user token and session ID in any instance)
this.rejectPendingTask(
task.id,
new Error(`Penpot session ${JSON.stringify(sessionId)} is not connected for this user.`)
);
} else if (Date.now() < resendDeadline) {
setTimeout(() => {
if (this.pendingTasks.has(task.id)) {
this.publishRemoteTask(task, userToken, sessionId, resendDeadline);
}
}, RESEND_INTERVAL_MS);
} else {
this.rejectPendingTask(
task.id,
new Error("The woken plugin closed its connection and did not reconnect in time.")
);
}
})
.catch((error) => {
this.rejectPendingTask(task.id, error instanceof Error ? error : new Error(String(error)));
});
}
// on timeout, release the response-channel subscription, since no response
// will arrive to trigger its self-unsubscribe.
this.startTaskTimeout(task, () => void redisBridge.unsubscribeFromResponse(task.id));
/** Follows the wake of a remote task's plugin, as reported by the instance holding its connection. */
private handleTaskNotice(task: AbstractPluginTask, userToken: string, sessionId: string, notice: TaskNotice): void {
if (!this.pendingTasks.has(task.id)) {
return;
}
const release = () => this.releaseRemoteTask(task);
switch (notice.notice) {
case "waking":
// the wake and the woken page's settling time do not count against the task
this.armTaskTimeout(task, WAKE_WAIT_MS + WAKE_SETTLE_MS, release);
return;
case "sent":
this.armTaskTimeout(task, 0, release);
return;
case "resend":
this.logger.info("The woken plugin of task %s is reconnecting; publishing the task again", task.id);
this.armTaskTimeout(task, WAKE_WAIT_MS, release);
this.publishRemoteTask(task, userToken, sessionId, Date.now() + WAKE_WAIT_MS);
return;
}
}
/** Releases the response-channel subscription of a remote task no response will settle. */
private releaseRemoteTask(task: AbstractPluginTask): void {
void this.redisBridge!.unsubscribeFromResponse(task.id);
}
/**
@ -445,7 +528,7 @@ export class PluginBridge implements TaskDispatchHost {
this.assertConnected(connection);
// the socket can be open while browser-throttled plugin JS cannot run tasks
const target = await this.sessionWaker.ensureResponsive(connection);
const target = await this.sessionWaker.ensureResponsive(connection, task);
// the connection may have closed after its liveness signal
this.assertConnected(target);
@ -454,6 +537,7 @@ export class PluginBridge implements TaskDispatchHost {
this.taskConnections.set(task.id, target);
target.socket.send(JSON.stringify(task.toRequest()));
this.startTaskTimeout(task);
task.notifySent();
}
/** Throws if the connection cannot receive tasks because its socket is not open. */
@ -465,20 +549,31 @@ export class PluginBridge implements TaskDispatchHost {
/** 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
this.armTaskTimeout(task, 0, onTimeout);
this.logger.info(`Sent task ${task.id}`);
}
/**
* Sets up, or restarts, the timeout that rejects a task if no response is received.
*
* @param task - The pending task
* @param extraMs - Time, in ms, granted on top of the task timeout
* @param onTimeout - Called when the timeout rejects the task
*/
private armTaskTimeout(task: AbstractPluginTask<any, any>, extraMs: number, onTimeout?: () => void): void {
clearTimeout(this.taskTimeouts.get(task.id));
const timeoutMs = this.taskTimeoutSecs * 1000 + extraMs;
const timeoutHandle = setTimeout(() => {
if (
this.rejectPendingTask(
task.id,
new Error(`Task ${task.id} timed out after ${this.taskTimeoutSecs} seconds`)
new Error(`Task ${task.id} timed out after ${timeoutMs / 1000} seconds`)
)
) {
onTimeout?.();
}
}, this.taskTimeoutSecs * 1000);
}, timeoutMs);
this.taskTimeouts.set(task.id, timeoutHandle);
this.logger.info(`Sent task ${task.id}`);
}
/**
@ -491,6 +586,7 @@ export class PluginBridge implements TaskDispatchHost {
socket.terminate();
}
this.wakeAgents.close();
await this.remoteWakeAgents?.close();
for (const taskId of this.pendingTasks.keys()) {
this.rejectPendingTask(taskId, new Error("MCP server is shutting down."));
}

6
mcp/packages/server/src/PluginTask.ts
View File

@ -82,6 +82,12 @@ export abstract class AbstractPluginTask<TParams = any, TResult extends PluginTa
* @param error - The error that occurred during task execution
*/
abstract rejectWithError(error: Error): void;
/** Called when the task waits for the plugin's tab to be woken before it is sent. */
notifyWaking(): void {}
/** Called once the task has been sent to the plugin. */
notifySent(): void {}
}
/**

90
mcp/packages/server/src/RedisBridge.test.ts
View File

@ -195,7 +195,30 @@ function instance(t: TestContext) {
});
return socket;
};
return { bridge, connect };
const connectAgent = async (sessionIds: string[]) => {
const socket = new WebSocket(`ws://127.0.0.1:${wsPort}?userToken=alice`);
sockets.push(socket);
await once(socket, "open");
socket.send(JSON.stringify({ type: "initialize-wake-agent", agentId: "agent", sessionIds }));
// the agent is reachable once registered and subscribed to the user's wake requests
while (
!bridge.wakeAgents.canWake("alice", sessionIds[0]) ||
![...subscriptions].some(([channel, subscribers]) => channel.includes(".wake.req.") && subscribers.size)
) {
await new Promise((resolve) => setImmediate(resolve));
}
return socket;
};
return { bridge, connect, connectAgent };
}
/** Reports a browser freeze from a plugin socket and waits until its instance has recorded it. */
async function freeze(bridge: PluginBridge, plugin: WebSocket, sessionId: string) {
const connection = bridge.getUserConnections("alice")!.getRegistered(sessionId)!;
plugin.send(JSON.stringify({ type: "freeze" }));
while (!connection.frozen) {
await new Promise((resolve) => setImmediate(resolve));
}
}
test("full Redis forwarding resolves the selected WebSocket's result", async (t) => {
@ -243,3 +266,68 @@ test("a reconnect on another instance closes the previous instance's connection
assert.equal((await closed)[0], 1008);
assert.equal((await requester.bridge.executePluginTask(new PluginTask("test", {}))).data, "current");
});
test("a frozen session is woken by a wake agent connected to another instance", async (t) => {
// the wake takes longer than the 0.1 s task timeout, which only starts once the task is sent
const owner = instance(t);
const holder = instance(t);
const requester = instance(t);
const plugin = await owner.connect("tab");
const agent = await holder.connectAgent(["tab"]);
await freeze(owner.bridge, plugin, "tab");
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" }));
});
assert.equal((await requester.bridge.executePluginTask(new PluginTask("test", {}), "tab")).data, "tab");
const wakeResponses = [...subscriptions].filter(([channel]) => channel.includes(".wake.res."));
assert.ok(wakeResponses.every(([, subscribers]) => subscribers.size === 0));
});
test("a frozen session fails at once when no instance holds a wake agent", async (t) => {
const owner = instance(t);
const requester = instance(t);
const plugin = await owner.connect("tab");
await freeze(owner.bridge, plugin, "tab");
await assert.rejects(
requester.bridge.executePluginTask(new PluginTask("test", {}), "tab"),
/frozen by the browser/
);
});
test("a woken page that reconnects to another instance runs the task there", async (t) => {
const owner = instance(t);
const holder = instance(t);
const other = instance(t);
const requester = instance(t);
const plugin = await owner.connect("tab", "success", "previous");
const agent = await holder.connectAgent(["tab"]);
await freeze(owner.bridge, plugin, "tab");
// the resumed page drops its socket and reconnects a moment later, through the load balancer
agent.on("message", (raw) => {
const request = JSON.parse(raw.toString());
agent.send(JSON.stringify({ type: "wake-result", requestId: request.requestId, outcome: "activated" }));
plugin.close(1001);
setTimeout(() => void other.connect("tab", "success", "reconnected"), 400);
});
assert.equal((await requester.bridge.executePluginTask(new PluginTask("test", {}), "tab")).data, "reconnected");
});
test("a task sent to a woken plugin times out with the plain task timeout", async (t) => {
const owner = instance(t);
const holder = instance(t);
const requester = instance(t);
const plugin = await owner.connect("tab", "silent");
const agent = await holder.connectAgent(["tab"]);
await freeze(owner.bridge, plugin, "tab");
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" }));
});
await assert.rejects(
requester.bridge.executePluginTask(new PluginTask("test", {}), "tab"),
/timed out after 0.1 seconds/
);
});

163
mcp/packages/server/src/RedisBridge.ts
View File

@ -1,5 +1,5 @@
import Redis from "ioredis";
import { PluginTaskRequest, PluginTaskResponse, PenpotSession } from "@penpot/mcp-common";
import { PluginTaskRequest, PluginTaskResponse, PenpotSession, WakeOutcome } from "@penpot/mcp-common";
import { createLogger } from "./logger";
import { z } from "zod";
@ -22,6 +22,28 @@ export interface SessionClaim {
sessionId: string;
}
/** Request to wake the tab holding a session, sent to the instances holding the user's wake agents. */
export interface WakeRelayRequest {
id: string;
sessionId: string;
/** maximum time, in ms, the agent may take to answer. */
timeoutMs: number;
}
/** One instance's answer to a {@link WakeRelayRequest}: the agent's outcome, or why it gave none. */
export type WakeRelayResponse = { outcome: WakeOutcome } | { error: string };
const wakeRelayRequestSchema = z.object({
id: z.string().min(1).max(128),
sessionId: z.string().min(1).max(128),
timeoutMs: z.number().int().positive(),
});
const wakeRelayResponseSchema = z.union([
z.object({ outcome: z.enum(["activated", "unknown-session", "tab-closed", "disabled"]) }),
z.object({ error: z.string() }),
]);
const sessionClaimSchema = z.object({
instanceId: z.string().min(1),
claimId: z.string().min(1),
@ -49,6 +71,30 @@ export type TaskRequestHandler = (request: PluginTaskRequest) => void;
*/
export type TaskResponseHandler = (response: PluginTaskResponse<any>) => void;
/**
* Progress of a forwarded task while the instance holding the plugin connection wakes the plugin's
* tab, published on the task's response channel before its response.
*
* - `waking`: the task waits for the tab to be woken
* - `sent`: the task has been sent to the woken plugin
* - `resend`: the woken page closed its connection, so the task must be published again for the
* page's new connection, wherever it reconnects
*/
export interface TaskNotice {
id: string;
notice: "waking" | "sent" | "resend";
}
/**
* Handler invoked for a task notice arriving on a subscribed response channel.
*/
export type TaskNoticeHandler = (notice: TaskNotice) => void;
const taskNoticeSchema = z.object({
id: z.string().min(1),
notice: z.enum(["waking", "sent", "resend"]),
});
/**
* Provides a Redis-backed transport for routing plugin task requests and responses
* between MCP server instances.
@ -128,12 +174,14 @@ export class RedisBridge {
* The response subscription is established *before* the request is published, to
* avoid a race in which the response would be published before the subscription is
* in place. The response handler is invoked at most once and the subscription is
* removed automatically upon delivery (response channels are single-use).
* removed automatically upon delivery (response channels are single-use). Task notices
* arriving before the response leave the subscription in place.
*
* @param userToken - The user token identifying the target plugin's request channel
* @param sessionId - The Penpot session identifying the target connection
* @param request - The serialized plugin task request, passed through verbatim
* @param onResponse - Handler invoked with the response when it arrives
* @param onNotice - Handler invoked with the task notices arriving before the response
* @returns The number of instances that received the request. A count of 0 means no
* instance is subscribed to the session's request channel (i.e. the plugin is not
* connected anywhere); the request was dropped, no response will ever arrive, and
@ -143,19 +191,32 @@ export class RedisBridge {
userToken: string,
sessionId: string,
request: PluginTaskRequest,
onResponse: TaskResponseHandler
onResponse: TaskResponseHandler,
onNotice?: TaskNoticeHandler
): Promise<number> {
const responseChannel = this.responseChannel(request.id);
const requestChannel = this.requestChannel(userToken, sessionId);
this.handlers.set(responseChannel, (rawMessage) => {
let message: unknown;
try {
message = JSON.parse(rawMessage);
} catch (error) {
this.logger.error(error, "Failed to parse task response message");
return;
}
const notice = taskNoticeSchema.safeParse(message);
if (notice.success) {
onNotice?.(notice.data);
return;
}
// a response channel is single-use: remove the handler and unsubscribe on delivery
this.handlers.delete(responseChannel);
void this.subscriber.unsubscribe(responseChannel);
try {
onResponse(JSON.parse(rawMessage) as PluginTaskResponse<any>);
onResponse(message as PluginTaskResponse<any>);
} catch (error) {
this.logger.error(error, "Failed to parse task response message");
this.logger.error(error, "Failed to handle task response message");
}
});
@ -207,6 +268,20 @@ export class RedisBridge {
.catch((error) => this.logger.error(error, "Failed to publish task response"));
}
/**
* Publishes a notice on the response channel for the given task ID.
*
* Used by the instance executing a forwarded task to tell the issuing instance how the
* wake of the plugin's tab proceeds.
*
* @param notice - The notice, carrying the ID of the originally requested task
*/
publishTaskNotice(notice: TaskNotice): void {
void this.publisher
.publish(this.responseChannel(notice.id), JSON.stringify(notice))
.catch((error) => this.logger.error(error, "Failed to publish task notice"));
}
/**
* Subscribes to task requests for the given user's Penpot session.
*
@ -343,6 +418,84 @@ export class RedisBridge {
await this.publisher.publish(this.sessionClaimChannel(userToken), JSON.stringify(claim));
}
private wakeChannel(userToken: string): string {
return `penpot.mcp.${this.tenant}.wake.req.${this.channelKey(userToken)}`;
}
private wakeResponseChannel(requestId: string): string {
return `penpot.mcp.${this.tenant}.wake.res.${this.channelKey(requestId)}`;
}
/**
* Subscribes to the answers to a wake request, then publishes the request to the instances
* holding the user's wake agents.
*
* @returns The number of instances that received the request. With 0 the response subscription
* has already been released.
*/
async sendWakeRequest(
userToken: string,
request: WakeRelayRequest,
onResponse: (response: WakeRelayResponse) => void
): Promise<number> {
const channel = this.wakeResponseChannel(request.id);
this.handlers.set(channel, (raw) => {
try {
onResponse(wakeRelayResponseSchema.parse(JSON.parse(raw)));
} catch (error) {
this.logger.error(error, "Invalid wake response");
}
});
let receiverCount: number;
try {
await this.subscriber.subscribe(channel);
receiverCount = await this.publisher.publish(this.wakeChannel(userToken), JSON.stringify(request));
} catch (error) {
await this.unsubscribeFromWakeResponses(request.id);
throw error;
}
if (receiverCount === 0) {
await this.unsubscribeFromWakeResponses(request.id);
}
return receiverCount;
}
async unsubscribeFromWakeResponses(requestId: string): Promise<void> {
const channel = this.wakeResponseChannel(requestId);
this.handlers.delete(channel);
await this.subscriber.unsubscribe(channel);
}
/** Subscribes to the wake requests for the given user's sessions. */
async subscribeToWakeRequests(userToken: string, handler: (request: WakeRelayRequest) => void): Promise<void> {
const channel = this.wakeChannel(userToken);
this.handlers.set(channel, (raw) => {
try {
handler(wakeRelayRequestSchema.parse(JSON.parse(raw)));
} catch (error) {
this.logger.error(error, "Invalid wake request");
}
});
try {
await this.subscriber.subscribe(channel);
} catch (error) {
await this.unsubscribeFromWakeRequests(userToken);
throw error;
}
}
async unsubscribeFromWakeRequests(userToken: string): Promise<void> {
const channel = this.wakeChannel(userToken);
this.handlers.delete(channel);
await this.subscriber.unsubscribe(channel);
}
publishWakeResponse(requestId: string, response: WakeRelayResponse): void {
void this.publisher
.publish(this.wakeResponseChannel(requestId), JSON.stringify(response))
.catch((error) => this.logger.error(error, "Failed to publish wake response"));
}
/** Closes both Redis connections. Call on server shutdown. */
async close(): Promise<void> {
await this.subscriber.quit();

23
mcp/packages/server/src/RemotePluginTask.ts
View File

@ -1,6 +1,7 @@
import { AbstractPluginTask } from "./PluginTask";
import { PluginTaskResult } from "@penpot/mcp-common";
import type { RedisBridge } from "./RedisBridge";
import { SessionReconnectingError } from "./wake/SessionWaker";
/**
* A plugin task whose outcome is forwarded back to a remote requester via Redis,
@ -18,8 +19,15 @@ import type { RedisBridge } from "./RedisBridge";
*
* It deliberately carries no result promise: settling the task *is* the side effect
* of publishing to Redis, and nothing awaits it locally.
*
* While the plugin's tab is woken, the task also publishes notices on the response channel,
* so that the requester can extend its timeout, and asks the requester to publish the task
* again when the woken page closed its connection to reconnect, possibly to another instance.
*/
export class RemotePluginTask extends AbstractPluginTask<any, PluginTaskResult<any>> {
/** whether the task waited for a wake, after which the requester expects a `sent` notice. */
private waking = false;
/**
* Creates a task that forwards its outcome to a Redis response channel.
*
@ -47,10 +55,25 @@ export class RemotePluginTask extends AbstractPluginTask<any, PluginTaskResult<a
}
rejectWithError(error: Error): void {
if (error instanceof SessionReconnectingError) {
this.redisBridge.publishTaskNotice({ id: this.originalTaskId, notice: "resend" });
return;
}
this.redisBridge.publishTaskResponse(this.originalTaskId, {
id: this.originalTaskId,
success: false,
error: error.message,
});
}
notifyWaking(): void {
this.waking = true;
this.redisBridge.publishTaskNotice({ id: this.originalTaskId, notice: "waking" });
}
notifySent(): void {
if (this.waking) {
this.redisBridge.publishTaskNotice({ id: this.originalTaskId, notice: "sent" });
}
}
}

19
mcp/packages/server/src/wake/LivenessWaiter.ts
View File

@ -2,7 +2,7 @@ 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. */
/** resolves with the connection that showed the plugin can run or that closed; rejects on timeout. */
readonly signal: Promise<PenpotConnection>;
/** Stops waiting, leaving the signal unsettled. */
@ -22,7 +22,8 @@ interface Waiter {
* 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}.
* owner of the connections reports both via {@link notify}. A wait also ends when a connection of
* the session closes ({@link notifyClosed}), since the replacement may connect elsewhere.
*/
export class LivenessWaiter {
private readonly waiters = new Map<string, Set<Waiter>>();
@ -59,6 +60,20 @@ export class LivenessWaiter {
if (connection.frozen) {
return;
}
this.settle(connection);
}
/**
* Reports a closed connection, ending the waits of its session with it.
*
* @param connection - The connection that closed
*/
notifyClosed(connection: PenpotConnection): void {
this.settle(connection);
}
/** Resolves the waits of the connection's session with the connection. */
private settle(connection: PenpotConnection): void {
const key = LivenessWaiter.sessionKey(connection);
const waiters = [...(this.waiters.get(key) ?? [])];
this.waiters.delete(key);

235
mcp/packages/server/src/wake/RedisWakeAgents.test.ts Normal file
View File

@ -0,0 +1,235 @@
import assert from "node:assert/strict";
import { randomUUID } from "node:crypto";
import { afterEach, beforeEach, test } from "node:test";
import type { WakeRequest } from "@penpot/mcp-common";
import type { RedisBridge, WakeRelayRequest, WakeRelayResponse } from "../RedisBridge";
import { RedisWakeAgents } from "./RedisWakeAgents";
import { WakeAgentRegistry, type WakeAgentSocket } from "./WakeAgentRegistry";
import { NoWakeAgentError } from "./WakeAgents";
/** In-memory Pub/Sub transport for wake requests, retaining Redis's per-subscriber delivery count. */
class WakeNetwork {
readonly requests = new Map<string, Map<WakePeer, (request: WakeRelayRequest) => void>>();
readonly responses = new Map<string, (response: WakeRelayResponse) => void>();
published = 0;
}
/** One MCP instance's connection to the test Pub/Sub network. */
class WakePeer {
constructor(readonly network: WakeNetwork) {}
asBridge(): RedisBridge {
return this as unknown as RedisBridge;
}
async subscribeToWakeRequests(token: string, handler: (request: WakeRelayRequest) => void) {
const subscribers = this.network.requests.get(token) ?? new Map();
subscribers.set(this, handler);
this.network.requests.set(token, subscribers);
}
async unsubscribeFromWakeRequests(token: string) {
this.network.requests.get(token)?.delete(this);
}
async sendWakeRequest(token: string, request: WakeRelayRequest, onResponse: (response: WakeRelayResponse) => void) {
this.network.published++;
this.network.responses.set(request.id, onResponse);
const subscribers = [...(this.network.requests.get(token)?.values() ?? [])];
for (const handler of subscribers) handler(request);
if (!subscribers.length) this.network.responses.delete(request.id);
return subscribers.length;
}
async unsubscribeFromWakeResponses(id: string) {
this.network.responses.delete(id);
}
publishWakeResponse(id: string, response: WakeRelayResponse) {
this.network.responses.get(id)?.(response);
}
}
/** Agent connection answering each wake request with a fixed outcome, or never when it is null. */
class FakeAgent {
readonly requests: WakeRequest[] = [];
onRequest?: () => void;
constructor(
private readonly registry: WakeAgentRegistry,
private readonly outcome: string | null
) {}
send(data: string): void {
const request: WakeRequest = JSON.parse(data);
this.requests.push(request);
this.onRequest?.();
if (this.outcome !== null) {
setImmediate(() =>
this.registry.handleMessage(this.asSocket(), {
type: "wake-result",
requestId: request.requestId,
outcome: this.outcome,
})
);
}
}
close(): void {}
ping(): void {}
asSocket(): WakeAgentSocket {
return this as unknown as WakeAgentSocket;
}
}
/** One MCP instance: its local agents and their routing over the network. */
class Instance {
readonly registry = new WakeAgentRegistry();
readonly peer: WakePeer;
readonly agents: RedisWakeAgents;
constructor(network: WakeNetwork) {
this.peer = new WakePeer(network);
this.agents = new RedisWakeAgents(this.registry, this.peer.asBridge());
this.registry.observe(this.agents);
}
async addAgent(sessionIds: string[], outcome: string | null = "activated", token: string | null = "alice") {
const agent = new FakeAgent(this.registry, outcome);
this.registry.add(agent.asSocket(), token, {
type: "initialize-wake-agent",
agentId: randomUUID(),
sessionIds,
});
await flush();
return agent;
}
async removeAgent(agent: FakeAgent) {
this.registry.remove(agent.asSocket());
await flush();
}
async close() {
this.registry.close();
await this.agents.close();
}
}
/** Lets pending subscription changes and agent answers run. */
async function flush() {
for (let i = 0; i < 5; i++) await new Promise((resolve) => setImmediate(resolve));
}
let network: WakeNetwork;
let instances: Instance[];
beforeEach(() => {
network = new WakeNetwork();
instances = [];
});
afterEach(async () => {
for (const instance of instances) await instance.close();
});
function instance(): Instance {
const created = new Instance(network);
instances.push(created);
return created;
}
test("wakes a session with an agent of this instance without asking other instances", async () => {
const local = instance();
const agent = await local.addAgent(["tab"]);
assert.equal(await local.agents.wake("alice", "tab", 1000), "activated");
assert.equal(agent.requests.length, 1);
assert.equal(network.published, 0);
});
test("wakes a session held by an agent of another instance", async () => {
const owner = instance();
const holder = instance();
const agent = await holder.addAgent(["tab"]);
assert.equal(await owner.agents.wake("alice", "tab", 1000), "activated");
assert.equal(agent.requests.length, 1);
assert.equal(network.responses.size, 0);
});
test("an activation wins over instances whose agents do not hold the session", async () => {
const owner = instance();
const other = instance();
const holder = instance();
const otherAgent = await other.addAgent(["other-tab"]);
await holder.addAgent(["tab"]);
assert.equal(await owner.agents.wake("alice", "tab", 1000), "activated");
assert.equal(otherAgent.requests.length, 0);
});
test("passes on a remote agent's outcome other than activation", async () => {
const owner = instance();
const holder = instance();
await holder.addAgent(["tab"], "disabled");
assert.equal(await owner.agents.wake("alice", "tab", 1000), "disabled");
});
test("fails at once when no instance holds an agent of the user", async () => {
const owner = instance();
const other = instance();
await other.addAgent(["tab"], "activated", "bob");
await assert.rejects(owner.agents.wake("alice", "tab", 60_000), NoWakeAgentError);
assert.equal(network.responses.size, 0);
});
test("fails as having no agent when no instance's agent holds the session", async () => {
const owner = instance();
const other = instance();
await other.addAgent(["other-tab"]);
await assert.rejects(owner.agents.wake("alice", "tab", 60_000), NoWakeAgentError);
});
test("does not ask other instances for a session without a user token", async () => {
const owner = instance();
await assert.rejects(owner.agents.wake(null, "tab", 1000), NoWakeAgentError);
assert.equal(network.published, 0);
});
test("passes on the failure of a remote agent that disconnects before answering", async () => {
const owner = instance();
const holder = instance();
const agent = await holder.addAgent(["tab"], null);
agent.onRequest = () => setImmediate(() => holder.registry.remove(agent.asSocket()));
await assert.rejects(owner.agents.wake("alice", "tab", 60_000), /disconnected before answering/);
});
test("fails when an instance does not answer in time and releases the response subscription", async () => {
const owner = instance();
const silent = new WakePeer(network);
await silent.subscribeToWakeRequests("alice", () => {});
await assert.rejects(owner.agents.wake("alice", "tab", 20), /No wake agent answered within 20ms/);
assert.equal(network.responses.size, 0);
});
test("subscribes to a user's wake requests only while holding an agent of the user", async () => {
const holder = instance();
const first = await holder.addAgent(["tab"]);
const second = await holder.addAgent(["other-tab"]);
assert.equal(network.requests.get("alice")?.has(holder.peer), true);
await holder.removeAgent(first);
assert.equal(network.requests.get("alice")?.has(holder.peer), true);
await holder.removeAgent(second);
assert.equal(network.requests.get("alice")?.has(holder.peer), false);
});
test("closing fails pending remote wake requests", async () => {
const owner = instance();
const silent = new WakePeer(network);
await silent.subscribeToWakeRequests("alice", () => {});
const wake = owner.agents.wake("alice", "tab", 60_000);
await flush();
await owner.agents.close();
await assert.rejects(wake, /shutting down/);
assert.equal(network.responses.size, 0);
});

171
mcp/packages/server/src/wake/RedisWakeAgents.ts Normal file
View File

@ -0,0 +1,171 @@
import { randomUUID } from "node:crypto";
import type { WakeOutcome } from "@penpot/mcp-common";
import type { RedisBridge, WakeRelayRequest, WakeRelayResponse } from "../RedisBridge";
import { createLogger } from "../logger";
import type { WakeAgentObserver, WakeAgentRegistry } from "./WakeAgentRegistry";
import { NoWakeAgentError, type WakeAgents } from "./WakeAgents";
/**
* Wake agents of all MCP instances, reached through Redis.
*
* A session is woken by an agent of this instance when one holds it. Otherwise the request is
* published to the instances holding agents of the user. Each answers with its agent's outcome
* when one of its agents holds the session, and with `unknown-session` at once otherwise. The
* first activation wins; without one, the result is decided once every instance has answered.
* An instance subscribes to a user's wake requests while it holds an agent of that user.
*/
export class RedisWakeAgents implements WakeAgents, WakeAgentObserver {
private readonly logger = createLogger("RedisWakeAgents");
private readonly subscribedUsers = new Set<string>();
private readonly subscriptionChanges = new Map<string, Promise<void>>();
private readonly pendingWakes = new Set<() => void>();
/**
* @param local - The wake agents connected to this instance
* @param redisBridge - The transport to the other instances
*/
constructor(
private readonly local: WakeAgentRegistry,
private readonly redisBridge: RedisBridge
) {}
async wake(userToken: string | null, sessionId: string, timeoutMs: number): Promise<WakeOutcome> {
if (this.local.canWake(userToken, sessionId)) {
return this.local.wake(userToken, sessionId, timeoutMs);
}
// instances share wake requests per user, so tokenless agents stay local
if (!userToken) {
throw new NoWakeAgentError();
}
return this.wakeRemotely(userToken, sessionId, timeoutMs);
}
agentsChanged(userToken: string | null): void {
if (!userToken) {
return;
}
const previous = this.subscriptionChanges.get(userToken) ?? Promise.resolve();
const current = previous
.then(() => this.updateSubscription(userToken))
.catch((error) => this.logger.error(error, "Failed to update the wake request subscription"));
this.subscriptionChanges.set(userToken, current);
void current.then(() => {
if (this.subscriptionChanges.get(userToken) === current) {
this.subscriptionChanges.delete(userToken);
}
});
}
/** Fails the pending remote wake requests and waits for subscription changes to finish. */
async close(): Promise<void> {
for (const fail of [...this.pendingWakes]) {
fail();
}
await Promise.allSettled(this.subscriptionChanges.values());
}
/** Subscribes to the user's wake requests while this instance holds an agent of the user. */
private async updateSubscription(userToken: string): Promise<void> {
const wanted = this.local.hasAgents(userToken);
if (wanted && !this.subscribedUsers.has(userToken)) {
await this.redisBridge.subscribeToWakeRequests(userToken, (request) => this.answer(userToken, request));
this.subscribedUsers.add(userToken);
} else if (!wanted && this.subscribedUsers.has(userToken)) {
this.subscribedUsers.delete(userToken);
await this.redisBridge.unsubscribeFromWakeRequests(userToken);
}
}
/** Answers a wake request from another instance (or this one) with the outcome of a local agent. */
private answer(userToken: string, request: WakeRelayRequest): void {
if (!this.local.canWake(userToken, request.sessionId)) {
this.redisBridge.publishWakeResponse(request.id, { outcome: "unknown-session" });
return;
}
this.local.wake(userToken, request.sessionId, request.timeoutMs).then(
(outcome) => this.redisBridge.publishWakeResponse(request.id, { outcome }),
(error) => this.redisBridge.publishWakeResponse(request.id, { error: (error as Error).message })
);
}
/** Asks the instances holding the user's agents to wake the session. */
private wakeRemotely(userToken: string, sessionId: string, timeoutMs: number): Promise<WakeOutcome> {
const request: WakeRelayRequest = { id: randomUUID(), sessionId, timeoutMs };
return new Promise<WakeOutcome>((resolve, reject) => {
const answers: WakeRelayResponse[] = [];
let expectedAnswers: number | undefined;
let settled = false;
const release = () => {
void this.redisBridge
.unsubscribeFromWakeResponses(request.id)
.catch((error) => this.logger.error(error, "Failed to release wake response subscription"));
};
const settle = (finish: () => void) => {
if (settled) {
return;
}
settled = true;
clearTimeout(timeout);
this.pendingWakes.delete(fail);
release();
finish();
};
const fail = () => settle(() => reject(new Error("The MCP server is shutting down.")));
const conclude = () => {
if (answers.some((answer) => "outcome" in answer && answer.outcome === "activated")) {
settle(() => resolve("activated"));
return;
}
if (expectedAnswers === undefined || answers.length < expectedAnswers) {
return;
}
settle(() => this.decide(answers, resolve, reject));
};
const timeout = setTimeout(
() => settle(() => reject(new Error(`No wake agent answered within ${timeoutMs}ms.`))),
timeoutMs
);
this.pendingWakes.add(fail);
// answers can arrive before publish returns its recipient count
this.redisBridge
.sendWakeRequest(userToken, request, (answer) => {
answers.push(answer);
conclude();
})
.then((count) => {
expectedAnswers = count;
conclude();
// a slow subscribe/publish may finish after the timeout
if (settled) {
release();
}
})
.catch((error) => settle(() => reject(error)));
});
}
/**
* Settles a remote wake request that no instance activated, once every instance has answered:
* with an agent's outcome if any agent held the session, else with an agent's failure, else
* as a session no agent holds.
*/
private decide(
answers: WakeRelayResponse[],
resolve: (outcome: WakeOutcome) => void,
reject: (error: Error) => void
): void {
for (const answer of answers) {
if ("outcome" in answer && answer.outcome !== "unknown-session") {
resolve(answer.outcome);
return;
}
}
for (const answer of answers) {
if ("error" in answer) {
reject(new Error(answer.error));
return;
}
}
reject(new NoWakeAgentError());
}
}

55
mcp/packages/server/src/wake/SessionWaker.test.ts
View File

@ -4,7 +4,7 @@ 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 { LocalSessionWaker, SessionReconnectingError } from "./SessionWaker";
import { WakeAgentRegistry, type WakeAgentSocket } from "./WakeAgentRegistry";
/** Agent connection answering each wake request with a fixed outcome. */
@ -74,6 +74,12 @@ function connection(state: { frozen?: boolean; lastHeartbeat?: number } = {}): P
function close(target: PenpotConnection): void {
target.ready = false;
open.remove(target);
liveness.notifyClosed(target);
}
/** Creates a waker for pages that may reconnect to any instance. */
function anyInstanceWaker(waitMs = 1000): LocalSessionWaker {
return new LocalSessionWaker(registry, liveness, { getUserConnections: () => open }, waitMs, 10, "any-instance");
}
function addAgent(outcome: string | null): FakeAgent {
@ -147,6 +153,53 @@ test("fails with the liveness error when the woken page closes its socket and do
await assert.rejects(waker.ensureResponsive(target), /has been frozen/);
});
test("waits for the new connection when the woken page closes its socket without a message", async () => {
const agent = addAgent("activated");
const target = connection();
let replacement: PenpotConnection | undefined;
agent.onWake = () => {
close(target);
setTimeout(() => {
replacement = connection({ frozen: false });
liveness.notify(replacement);
}, 30);
};
assert.equal(await waker.ensureResponsive(target), replacement);
});
test("tells the task's observer when it waits for a wake", async () => {
const agent = addAgent("activated");
const target = connection();
agent.onWake = () => receive(target);
let waking = 0;
await waker.ensureResponsive(target, { notifyWaking: () => waking++ });
assert.equal(waking, 1);
await waker.ensureResponsive(target, { notifyWaking: () => waking++ });
assert.equal(waking, 1);
});
test("fails as reconnecting when the woken page closes its socket and may reconnect to any instance", async () => {
const agent = addAgent("activated");
const target = connection();
agent.onWake = () => close(target);
const started = Date.now();
await assert.rejects(anyInstanceWaker().ensureResponsive(target), SessionReconnectingError);
// the close ends the wait at once instead of the 1000 ms wait for a signal
assert.ok(Date.now() - started < 500);
});
test("returns a connection that reconnected to this instance while the woken page settled", async () => {
const agent = addAgent("activated");
const target = connection();
let replacement: PenpotConnection | undefined;
agent.onWake = () => {
close(target);
replacement = connection({ frozen: false });
liveness.notify(replacement);
};
assert.equal(await anyInstanceWaker().ensureResponsive(target), replacement);
});
test("wakes a session whose heartbeat is stale", async () => {
const agent = addAgent("activated");
const target = connection({ frozen: false, lastHeartbeat: Date.now() - 60_000 });

66
mcp/packages/server/src/wake/SessionWaker.ts
View File

@ -2,7 +2,7 @@ import { assertPluginResponsive, type PenpotConnection } from "../PenpotConnecti
import type { UserPenpotConnections } from "../UserPenpotConnections";
import { createLogger } from "../logger";
import type { LivenessWaiter } from "./LivenessWaiter";
import type { WakeAgentRegistry } from "./WakeAgentRegistry";
import { NoWakeAgentError, type WakeAgents } from "./WakeAgents";
/** the maximum time, in ms, to wait for a woken plugin to show it can run. */
export const WAKE_WAIT_MS = 10_000;
@ -13,6 +13,26 @@ export const WAKE_WAIT_MS = 10_000;
*/
export const WAKE_SETTLE_MS = 250;
/** Observer of a task that waits for the plugin's tab to be woken. */
export interface WakeObserver {
/** Called when the task starts waiting for the wake. */
notifyWaking(): void;
}
/**
* Where a woken page that replaced its WebSocket may reconnect: to this instance only, or to any
* instance behind the same endpoint.
*/
export type ReconnectScope = "this-instance" | "any-instance";
/** Error raised when a woken page closed its connection and may reconnect to another instance. */
export class SessionReconnectingError extends Error {
constructor() {
super("The woken plugin closed its connection and is reconnecting.");
this.name = "SessionReconnectingError";
}
}
/**
* Strategy ensuring that a plugin connection can run a task before the task is sent to it.
*/
@ -21,11 +41,14 @@ 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
* @param observer - The observer of the task, told when the task waits for a wake
* @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 SessionReconnectingError if the woken page closed its connection and may reconnect
* to another instance
* @throws Error the error of {@link assertPluginResponsive} if the plugin cannot run tasks
*/
ensureResponsive(connection: PenpotConnection): Promise<PenpotConnection>;
ensureResponsive(connection: PenpotConnection, observer?: WakeObserver): Promise<PenpotConnection>;
}
/** Lookup of the ready connections of each user. */
@ -34,11 +57,13 @@ export interface SessionConnections {
}
/**
* 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.
* Session waker that asks a wake agent to wake the tab of an unresponsive plugin, then waits for a
* fresh liveness signal from the plugin on this instance.
*
* 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.
* connection at that point. If the page closed its socket on resume, the waker waits for its new
* connection when pages reconnect to this instance, and fails with
* {@link SessionReconnectingError} when they may reconnect to any instance.
* 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.
*/
@ -47,28 +72,28 @@ export class LocalSessionWaker implements SessionWaker {
private readonly wakes = new Map<string, Promise<PenpotConnection>>();
/**
* @param registry - The wake agents connected to this instance
* @param agents - The wake agents that can wake the session
* @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
* @param reconnectScope - Where a woken page that replaced its WebSocket may reconnect
*/
constructor(
private readonly registry: WakeAgentRegistry,
private readonly agents: WakeAgents,
private readonly liveness: LivenessWaiter,
private readonly connections: SessionConnections,
private readonly waitMs: number = WAKE_WAIT_MS,
private readonly settleMs: number = WAKE_SETTLE_MS
private readonly settleMs: number = WAKE_SETTLE_MS,
private readonly reconnectScope: ReconnectScope = "this-instance"
) {}
async ensureResponsive(connection: PenpotConnection): Promise<PenpotConnection> {
async ensureResponsive(connection: PenpotConnection, observer?: WakeObserver): Promise<PenpotConnection> {
try {
assertPluginResponsive(connection, Date.now());
return connection;
} catch (error) {
if (!this.registry.canWake(connection.userToken, connection.session.sessionId)) {
throw error;
}
observer?.notifyWaking();
const key = JSON.stringify([connection.userToken, connection.session.sessionId]);
let wake = this.wakes.get(key);
if (!wake) {
@ -78,7 +103,12 @@ export class LocalSessionWaker implements SessionWaker {
try {
return await wake;
} catch (wakeError) {
this.logger.info("Could not wake the plugin session: %s", (wakeError as Error).message);
if (wakeError instanceof SessionReconnectingError) {
throw wakeError;
}
if (!(wakeError instanceof NoWakeAgentError)) {
this.logger.info("Could not wake the plugin session: %s", (wakeError as Error).message);
}
throw error;
}
}
@ -98,7 +128,7 @@ export class LocalSessionWaker implements SessionWaker {
// 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);
const outcome = await this.agents.wake(connection.userToken, connection.session.sessionId, this.waitMs);
if (outcome !== "activated") {
throw new Error(`The wake agent answered ${JSON.stringify(outcome)}.`);
}
@ -114,7 +144,10 @@ export class LocalSessionWaker implements SessionWaker {
return live;
}
/** Returns the session's open connection, waiting for a new one until the deadline if there is none. */
/**
* Returns the session's open connection. Without one, waits for a new connection until the
* deadline, or fails at once when the page may reconnect to another instance.
*/
private async waitForOpenConnection(connection: PenpotConnection, deadline: number): Promise<PenpotConnection> {
const { userToken } = connection;
const { sessionId } = connection.session;
@ -123,6 +156,9 @@ export class LocalSessionWaker implements SessionWaker {
if (open) {
return open;
}
if (this.reconnectScope === "any-instance") {
throw new SessionReconnectingError();
}
const remaining = deadline - Date.now();
if (remaining <= 0) {
throw new Error("The woken plugin closed its connection and did not reconnect in time.");

16
mcp/packages/server/src/wake/WakeAgentRegistry.test.ts
View File

@ -2,6 +2,7 @@ import assert from "node:assert/strict";
import { afterEach, beforeEach, test } from "node:test";
import type { WakeRequest } from "@penpot/mcp-common";
import { WakeAgentRegistry, type WakeAgentSocket } from "./WakeAgentRegistry";
import { NoWakeAgentError } from "./WakeAgents";
/** In-memory agent connection recording what the registry sends. */
class FakeSocket {
@ -66,7 +67,20 @@ test("passes on outcomes other than activation", async () => {
test("only wakes sessions registered under the same user token", async () => {
addAgent("bob", ["tab"]);
assert.equal(registry.canWake("alice", "tab"), false);
await assert.rejects(registry.wake("alice", "tab", 1000), /No wake agent/);
await assert.rejects(registry.wake("alice", "tab", 1000), NoWakeAgentError);
});
test("tells observers about the users whose agents connect or leave", () => {
const changes: (string | null)[] = [];
registry.observe({ agentsChanged: (token) => changes.push(token) });
const alice = addAgent("alice", ["tab"]);
addAgent(null, ["other"]);
assert.equal(registry.hasAgents("alice"), true);
assert.equal(registry.hasAgents("bob"), false);
registry.remove(alice.asSocket());
assert.equal(registry.hasAgents("alice"), false);
assert.equal(registry.hasAgents(null), true);
assert.deepEqual(changes, ["alice", null, "alice"]);
});
test("single-user agents wake sessions without a token", async () => {

45
mcp/packages/server/src/wake/WakeAgentRegistry.ts
View File

@ -3,6 +3,7 @@ import type { WebSocket } from "ws";
import type { WakeAgentInit, WakeOutcome, WakeRequest } from "@penpot/mcp-common";
import { createLogger } from "../logger";
import { wakeAgentMessageSchema } from "./WakeAgentProtocol";
import { NoWakeAgentError, type WakeAgents } from "./WakeAgents";
const KEEP_ALIVE_TIME_MS = 30_000;
@ -18,6 +19,12 @@ interface WakeAgent {
sessionIds: Set<string>;
}
/** Observer of the users whose wake agents connect to or leave this instance. */
export interface WakeAgentObserver {
/** Called after an agent of the user is added or removed. */
agentsChanged(userToken: string | null): void;
}
/** A wake request awaiting the agent's result. */
interface PendingWake {
readonly agent: WakeAgent;
@ -33,16 +40,22 @@ interface PendingWake {
* Agents are indexed by their connection and scoped to the user token they connected with:
* an agent can only be asked to wake sessions it registered under the same token.
*/
export class WakeAgentRegistry {
export class WakeAgentRegistry implements WakeAgents {
private readonly logger = createLogger("WakeAgentRegistry");
private readonly agentsBySocket = new Map<WakeAgentSocket, WakeAgent>();
private readonly pendingWakes = new Map<string, PendingWake>();
private readonly observers: WakeAgentObserver[] = [];
/**
* @param keepAliveMs - Interval, in ms, of the protocol pings sent to each agent
*/
constructor(private readonly keepAliveMs: number = KEEP_ALIVE_TIME_MS) {}
/** Registers an observer of the users whose agents connect or leave. */
observe(observer: WakeAgentObserver): void {
this.observers.push(observer);
}
/** Whether the socket belongs to a registered wake agent. */
has(socket: WakeAgentSocket): boolean {
return this.agentsBySocket.has(socket);
@ -75,6 +88,7 @@ export class WakeAgentRegistry {
sessionIds: new Set(init.sessionIds),
});
this.logger.info("Wake agent registered with %d session(s)", init.sessionIds.length);
this.notify(userToken);
}
/**
@ -138,6 +152,17 @@ export class WakeAgentRegistry {
}
}
this.logger.info("Wake agent removed");
this.notify(agent.userToken);
}
/** Whether an agent of the user is connected. */
hasAgents(userToken: string | null): boolean {
for (const agent of this.agentsBySocket.values()) {
if (agent.userToken === userToken) {
return true;
}
}
return false;
}
/** Whether a connected agent of the user can wake the session. */
@ -145,19 +170,11 @@ export class WakeAgentRegistry {
return this.findAgent(userToken, sessionId) !== undefined;
}
/**
* Asks the agent holding a session to wake its tab.
*
* @param userToken - The token of the user owning the session
* @param sessionId - The session to wake
* @param timeoutMs - Maximum time, in ms, to wait for the agent's result
* @returns The outcome reported by the agent
* @throws Error if no agent can wake the session, the agent disconnects, or the wait runs out
*/
/** Asks the agent of this instance holding a session to wake its tab. */
wake(userToken: string | null, sessionId: string, timeoutMs: number): Promise<WakeOutcome> {
const agent = this.findAgent(userToken, sessionId);
if (!agent) {
return Promise.reject(new Error("No wake agent is connected for this session."));
return Promise.reject(new NoWakeAgentError());
}
const requestId = randomUUID();
return new Promise<WakeOutcome>((resolve, reject) => {
@ -179,6 +196,12 @@ export class WakeAgentRegistry {
}
}
private notify(userToken: string | null): void {
for (const observer of this.observers) {
observer.agentsChanged(userToken);
}
}
/** Finds the most recently registered agent of the user that can wake the session. */
private findAgent(userToken: string | null, sessionId: string): WakeAgent | undefined {
let found: WakeAgent | undefined;

26
mcp/packages/server/src/wake/WakeAgents.ts Normal file
View File

@ -0,0 +1,26 @@
import type { WakeOutcome } from "@penpot/mcp-common";
/**
* Wake agents that can be asked to wake the browser tab holding a Penpot session.
*/
export interface WakeAgents {
/**
* Asks an agent holding the session to wake its tab.
*
* @param userToken - The token of the user owning the session, or null in single-user mode
* @param sessionId - The session to wake
* @param timeoutMs - Maximum time, in ms, to wait for the agent's result
* @returns The outcome reported by the agent
* @throws NoWakeAgentError if no agent holds the session
* @throws Error if the agent disconnects or does not answer in time
*/
wake(userToken: string | null, sessionId: string, timeoutMs: number): Promise<WakeOutcome>;
}
/** Error raised when no wake agent holds the session to wake. */
export class NoWakeAgentError extends Error {
constructor() {
super("No wake agent is connected for this session.");
this.name = "NoWakeAgentError";
}
}