🐛 Fix MCP session discovery for plugins connected from the dashboard (#12865)

* 🐛 Fix MCP session discovery for plugins connected from the dashboard

With Redis-backed routing, execute_code without a sessionId timed out
with SessionDiscoveryError when the plugin ran on the dashboard. The
discovery reply schema required a string fileId and fileName, so replies
listing a session with no open file were dropped, and discovery waited
for its timeout.

Share one session schema between the plugin handshake and discovery
replies. It accepts a null file and keeps workspaceState. An unreadable
reply now fails discovery at once with an error that names the problem
and says a retry will not help. Timeout and Redis failures state their
reason.

Closes #12860

* 🐛 Stop blaming server load when MCP discovery fails on shutdown

Closing the dispatcher rejected pending discoveries with the default advice,
which blamed temporary server load. Shutdown now gives its own advice.

The dispatcher tests now cover each discovery failure path: an unreadable
reply alongside a valid one rejects with InvalidDiscoveryResponseError and
does not dispatch, a late unreadable reply leaves a settled discovery
unchanged, and timeout, Redis failure and shutdown each state their reason.
This commit is contained in:
Alonso Torres 2026-11-11 12:10:34 +01:00 committed by GitHub
parent adda9ad908
commit 1599e096da
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 224 additions and 45 deletions

View File

@ -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<PenpotSession>;

View File

@ -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;
}

View File

@ -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);

View File

@ -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<number> {
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);

View File

@ -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<string, Map<RedisPeer, (request: PluginTaskRequest) => void>>();
readonly discovery = new Map<string, Map<RedisPeer, (request: SessionDiscoveryRequest) => void>>();
readonly taskResponses = new Map<string, (response: PluginTaskResponse<unknown>) => void>();
readonly discoveryResponses = new Map<string, (response: SessionDiscoveryResponse) => void>();
readonly discoveryResponses = new Map<string, DiscoveryReplyHandlers>();
/** reply callbacks of the latest discovery request, kept after its subscription is released. */
lastDiscoveryReplyHandlers?: DiscoveryReplyHandlers;
readonly claims = new Map<string, Map<RedisPeer, (claim: SessionClaim) => 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<void>((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 () => {

View File

@ -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 {

View File

@ -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";
}
}