From e9588fea49011bc29a0beacadbbac9426e62b3f3 Mon Sep 17 00:00:00 2001 From: "alonso.torres" Date: Mon, 26 Oct 2026 10:24:52 +0000 Subject: [PATCH] :sparkles: Accept browser extensions as MCP wake agents A browser freezes background tabs, and a frozen Penpot tab cannot run the MCP plugin. A browser extension can wake such a tab, so let it connect to the MCP server as a wake agent. Define the messages the extension exchanges with the Penpot page and with the server: the page's session announcement, and the agent's initialization, session list, heartbeat and wake result, and the server's wake request. The server checks the agent messages with zod schemas. The plugin WebSocket endpoint accepts agents: a socket whose first message is initialize-wake-agent goes to a wake agent registry, which then handles its messages, keep-alive pings and close. It never becomes a plugin connection, so it is not listed as a Penpot session and does not displace the plugin connection of the same session. Agents go through the same user token checks as plugins. The registry keeps the sessions each agent can wake, sends wake requests and matches their results by request ID. A request fails when no agent holds the session, the agent disconnects or it does not answer in time. An agent may only wake sessions it registered under the same token, and a reconnecting agent replaces its previous connection. --- mcp/packages/common/src/types.ts | 99 +++++++++ .../server/src/PluginBridge.sessions.test.ts | 65 ++++++ mcp/packages/server/src/PluginBridge.ts | 22 ++ .../server/src/wake/WakeAgentProtocol.test.ts | 49 +++++ .../server/src/wake/WakeAgentProtocol.ts | 38 ++++ .../server/src/wake/WakeAgentRegistry.test.ts | 145 +++++++++++++ .../server/src/wake/WakeAgentRegistry.ts | 201 ++++++++++++++++++ 7 files changed, 619 insertions(+) create mode 100644 mcp/packages/server/src/wake/WakeAgentProtocol.test.ts create mode 100644 mcp/packages/server/src/wake/WakeAgentProtocol.ts create mode 100644 mcp/packages/server/src/wake/WakeAgentRegistry.test.ts create mode 100644 mcp/packages/server/src/wake/WakeAgentRegistry.ts diff --git a/mcp/packages/common/src/types.ts b/mcp/packages/common/src/types.ts index f9d7cde400..cfb5dbba79 100644 --- a/mcp/packages/common/src/types.ts +++ b/mcp/packages/common/src/types.ts @@ -97,3 +97,102 @@ export interface ExecuteCodeTaskResultData { */ log: string; } + +/** + * Status of the Penpot MCP plugin connection, as reported to the Penpot app. + */ +export type McpConnectionStatus = "idle" | "connecting" | "connected" | "disconnected" | "error"; + +/** + * Message posted by a browser extension's content script to ask the Penpot page for its MCP session. + */ +export interface ExtensionHello { + type: "penpot-mcp-extension:hello"; +} + +/** + * Message posted by the Penpot page when it starts listening for an {@link ExtensionHello}. + * + * A content script that said hello before the page was listening says it again. + */ +export interface PageReady { + type: "penpot-mcp:ready"; +} + +/** + * Message posted by the Penpot page to announce its MCP session to a browser extension. + * + * The page sends it in reply to {@link ExtensionHello} and on every connection status change. + */ +export interface SessionAnnouncement { + type: "penpot-mcp:session"; + /** the MCP session ID of the page, or null while it is not known. */ + sessionId: string | null; + /** the URL of the MCP server's WebSocket endpoint, with a ws(s) or http(s) scheme. */ + serverUrl: string; + /** the user's MCP access token, or null when the server needs none. */ + token: string | null; + status: McpConnectionStatus; +} + +/** + * Outcome of a wake request handled by a browser extension. + * + * - `activated`: the extension woke the tab holding the session, in place or by bringing it forward + * - `unknown-session`: the extension holds no tab for the session + * - `tab-closed`: the tab holding the session no longer exists + * - `disabled`: the user has paused waking in the extension + */ +export type WakeOutcome = "activated" | "unknown-session" | "tab-closed" | "disabled"; + +/** + * First message sent by a wake agent (a browser extension that can wake Penpot tabs) on its + * WebSocket connection to the MCP server. + */ +export interface WakeAgentInit { + type: "initialize-wake-agent"; + /** stable identifier of the agent; a reconnecting agent replaces its previous connection. */ + agentId: string; + /** the sessions the agent can wake. */ + sessionIds: string[]; +} + +/** + * Message replacing the full list of sessions a wake agent can wake. + */ +export interface WakeAgentSessions { + type: "wake-sessions"; + sessionIds: string[]; +} + +/** + * Message sent periodically by a wake agent to keep its connection active. + */ +export interface WakeAgentHeartbeat { + type: "heartbeat"; +} + +/** + * Request sent by the MCP server asking a wake agent to wake the tab holding a session. + */ +export interface WakeRequest { + type: "wake"; + requestId: string; + sessionId: string; +} + +/** + * Reply of a wake agent to a {@link WakeRequest}. + */ +export interface WakeResult { + type: "wake-result"; + requestId: string; + outcome: WakeOutcome; + /** what the agent did, for the server log. */ + details?: string; +} + +/** + * Message sent by a wake agent to the MCP server. + */ +export type WakeAgentMessage = WakeAgentInit | WakeAgentSessions | WakeAgentHeartbeat | WakeResult; diff --git a/mcp/packages/server/src/PluginBridge.sessions.test.ts b/mcp/packages/server/src/PluginBridge.sessions.test.ts index 1a42241169..1046fdb062 100644 --- a/mcp/packages/server/src/PluginBridge.sessions.test.ts +++ b/mcp/packages/server/src/PluginBridge.sessions.test.ts @@ -169,3 +169,68 @@ test("a dashboard session remains routable while its workspace metadata changes" assert.deepEqual(bridge.getUserConnections("alice")!.getSessions(), [{ sessionId: "tab", ...context }]); } }); + +async function connectAgent(sessionIds: string[], token = "alice", agentId = "agent") { + const socket = new WebSocket(`ws://127.0.0.1:${port}?userToken=${token}`); + sockets.push(socket); + await once(socket, "open"); + socket.send(JSON.stringify({ type: "initialize-wake-agent", agentId, sessionIds })); + // the agent is registered once a message sent after its initialization has been processed + while (!bridge.wakeAgents.canWake(multiUser ? token : null, sessionIds[0])) { + await new Promise((resolve) => setImmediate(resolve)); + } + return socket; +} + +/** Replies to the next wake request on an agent socket with the given outcome. */ +function answerWake(agent: WebSocket, outcome: string) { + agent.once("message", (raw) => { + const request = JSON.parse(raw.toString()); + agent.send(JSON.stringify({ type: "wake-result", requestId: request.requestId, outcome })); + }); +} + +test("a wake agent connection is not listed as a Penpot session", async () => { + await connectAgent(["tab"]); + await assert.rejects(bridge.executePluginTask(new PluginTask("test", {})), /No Penpot sessions/); +}); + +test("a wake agent answers wake requests over its connection", async () => { + const agent = await connectAgent(["tab"]); + answerWake(agent, "activated"); + assert.equal(await bridge.wakeAgents.wake("alice", "tab", 1000), "activated"); +}); + +test("a wake agent can only wake sessions under the token it connected with", async () => { + await connectAgent(["tab"], "bob"); + assert.equal(bridge.wakeAgents.canWake("alice", "tab"), false); +}); + +test("a plugin and a wake agent for the same session do not displace each other", async () => { + const agent = await connectAgent(["tab"]); + await connect("tab"); + assert.equal(agent.readyState, WebSocket.OPEN); + assert.equal((await bridge.executePluginTask(new PluginTask("test", {}), "tab")).data, "tab"); + assert.equal(bridge.wakeAgents.canWake("alice", "tab"), true); +}); + +test("closing a wake agent connection unregisters it", async () => { + const agent = await connectAgent(["tab"]); + agent.close(); + while (bridge.wakeAgents.canWake("alice", "tab")) { + await new Promise((resolve) => setImmediate(resolve)); + } +}); + +test("single-user wake agents are registered without a token", async () => { + multiUser = false; + await connectAgent(["tab"], "ignored"); + assert.equal(bridge.wakeAgents.canWake(null, "tab"), true); +}); + +test("a wake agent connection without a token is rejected in multi-user mode", async () => { + const socket = new WebSocket(`ws://127.0.0.1:${port}`); + sockets.push(socket); + const [code] = await once(socket, "close"); + assert.equal(code, 1008); +}); diff --git a/mcp/packages/server/src/PluginBridge.ts b/mcp/packages/server/src/PluginBridge.ts index 913b28b26a..b4b8e084b3 100644 --- a/mcp/packages/server/src/PluginBridge.ts +++ b/mcp/packages/server/src/PluginBridge.ts @@ -11,6 +11,8 @@ import { UserPenpotConnections } from "./UserPenpotConnections"; import type { TaskDispatcher, TaskDispatchHost } from "./dispatch/TaskDispatcher"; import { SingleInstanceTaskDispatcher } from "./dispatch/SingleInstanceTaskDispatcher"; import { MultiInstanceTaskDispatcher } from "./dispatch/MultiInstanceTaskDispatcher"; +import { WakeAgentRegistry } from "./wake/WakeAgentRegistry"; +import { wakeAgentInitSchema } from "./wake/WakeAgentProtocol"; const KEEP_ALIVE_TIME = 30000; // 30 seconds const INITIALIZATION_TIMEOUT_MS = 10_000; @@ -78,6 +80,8 @@ export class PluginBridge implements TaskDispatchHost { private readonly taskTimeouts: Map = new Map(); private readonly taskConnections: Map = new Map(); private readonly dispatcher: TaskDispatcher; + /** the browser extensions that can wake the tabs of frozen plugin sessions. */ + public readonly wakeAgents = new WakeAgentRegistry(KEEP_ALIVE_TIME); /** * Creates the plugin bridge and starts its WebSocket server. @@ -138,11 +142,26 @@ export class PluginBridge implements TaskDispatchHost { this.logger.debug("Received WebSocket message: %s", data.toString()); try { const message = JSON.parse(data.toString()); + if (this.wakeAgents.has(ws)) { + this.wakeAgents.handleMessage(ws, message); + return; + } const connection = this.connectionsBySocket.get(ws); if (!connection) { if (initializing) { return; } + // a browser extension connects as a wake agent rather than as a plugin + const agentInitialization = wakeAgentInitSchema.safeParse(message); + if (agentInitialization.success) { + clearTimeout(initializationTimeout); + this.wakeAgents.add( + ws, + this.mcpServer.isMultiUserMode() ? userToken : null, + agentInitialization.data + ); + return; + } const initialization = connectionInitSchema.safeParse(message); if (!initialization.success) { ws.close(1008, "Expected connection metadata; please update the Penpot MCP plugin."); @@ -190,12 +209,14 @@ export class PluginBridge implements TaskDispatchHost { ws.on("close", () => { this.logger.info("WebSocket connection closed"); clearTimeout(initializationTimeout); + this.wakeAgents.remove(ws); this.removeConnection(ws); }); ws.on("error", (error) => { this.logger.error(error, "WebSocket connection error"); clearTimeout(initializationTimeout); + this.wakeAgents.remove(ws); this.removeConnection(ws); ws.terminate(); }); @@ -478,6 +499,7 @@ export class PluginBridge implements TaskDispatchHost { this.removeConnection(socket); socket.terminate(); } + this.wakeAgents.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/wake/WakeAgentProtocol.test.ts b/mcp/packages/server/src/wake/WakeAgentProtocol.test.ts new file mode 100644 index 0000000000..935fd7849f --- /dev/null +++ b/mcp/packages/server/src/wake/WakeAgentProtocol.test.ts @@ -0,0 +1,49 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import type { z } from "zod"; +import type { WakeAgentInit, WakeAgentMessage } from "@penpot/mcp-common"; +import { MAX_WAKE_AGENT_SESSIONS, wakeAgentInitSchema, wakeAgentMessageSchema } from "./WakeAgentProtocol"; + +// parsed messages must satisfy the shared protocol types +const _init: WakeAgentInit = {} as z.infer; +const _message: WakeAgentMessage = {} as z.infer; +void _init; +void _message; + +test("accepts an agent initialization with the sessions it can wake", () => { + const init = { type: "initialize-wake-agent", agentId: "agent", sessionIds: ["tab-1", "tab-2"] }; + assert.deepEqual(wakeAgentInitSchema.parse(init), init); +}); + +test("rejects a plugin initialization as an agent initialization", () => { + const init = { type: "initialize", session: { sessionId: "tab", fileId: null, fileName: null } }; + assert.equal(wakeAgentInitSchema.safeParse(init).success, false); +}); + +test("rejects an agent initialization with too many sessions", () => { + const sessionIds = Array.from({ length: MAX_WAKE_AGENT_SESSIONS + 1 }, (_, i) => `tab-${i}`); + const init = { type: "initialize-wake-agent", agentId: "agent", sessionIds }; + assert.equal(wakeAgentInitSchema.safeParse(init).success, false); +}); + +test("accepts session updates, heartbeats and wake results", () => { + for (const message of [ + { type: "wake-sessions", sessionIds: [] }, + { type: "heartbeat" }, + { type: "wake-result", requestId: "request", outcome: "tab-closed" }, + { type: "wake-result", requestId: "request", outcome: "activated", details: "Brought the tab forward." }, + ]) { + assert.deepEqual(wakeAgentMessageSchema.parse(message), message); + } +}); + +test("cuts long wake result details", () => { + const result = { type: "wake-result", requestId: "request", outcome: "activated", details: "x".repeat(2000) }; + const parsed = wakeAgentMessageSchema.parse(result); + assert.equal(parsed.type === "wake-result" && parsed.details?.length, 500); +}); + +test("rejects a wake result with an unknown outcome", () => { + const result = { type: "wake-result", requestId: "request", outcome: "maybe" }; + assert.equal(wakeAgentMessageSchema.safeParse(result).success, false); +}); diff --git a/mcp/packages/server/src/wake/WakeAgentProtocol.ts b/mcp/packages/server/src/wake/WakeAgentProtocol.ts new file mode 100644 index 0000000000..a5df61c86b --- /dev/null +++ b/mcp/packages/server/src/wake/WakeAgentProtocol.ts @@ -0,0 +1,38 @@ +import { z } from "zod"; + +/** the maximum number of sessions a single wake agent may register. */ +export const MAX_WAKE_AGENT_SESSIONS = 256; + +/** the length to which the server cuts the details of a wake result. */ +const MAX_WAKE_DETAILS_LENGTH = 500; + +const sessionIdSchema = z.string().min(1).max(128); + +const sessionIdsSchema = z.array(sessionIdSchema).max(MAX_WAKE_AGENT_SESSIONS); + +/** Schema of the first message of a wake agent connection. */ +export const wakeAgentInitSchema = z.object({ + type: z.literal("initialize-wake-agent"), + agentId: z.string().min(1).max(128), + sessionIds: sessionIdsSchema, +}); + +/** Schema of the messages a wake agent sends after its initialization. */ +export const wakeAgentMessageSchema = z.discriminatedUnion("type", [ + z.object({ + type: z.literal("wake-sessions"), + sessionIds: sessionIdsSchema, + }), + z.object({ + type: z.literal("heartbeat"), + }), + z.object({ + type: z.literal("wake-result"), + requestId: z.string().min(1).max(128), + outcome: z.enum(["activated", "unknown-session", "tab-closed", "disabled"]), + details: z + .string() + .transform((details) => details.slice(0, MAX_WAKE_DETAILS_LENGTH)) + .optional(), + }), +]); diff --git a/mcp/packages/server/src/wake/WakeAgentRegistry.test.ts b/mcp/packages/server/src/wake/WakeAgentRegistry.test.ts new file mode 100644 index 0000000000..82e13b4e99 --- /dev/null +++ b/mcp/packages/server/src/wake/WakeAgentRegistry.test.ts @@ -0,0 +1,145 @@ +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"; + +/** In-memory agent connection recording what the registry sends. */ +class FakeSocket { + readonly sent: WakeRequest[] = []; + closeCode?: number; + pings = 0; + + send(data: string): void { + this.sent.push(JSON.parse(data)); + } + + close(code?: number): void { + this.closeCode = code; + } + + ping(): void { + this.pings++; + } + + asSocket(): WakeAgentSocket { + return this as unknown as WakeAgentSocket; + } +} + +let registry: WakeAgentRegistry; + +beforeEach(() => { + registry = new WakeAgentRegistry(10); +}); + +afterEach(() => { + registry.close(); +}); + +function addAgent(token: string | null, sessionIds: string[], agentId = "agent"): FakeSocket { + const socket = new FakeSocket(); + registry.add(socket.asSocket(), token, { type: "initialize-wake-agent", agentId, sessionIds }); + return socket; +} + +function reply(socket: FakeSocket, outcome: string, requestId = socket.sent.at(-1)!.requestId): void { + registry.handleMessage(socket.asSocket(), { type: "wake-result", requestId, outcome }); +} + +test("sends a wake request to the agent holding the session and returns its outcome", async () => { + const socket = addAgent("alice", ["tab"]); + const outcome = registry.wake("alice", "tab", 1000); + assert.equal(socket.sent.length, 1); + assert.equal(socket.sent[0].type, "wake"); + assert.equal(socket.sent[0].sessionId, "tab"); + reply(socket, "activated"); + assert.equal(await outcome, "activated"); +}); + +test("passes on outcomes other than activation", async () => { + const socket = addAgent("alice", ["tab"]); + const outcome = registry.wake("alice", "tab", 1000); + reply(socket, "disabled"); + assert.equal(await outcome, "disabled"); +}); + +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/); +}); + +test("single-user agents wake sessions without a token", async () => { + const socket = addAgent(null, ["tab"]); + assert.equal(registry.canWake(null, "tab"), true); + assert.equal(registry.canWake("alice", "tab"), false); + const outcome = registry.wake(null, "tab", 1000); + reply(socket, "activated"); + assert.equal(await outcome, "activated"); +}); + +test("a session list update replaces the sessions an agent can wake", () => { + const socket = addAgent("alice", ["old"]); + registry.handleMessage(socket.asSocket(), { type: "wake-sessions", sessionIds: ["new"] }); + assert.equal(registry.canWake("alice", "old"), false); + assert.equal(registry.canWake("alice", "new"), true); +}); + +test("fails a wake request the agent does not answer in time", async () => { + addAgent("alice", ["tab"]); + await assert.rejects(registry.wake("alice", "tab", 20), /did not answer within 20ms/); +}); + +test("fails pending wake requests when the agent disconnects", async () => { + const socket = addAgent("alice", ["tab"]); + const outcome = registry.wake("alice", "tab", 1000); + registry.remove(socket.asSocket()); + await assert.rejects(outcome, /disconnected before answering/); + assert.equal(registry.canWake("alice", "tab"), false); +}); + +test("ignores a result sent by another agent for the request", async () => { + const holder = addAgent("alice", ["tab"], "holder"); + const other = addAgent("alice", [], "other"); + const outcome = registry.wake("alice", "tab", 1000); + reply(other, "activated", holder.sent[0].requestId); + reply(holder, "tab-closed"); + assert.equal(await outcome, "tab-closed"); +}); + +test("ignores invalid messages and results for unknown requests", () => { + const socket = addAgent("alice", ["tab"]); + registry.handleMessage(socket.asSocket(), { type: "wake-sessions", sessionIds: "tab" }); + registry.handleMessage(socket.asSocket(), { type: "wake-result", requestId: "unknown", outcome: "activated" }); + registry.handleMessage(socket.asSocket(), { type: "heartbeat" }); + assert.equal(registry.canWake("alice", "tab"), true); +}); + +test("a reconnecting agent replaces its previous connection", async () => { + const previous = addAgent("alice", ["tab"]); + const current = addAgent("alice", ["tab"]); + assert.equal(previous.closeCode, 1000); + assert.equal(registry.has(previous.asSocket()), false); + const outcome = registry.wake("alice", "tab", 1000); + assert.equal(previous.sent.length, 0); + reply(current, "activated"); + assert.equal(await outcome, "activated"); +}); + +test("the same agent ID under another user token is a separate agent", () => { + const bob = addAgent("bob", ["bob-tab"]); + addAgent("alice", ["alice-tab"]); + assert.equal(bob.closeCode, undefined); + assert.equal(registry.canWake("bob", "bob-tab"), true); +}); + +test("pings each agent to keep its connection alive", async () => { + const socket = addAgent("alice", ["tab"]); + while (socket.pings < 2) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + registry.remove(socket.asSocket()); + const pings = socket.pings; + await new Promise((resolve) => setTimeout(resolve, 50)); + assert.equal(socket.pings, pings); +}); diff --git a/mcp/packages/server/src/wake/WakeAgentRegistry.ts b/mcp/packages/server/src/wake/WakeAgentRegistry.ts new file mode 100644 index 0000000000..8b8e03d79b --- /dev/null +++ b/mcp/packages/server/src/wake/WakeAgentRegistry.ts @@ -0,0 +1,201 @@ +import { randomUUID } from "node:crypto"; +import type { WebSocket } from "ws"; +import type { WakeAgentInit, WakeOutcome, WakeRequest } from "@penpot/mcp-common"; +import { createLogger } from "../logger"; +import { wakeAgentMessageSchema } from "./WakeAgentProtocol"; + +const KEEP_ALIVE_TIME_MS = 30_000; + +/** Socket operations the registry uses on a wake agent connection. */ +export type WakeAgentSocket = Pick; + +/** One connected wake agent and the sessions it can wake. */ +interface WakeAgent { + readonly socket: WakeAgentSocket; + readonly userToken: string | null; + readonly agentId: string; + readonly pingInterval: NodeJS.Timeout; + sessionIds: Set; +} + +/** A wake request awaiting the agent's result. */ +interface PendingWake { + readonly agent: WakeAgent; + readonly resolve: (outcome: WakeOutcome) => void; + readonly reject: (error: Error) => void; + readonly timeout: NodeJS.Timeout; +} + +/** + * Registry of connected wake agents, i.e. browser extensions that can wake the tabs holding + * Penpot sessions. + * + * 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 { + private readonly logger = createLogger("WakeAgentRegistry"); + private readonly agentsBySocket = new Map(); + private readonly pendingWakes = new Map(); + + /** + * @param keepAliveMs - Interval, in ms, of the protocol pings sent to each agent + */ + constructor(private readonly keepAliveMs: number = KEEP_ALIVE_TIME_MS) {} + + /** Whether the socket belongs to a registered wake agent. */ + has(socket: WakeAgentSocket): boolean { + return this.agentsBySocket.has(socket); + } + + /** + * Registers an initialized wake agent connection. + * + * An earlier connection of the same agent for the same user is closed, since a reconnecting + * agent replaces it. + * + * @param socket - The agent's connection + * @param userToken - The user token the agent connected with, or null in single-user mode + * @param init - The agent's initialization message + */ + add(socket: WakeAgentSocket, userToken: string | null, init: WakeAgentInit): void { + for (const other of [...this.agentsBySocket.values()]) { + if (other.userToken === userToken && other.agentId === init.agentId) { + this.logger.info("Closing wake agent connection replaced by a newer one"); + this.remove(other.socket); + other.socket.close(1000, "Replaced by a newer connection of this wake agent."); + } + } + const pingInterval = setInterval(() => socket.ping(), this.keepAliveMs); + this.agentsBySocket.set(socket, { + socket, + userToken, + agentId: init.agentId, + pingInterval, + sessionIds: new Set(init.sessionIds), + }); + this.logger.info("Wake agent registered with %d session(s)", init.sessionIds.length); + } + + /** + * Handles a message received from a registered agent. + * + * @param socket - The agent's connection + * @param message - The parsed JSON message + */ + handleMessage(socket: WakeAgentSocket, message: unknown): void { + const agent = this.agentsBySocket.get(socket); + if (!agent) { + return; + } + const parsed = wakeAgentMessageSchema.safeParse(message); + if (!parsed.success) { + this.logger.warn("Ignoring invalid wake agent message"); + return; + } + const data = parsed.data; + switch (data.type) { + case "wake-sessions": + agent.sessionIds = new Set(data.sessionIds); + return; + case "heartbeat": + return; + case "wake-result": { + const pending = this.pendingWakes.get(data.requestId); + if (pending?.agent !== agent) { + return; + } + this.settle(data.requestId); + this.logger.info( + "Wake request %s answered: %s (%s)", + data.requestId, + data.outcome, + data.details ?? "no details" + ); + pending.resolve(data.outcome); + return; + } + } + } + + /** + * Removes an agent connection, failing its pending wake requests. Safe to call with a socket + * that is not (or no longer) registered. + * + * @param socket - The agent's connection + */ + remove(socket: WakeAgentSocket): void { + const agent = this.agentsBySocket.get(socket); + if (!agent) { + return; + } + clearInterval(agent.pingInterval); + this.agentsBySocket.delete(socket); + for (const [requestId, pending] of [...this.pendingWakes]) { + if (pending.agent === agent) { + this.settle(requestId); + pending.reject(new Error("The wake agent disconnected before answering.")); + } + } + this.logger.info("Wake agent removed"); + } + + /** Whether a connected agent of the user can wake the session. */ + canWake(userToken: string | null, sessionId: string): boolean { + 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 + */ + 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.")); + } + const requestId = randomUUID(); + return new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + this.settle(requestId); + reject(new Error(`The wake agent did not answer within ${timeoutMs}ms.`)); + }, timeoutMs); + this.pendingWakes.set(requestId, { agent, resolve, reject, timeout }); + const request: WakeRequest = { type: "wake", requestId, sessionId }; + agent.socket.send(JSON.stringify(request)); + this.logger.info("Sent wake request %s", requestId); + }); + } + + /** Removes all agents and fails their pending wake requests. */ + close(): void { + for (const socket of [...this.agentsBySocket.keys()]) { + this.remove(socket); + } + } + + /** 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; + for (const agent of this.agentsBySocket.values()) { + if (agent.userToken === userToken && agent.sessionIds.has(sessionId)) { + found = agent; + } + } + return found; + } + + /** Releases the correlation state of a pending wake request. */ + private settle(requestId: string): void { + const pending = this.pendingWakes.get(requestId); + if (pending) { + clearTimeout(pending.timeout); + this.pendingWakes.delete(requestId); + } + } +}