diff --git a/mcp/packages/server/src/PenpotSessionSchema.ts b/mcp/packages/server/src/PenpotSessionSchema.ts new file mode 100644 index 0000000000..b70da49330 --- /dev/null +++ b/mcp/packages/server/src/PenpotSessionSchema.ts @@ -0,0 +1,14 @@ +import type { PenpotSession } from "@penpot/mcp-common"; +import { z } from "zod"; + +/** + * Schema of a Penpot session as reported by the plugin and exchanged between MCP instances. + * + * A session without an open file (e.g. a plugin connected from the dashboard) has null `fileId` and `fileName`. + */ +export const penpotSessionSchema = z.object({ + sessionId: z.string().min(1).max(128), + fileId: z.string().min(1).max(128).nullable(), + fileName: z.string().nullable(), + workspaceState: z.enum(["none", "loading", "ready"]).optional(), +}) satisfies z.ZodType; diff --git a/mcp/packages/server/src/PluginBridge.ts b/mcp/packages/server/src/PluginBridge.ts index 340016f44f..8bf426e9a8 100644 --- a/mcp/packages/server/src/PluginBridge.ts +++ b/mcp/packages/server/src/PluginBridge.ts @@ -10,6 +10,7 @@ import { UnclaimedRollbacksMessage, } from "@penpot/mcp-common"; import { createLogger } from "./logger"; +import { penpotSessionSchema } from "./PenpotSessionSchema"; import type { PenpotMcpServer } from "./PenpotMcpServer"; import type { RedisBridge } from "./RedisBridge"; import type { PenpotConnection, PluginLivenessState } from "./PenpotConnection"; @@ -39,12 +40,7 @@ export const HEARTBEAT_STALE_THRESHOLD_MS = 30000; const connectionInitSchema = z.object({ type: z.literal("initialize"), - session: z.object({ - sessionId: z.string().min(1).max(128), - fileId: z.string().min(1).max(128).nullable(), - fileName: z.string().nullable(), - workspaceState: z.enum(["none", "loading", "ready"]).optional(), - }), + session: penpotSessionSchema, }); /** @@ -191,9 +187,7 @@ export class PluginBridge implements TaskDispatchHost { } connection.frozen = false; if (message?.type === "context-update") { - const context = connectionInitSchema.shape.session - .omit({ sessionId: true }) - .parse(message.context); + const context = penpotSessionSchema.omit({ sessionId: true }).parse(message.context); Object.assign(connection.session, context); return; } diff --git a/mcp/packages/server/src/RedisBridge.test.ts b/mcp/packages/server/src/RedisBridge.test.ts index 52080fcc3c..c64650ef1e 100644 --- a/mcp/packages/server/src/RedisBridge.test.ts +++ b/mcp/packages/server/src/RedisBridge.test.ts @@ -7,6 +7,8 @@ import { RedisBridge } from "./RedisBridge"; import { PluginBridge } from "./PluginBridge"; import { PluginTask } from "./PluginTask"; import type { PenpotMcpServer } from "./PenpotMcpServer"; +import type { PenpotSession } from "@penpot/mcp-common"; +import { InvalidDiscoveryResponseError } from "./dispatch/TaskDispatcher"; // Replace only the Redis command boundary; exercise the real transport, // dispatchers, task tracking, and WebSockets above it. @@ -105,8 +107,11 @@ test("Redis discovery receives one reply per subscribed instance and returns pub second.publishDiscoveryResponse(id, { instanceId: "second", sessions: [] }); }); const responders: string[] = []; - const count = await requester.sendDiscoveryRequest("alice", "query", (response) => - responders.push(response.instanceId) + const count = await requester.sendDiscoveryRequest( + "alice", + "query", + (response) => responders.push(response.instanceId), + assert.fail ); assert.equal(count, 2); assert.deepEqual(responders.sort(), ["first", "second"]); @@ -140,12 +145,12 @@ test("a failed discovery subscription is cleaned up and can be retried", async ( owner.subscribeToDiscovery("alice", () => {}), /Subscription failed/ ); - assert.equal(await caller.sendDiscoveryRequest("alice", "after-failure", () => {}), 0); + assert.equal(await caller.sendDiscoveryRequest("alice", "after-failure", () => {}, assert.fail), 0); await caller.unsubscribeFromDiscoveryResponses("after-failure"); await owner.subscribeToDiscovery("alice", ({ id }) => { owner.publishDiscoveryResponse(id, { instanceId: "owner", sessions: [] }); }); - assert.equal(await caller.sendDiscoveryRequest("alice", "retry", () => {}), 1); + assert.equal(await caller.sendDiscoveryRequest("alice", "retry", () => {}, assert.fail), 1); await caller.unsubscribeFromDiscoveryResponses("retry"); }); @@ -172,13 +177,14 @@ function instance(t: TestContext) { const connect = async ( sessionId: string, outcome: "success" | "error" | "silent" = "success", - reply: string = sessionId + reply: string = sessionId, + file: { fileId: string | null; fileName: string | null } = { fileId: "file", fileName: "Design" } ) => { const socket = new WebSocket(`ws://127.0.0.1:${wsPort}?userToken=alice`); sockets.push(socket); await once(socket, "open"); const initialized = once(socket, "message", { signal: AbortSignal.timeout(2000) }); - socket.send(JSON.stringify({ type: "initialize", session: { sessionId, fileId: "file", fileName: "Design" } })); + socket.send(JSON.stringify({ type: "initialize", session: { sessionId, ...file } })); await initialized; socket.on("message", (raw) => { const request = JSON.parse(raw.toString()); @@ -214,6 +220,34 @@ test("full Redis discovery implicitly selects a sole remote WebSocket", async (t assert.equal((await requester.bridge.executePluginTask(new PluginTask("test", {}))).data, "only"); }); +test("full Redis discovery implicitly selects a sole remote WebSocket without an open file", async (t) => { + const owner = instance(t); + const requester = instance(t); + await owner.connect("dashboard", "success", "dashboard", { fileId: null, fileName: null }); + assert.equal((await requester.bridge.executePluginTask(new PluginTask("test", {}))).data, "dashboard"); +}); + +test("an unreadable discovery reply fails discovery at once without suggesting a retry", async (t) => { + const responder = transport(t); + const requester = instance(t); + await responder.subscribeToDiscovery("alice", ({ id }) => { + responder.publishDiscoveryResponse(id, { + instanceId: "other-version", + sessions: [{ sessionId: "tab", fileId: 42 } as unknown as PenpotSession], + }); + }); + const started = Date.now(); + await assert.rejects(requester.bridge.executePluginTask(new PluginTask("test", {})), (error: Error) => { + assert.ok(error instanceof InvalidDiscoveryResponseError); + assert.match(error.message, /cannot be read/); + assert.match(error.message, /fileId/); + assert.match(error.message, /Retrying will not help/); + assert.doesNotMatch(error.message, /temporary server load/); + return true; + }); + assert.ok(Date.now() - started < 1000); +}); + test("forwarded plugin failures retain the regular execution failure route", async (t) => { const owner = instance(t); const requester = instance(t); diff --git a/mcp/packages/server/src/RedisBridge.ts b/mcp/packages/server/src/RedisBridge.ts index b653c70faf..ca3a45364d 100644 --- a/mcp/packages/server/src/RedisBridge.ts +++ b/mcp/packages/server/src/RedisBridge.ts @@ -1,6 +1,7 @@ import Redis from "ioredis"; import { PluginTaskRequest, PluginTaskResponse, PenpotSession } from "@penpot/mcp-common"; import { createLogger } from "./logger"; +import { penpotSessionSchema } from "./PenpotSessionSchema"; import { z } from "zod"; /** Query for the user's sessions held by each subscribed MCP instance. */ @@ -30,13 +31,7 @@ const sessionClaimSchema = z.object({ const discoveryResponseSchema = z.object({ instanceId: z.string().min(1), - sessions: z.array( - z.object({ - sessionId: z.string().min(1).max(128), - fileId: z.string().min(1).max(128), - fileName: z.string(), - }) - ), + sessions: z.array(penpotSessionSchema), }); /** @@ -248,19 +243,29 @@ export class RedisBridge { return `penpot.mcp.${this.tenant}.discovery.res.${this.channelKey(requestId)}`; } - /** Subscribes to replies before publishing; returns the actual recipient count. */ + /** + * Subscribes to replies before publishing; returns the actual recipient count. + * + * A reply that cannot be parsed is passed to `onInvalidResponse` with a readable description of the problem. + */ async sendDiscoveryRequest( userToken: string, requestId: string, - onResponse: (response: SessionDiscoveryResponse) => void + onResponse: (response: SessionDiscoveryResponse) => void, + onInvalidResponse: (error: Error) => void ): Promise { const channel = this.discoveryResponseChannel(requestId); this.handlers.set(channel, (raw) => { + let response: SessionDiscoveryResponse; try { - onResponse(discoveryResponseSchema.parse(JSON.parse(raw))); + response = discoveryResponseSchema.parse(JSON.parse(raw)); } catch (error) { this.logger.error(error, "Invalid discovery response"); + const reason = error instanceof z.ZodError ? z.prettifyError(error) : String(error); + onInvalidResponse(new Error(reason)); + return; } + onResponse(response); }); try { await this.subscriber.subscribe(channel); diff --git a/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.test.ts b/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.test.ts index 4c1b153916..18f3c657bf 100644 --- a/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.test.ts +++ b/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.test.ts @@ -6,14 +6,25 @@ import { PluginTask, type AbstractPluginTask } from "../PluginTask"; import type { RedisBridge, SessionClaim, SessionDiscoveryRequest, SessionDiscoveryResponse } from "../RedisBridge"; import { UserPenpotConnections } from "../UserPenpotConnections"; import { MultiInstanceTaskDispatcher } from "./MultiInstanceTaskDispatcher"; -import { SessionDiscoveryError, type TaskDispatchHost } from "./TaskDispatcher"; +import { InvalidDiscoveryResponseError, SessionDiscoveryError, type TaskDispatchHost } from "./TaskDispatcher"; + +/** Parse failure reported for a discovery reply that the requester cannot read. */ +const INVALID_REPLY_REASON = "sessions[0].fileId: expected string"; + +/** Reply callbacks registered by one discovery request. */ +interface DiscoveryReplyHandlers { + onResponse: (response: SessionDiscoveryResponse) => void; + onInvalidResponse: (error: Error) => void; +} /** In-memory Pub/Sub transport retaining Redis's per-subscriber delivery count. */ class RedisNetwork { readonly tasks = new Map void>>(); readonly discovery = new Map void>>(); readonly taskResponses = new Map) => void>(); - readonly discoveryResponses = new Map void>(); + readonly discoveryResponses = new Map(); + /** reply callbacks of the latest discovery request, kept after its subscription is released. */ + lastDiscoveryReplyHandlers?: DiscoveryReplyHandlers; readonly claims = new Map void>>(); failDiscovery = false; dropClaims = false; @@ -25,6 +36,7 @@ class RedisNetwork { class RedisPeer { replyToDiscovery = true; repeatDiscoveryReplies = false; + sendInvalidDiscoveryReplies = false; constructor(readonly network: RedisNetwork) {} @@ -72,9 +84,16 @@ class RedisPeer { this.network.discovery.get(token)?.delete(this); } - async sendDiscoveryRequest(token: string, id: string, onResponse: (response: SessionDiscoveryResponse) => void) { + async sendDiscoveryRequest( + token: string, + id: string, + onResponse: (response: SessionDiscoveryResponse) => void, + onInvalidResponse: (error: Error) => void + ) { if (this.network.failDiscovery) throw new Error("Redis unavailable for discovery"); - this.network.discoveryResponses.set(id, onResponse); + const handlers = { onResponse, onInvalidResponse }; + this.network.discoveryResponses.set(id, handlers); + this.network.lastDiscoveryReplyHandlers = handlers; await this.network.beforeDiscoveryPublish?.(); const subscribers = this.network.discovery.get(token) ?? new Map(); for (const [peer, handler] of subscribers) { @@ -87,7 +106,12 @@ class RedisPeer { } publishDiscoveryResponse(id: string, response: SessionDiscoveryResponse) { - this.network.discoveryResponses.get(id)?.(response); + const handlers = this.network.discoveryResponses.get(id); + if (this.sendInvalidDiscoveryReplies) { + handlers?.onInvalidResponse(new Error(INVALID_REPLY_REASON)); + } else { + handlers?.onResponse(response); + } } async unsubscribeFromDiscoveryResponses(id: string) { @@ -244,6 +268,7 @@ test("incomplete discovery never dispatches even when it yields exactly one sess slow.redis.replyToDiscovery = false; await assert.rejects(new Instance(network).execute("alice"), (error: Error) => { assert.ok(error instanceof SessionDiscoveryError); + assert.match(error.message, /not every MCP server instance replied within 20ms/); assert.match(error.message, /temporary server load.*retry/); return true; }); @@ -265,7 +290,65 @@ test("duplicate replies from one instance do not substitute for a missing reply" test("discovery transport failures use the retryable discovery error", async () => { const network = new RedisNetwork(); network.failDiscovery = true; - await assert.rejects(new Instance(network).execute("alice"), SessionDiscoveryError); + await assert.rejects(new Instance(network).execute("alice"), (error: Error) => { + assert.ok(error instanceof SessionDiscoveryError); + assert.match(error.message, /the Redis request failed \(Redis unavailable for discovery\)/); + assert.match(error.message, /temporary server load.*retry/); + return true; + }); +}); + +test("an invalid reply fails discovery without dispatching even when another instance replies validly", async () => { + const network = new RedisNetwork(); + const responsive = new Instance(network); + const otherVersion = new Instance(network); + await responsive.connect("alice", "visible"); + await otherVersion.connect("alice", "unreadable"); + otherVersion.redis.sendInvalidDiscoveryReplies = true; + await assert.rejects(new Instance(network).execute("alice"), (error: Error) => { + assert.ok(error instanceof InvalidDiscoveryResponseError); + assert.ok(error.message.includes(INVALID_REPLY_REASON)); + assert.match(error.message, /Retrying will not help/); + return true; + }); + assert.deepEqual(responsive.executions, []); + assert.deepEqual(otherVersion.executions, []); + assert.equal(network.discoveryResponses.size, 0); +}); + +test("an invalid reply after discovery settled leaves its outcome unchanged", async () => { + const network = new RedisNetwork(); + const owner = new Instance(network); + await owner.connect("alice", "only"); + const requester = new Instance(network); + const sessions = requester.dispatcher.discoverSessions("alice"); + const late = network.lastDiscoveryReplyHandlers!; + assert.deepEqual( + (await sessions).map((session) => session.sessionId), + ["only"] + ); + late.onInvalidResponse(new Error(INVALID_REPLY_REASON)); + assert.equal((await requester.execute("alice")).data, "only"); +}); + +test("closing the dispatcher fails pending discoveries without blaming server load", async () => { + const network = new RedisNetwork(); + const silent = new Instance(network); + await silent.connect("alice", "tab"); + silent.redis.replyToDiscovery = false; + const requester = new Instance(network); + const rejected = assert.rejects(requester.execute("alice"), (error: Error) => { + assert.ok(error instanceof SessionDiscoveryError); + assert.match(error.message, /the MCP server is shutting down/); + assert.match(error.message, /You can retry the request/); + assert.doesNotMatch(error.message, /temporary server load/); + return true; + }); + await new Promise((resolve) => setImmediate(resolve)); + await requester.dispatcher.close(); + await rejected; + assert.deepEqual(silent.executions, []); + assert.equal(network.discoveryResponses.size, 0); }); test("disconnecting one tab retains discovery and dispatch for the remaining tab", async () => { diff --git a/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.ts b/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.ts index 10c895e659..891bb12d77 100644 --- a/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.ts +++ b/mcp/packages/server/src/dispatch/MultiInstanceTaskDispatcher.ts @@ -5,7 +5,12 @@ import type { AbstractPluginTask } from "../PluginTask"; import { RemotePluginTask } from "../RemotePluginTask"; import type { RedisBridge, SessionClaim, SessionDiscoveryResponse } from "../RedisBridge"; import { createLogger } from "../logger"; -import { SessionDiscoveryError, TaskDispatcher, type TaskDispatchHost } from "./TaskDispatcher"; +import { + InvalidDiscoveryResponseError, + SessionDiscoveryError, + TaskDispatcher, + type TaskDispatchHost, +} from "./TaskDispatcher"; /** Routes tasks over Redis and discovers live sessions across MCP instances. */ export class MultiInstanceTaskDispatcher extends TaskDispatcher { @@ -44,10 +49,10 @@ export class MultiInstanceTaskDispatcher extends TaskDispatcher { resolve = res; reject = rej; }); - const fail = () => { + const fail = (error: SessionDiscoveryError) => { if (!settled) { settled = true; - reject(new SessionDiscoveryError()); + reject(error); } }; const complete = () => { @@ -56,8 +61,18 @@ export class MultiInstanceTaskDispatcher extends TaskDispatcher { resolve(this.uniqueSessions([...responses.values()].flat())); } }; - const timeout = setTimeout(fail, this.discoveryTimeoutMs); - this.pendingDiscoveries.set(requestId, fail); + const timeout = setTimeout( + () => + fail( + new SessionDiscoveryError( + `not every MCP server instance replied within ${this.discoveryTimeoutMs}ms` + ) + ), + this.discoveryTimeoutMs + ); + this.pendingDiscoveries.set(requestId, () => + fail(new SessionDiscoveryError("the MCP server is shutting down", "You can retry the request.")) + ); const releaseResponses = () => { void this.redisBridge .unsubscribeFromDiscoveryResponses(requestId) @@ -65,18 +80,29 @@ export class MultiInstanceTaskDispatcher extends TaskDispatcher { }; // Replies can arrive before publish returns its recipient count. void this.redisBridge - .sendDiscoveryRequest(token, requestId, (response: SessionDiscoveryResponse) => { - if (settled) return; - responses.set(response.instanceId, response.sessions); - complete(); - }) + .sendDiscoveryRequest( + token, + requestId, + (response: SessionDiscoveryResponse) => { + if (settled) return; + responses.set(response.instanceId, response.sessions); + complete(); + }, + (error) => fail(new InvalidDiscoveryResponseError(error)) + ) .then((count) => { expectedResponses = count; complete(); // A slow subscribe/publish may finish after the discovery timeout. if (settled) releaseResponses(); }) - .catch(fail); + .catch((error) => + fail( + new SessionDiscoveryError( + `the Redis request failed (${error instanceof Error ? error.message : error})` + ) + ) + ); try { return await result; } finally { diff --git a/mcp/packages/server/src/dispatch/TaskDispatcher.ts b/mcp/packages/server/src/dispatch/TaskDispatcher.ts index 062af6b80c..b9ea7b617a 100644 --- a/mcp/packages/server/src/dispatch/TaskDispatcher.ts +++ b/mcp/packages/server/src/dispatch/TaskDispatcher.ts @@ -40,10 +40,33 @@ export abstract class TaskDispatcher { } } -/** Error returned when discovery cannot establish the complete set of sessions. */ +/** + * Error returned when discovery cannot establish the complete set of sessions. + * + * The default message marks the failure as transient; subclasses describe failures that a retry does not fix. + */ export class SessionDiscoveryError extends Error { - constructor() { - super("Penpot session discovery failed. This may be due to temporary server load. You can retry the request."); + /** + * @param reason - what prevented discovery from completing + * @param advice - what the caller can do about it + */ + constructor(reason: string, advice = "This may be due to temporary server load. You can retry the request.") { + super(`Penpot session discovery failed: ${reason}. ${advice}`); this.name = "SessionDiscoveryError"; } } + +/** Error returned when an MCP server instance replies to discovery with a session list that cannot be read. */ +export class InvalidDiscoveryResponseError extends SessionDiscoveryError { + /** + * @param cause - the parse failure of the reply + */ + constructor(cause: Error) { + super( + `an MCP server instance sent a session list that cannot be read (${cause.message})`, + "Retrying will not help; the MCP server instances may be running different versions. " + + "To work around it, ask the user for the session ID shown in the Penpot MCP menu and pass it as sessionId." + ); + this.name = "InvalidDiscoveryResponseError"; + } +}