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