From b132939b7139e4435ee939b3a76709587176c067 Mon Sep 17 00:00:00 2001 From: "alonso.torres" Date: Mon, 26 Oct 2026 10:25:57 +0000 Subject: [PATCH] :sparkles: Route MCP wake requests across instances through Redis 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. --- .../server/src/PluginBridge.sessions.test.ts | 14 +- mcp/packages/server/src/PluginBridge.ts | 136 ++++++++-- mcp/packages/server/src/PluginTask.ts | 6 + mcp/packages/server/src/RedisBridge.test.ts | 90 ++++++- mcp/packages/server/src/RedisBridge.ts | 163 +++++++++++- mcp/packages/server/src/RemotePluginTask.ts | 23 ++ .../server/src/wake/LivenessWaiter.ts | 19 +- .../server/src/wake/RedisWakeAgents.test.ts | 235 ++++++++++++++++++ .../server/src/wake/RedisWakeAgents.ts | 171 +++++++++++++ .../server/src/wake/SessionWaker.test.ts | 55 +++- mcp/packages/server/src/wake/SessionWaker.ts | 66 +++-- .../server/src/wake/WakeAgentRegistry.test.ts | 16 +- .../server/src/wake/WakeAgentRegistry.ts | 45 +++- mcp/packages/server/src/wake/WakeAgents.ts | 26 ++ 14 files changed, 1008 insertions(+), 57 deletions(-) create mode 100644 mcp/packages/server/src/wake/RedisWakeAgents.test.ts create mode 100644 mcp/packages/server/src/wake/RedisWakeAgents.ts create mode 100644 mcp/packages/server/src/wake/WakeAgents.ts diff --git a/mcp/packages/server/src/PluginBridge.sessions.test.ts b/mcp/packages/server/src/PluginBridge.sessions.test.ts index 20874c60fc..037be85ae6 100644 --- a/mcp/packages/server/src/PluginBridge.sessions.test.ts +++ b/mcp/packages/server/src/PluginBridge.sessions.test.ts @@ -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); diff --git a/mcp/packages/server/src/PluginBridge.ts b/mcp/packages/server/src/PluginBridge.ts index 4c6125ebf0..21ffa4619f 100644 --- a/mcp/packages/server/src/PluginBridge.ts +++ b/mcp/packages/server/src/PluginBridge.ts @@ -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, 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, 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.")); } diff --git a/mcp/packages/server/src/PluginTask.ts b/mcp/packages/server/src/PluginTask.ts index 30a88c73de..b01d6cf8c9 100644 --- a/mcp/packages/server/src/PluginTask.ts +++ b/mcp/packages/server/src/PluginTask.ts @@ -82,6 +82,12 @@ export abstract class AbstractPluginTask { + 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/ + ); +}); diff --git a/mcp/packages/server/src/RedisBridge.ts b/mcp/packages/server/src/RedisBridge.ts index b653c70faf..745e39d39e 100644 --- a/mcp/packages/server/src/RedisBridge.ts +++ b/mcp/packages/server/src/RedisBridge.ts @@ -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) => 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 { 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); + onResponse(message as PluginTaskResponse); } 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 { + 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 { + 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 { + 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 { + 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 { await this.subscriber.quit(); diff --git a/mcp/packages/server/src/RemotePluginTask.ts b/mcp/packages/server/src/RemotePluginTask.ts index 368623a5a0..c5580b5cde 100644 --- a/mcp/packages/server/src/RemotePluginTask.ts +++ b/mcp/packages/server/src/RemotePluginTask.ts @@ -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> { + /** 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; /** 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>(); @@ -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); diff --git a/mcp/packages/server/src/wake/RedisWakeAgents.test.ts b/mcp/packages/server/src/wake/RedisWakeAgents.test.ts new file mode 100644 index 0000000000..a3cfee8942 --- /dev/null +++ b/mcp/packages/server/src/wake/RedisWakeAgents.test.ts @@ -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 void>>(); + readonly responses = new Map 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); +}); diff --git a/mcp/packages/server/src/wake/RedisWakeAgents.ts b/mcp/packages/server/src/wake/RedisWakeAgents.ts new file mode 100644 index 0000000000..25f8d11275 --- /dev/null +++ b/mcp/packages/server/src/wake/RedisWakeAgents.ts @@ -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(); + private readonly subscriptionChanges = new Map>(); + 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 { + 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 { + 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 { + 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 { + const request: WakeRelayRequest = { id: randomUUID(), sessionId, timeoutMs }; + return new Promise((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()); + } +} diff --git a/mcp/packages/server/src/wake/SessionWaker.test.ts b/mcp/packages/server/src/wake/SessionWaker.test.ts index e026513c1d..9a57d6c9f1 100644 --- a/mcp/packages/server/src/wake/SessionWaker.test.ts +++ b/mcp/packages/server/src/wake/SessionWaker.test.ts @@ -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 }); diff --git a/mcp/packages/server/src/wake/SessionWaker.ts b/mcp/packages/server/src/wake/SessionWaker.ts index 0cf2521dd2..fb28a3f877 100644 --- a/mcp/packages/server/src/wake/SessionWaker.ts +++ b/mcp/packages/server/src/wake/SessionWaker.ts @@ -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; + ensureResponsive(connection: PenpotConnection, observer?: WakeObserver): Promise; } /** 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>(); /** - * @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 { + async ensureResponsive(connection: PenpotConnection, observer?: WakeObserver): Promise { 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 { 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."); diff --git a/mcp/packages/server/src/wake/WakeAgentRegistry.test.ts b/mcp/packages/server/src/wake/WakeAgentRegistry.test.ts index 82e13b4e99..ae46c6ced0 100644 --- a/mcp/packages/server/src/wake/WakeAgentRegistry.test.ts +++ b/mcp/packages/server/src/wake/WakeAgentRegistry.test.ts @@ -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 () => { diff --git a/mcp/packages/server/src/wake/WakeAgentRegistry.ts b/mcp/packages/server/src/wake/WakeAgentRegistry.ts index 8b8e03d79b..a40a1d3649 100644 --- a/mcp/packages/server/src/wake/WakeAgentRegistry.ts +++ b/mcp/packages/server/src/wake/WakeAgentRegistry.ts @@ -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; } +/** 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(); private readonly pendingWakes = new Map(); + 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 { 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((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; diff --git a/mcp/packages/server/src/wake/WakeAgents.ts b/mcp/packages/server/src/wake/WakeAgents.ts new file mode 100644 index 0000000000..0de22bed6c --- /dev/null +++ b/mcp/packages/server/src/wake/WakeAgents.ts @@ -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; +} + +/** 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"; + } +}