mirror of
https://github.com/penpot/penpot.git
synced 2026-11-06 10:00:38 +00:00
d09bc9c3ee
The MCP plugin sends the rollbacks it saw with its next task response and then forgets them. When that response came after the server had given up on the task, the server dropped it and the agent never heard about the rollback. A reload of the tab or of the plugin lost them too. The server now sends the rollbacks of such a late response back to the plugin, which reports them again with its next response. The plugin keeps the rollbacks not reported yet in its local storage, one entry per file, so a reload does not lose them. Both MCP plugin manifests ask for the allow:localstorage permission for that.
419 lines
16 KiB
TypeScript
419 lines
16 KiB
TypeScript
import assert from "node:assert/strict";
|
|
import { randomUUID } from "node:crypto";
|
|
import { readFileSync } from "node:fs";
|
|
import test from "node:test";
|
|
import { runInNewContext } from "node:vm";
|
|
import ts from "typescript";
|
|
import * as ErrorUtils from "./ErrorUtils.ts";
|
|
import { shouldReconnectAfterClose } from "./ReconnectPolicy.ts";
|
|
import { SessionId } from "./SessionId.ts";
|
|
import { SessionIdDisplay } from "./SessionIdDisplay.ts";
|
|
import * as VersionCompatibility from "./VersionCompatibility.ts";
|
|
|
|
/** Compiles the actual plugin entry point for a simulated browser boundary. */
|
|
function compile(name: string) {
|
|
return ts.transpileModule(readFileSync(new URL(name, import.meta.url), "utf8"), {
|
|
compilerOptions: { module: ts.ModuleKind.CommonJS, target: ts.ScriptTarget.ES2020 },
|
|
}).outputText;
|
|
}
|
|
|
|
/** Exports of a compiled module that Node's type stripping cannot load directly. */
|
|
function compiledModule(name: string) {
|
|
const exports = {};
|
|
runInNewContext(compile(name), { exports });
|
|
return exports;
|
|
}
|
|
|
|
/** Resolves the imports of `plugin.ts`, stubbing the code execution handler and the rollback log. */
|
|
function pluginRequire(name: string) {
|
|
if (name === "./task-handlers/ExecuteCodeTaskHandler") return { ExecuteCodeTaskHandler: class {} };
|
|
if (name === "./TaskHandler") return compiledModule("TaskHandler.ts");
|
|
if (name === "./ErrorUtils") return ErrorUtils;
|
|
if (name === "./VersionCompatibility") return VersionCompatibility;
|
|
if (name === "./RollbackLog")
|
|
return {
|
|
LocalStorageRollbackStore: class {},
|
|
RollbackLog: class {
|
|
listen() {}
|
|
},
|
|
};
|
|
throw new Error(`Unexpected import: ${name}`);
|
|
}
|
|
|
|
function pluginUi() {
|
|
const messages: any[] = [];
|
|
const sockets: FakeSocket[] = [];
|
|
const events = new Map<string, (event?: any) => void>();
|
|
const timers = new Map<number, () => void>();
|
|
let timerId = 0;
|
|
class FakeSocket {
|
|
static OPEN = 1;
|
|
static CONNECTING = 0;
|
|
readyState = 0;
|
|
sent: any[] = [];
|
|
onopen?: () => void;
|
|
onclose?: (event: any) => void;
|
|
onmessage?: (event: any) => void;
|
|
constructor(_url: string) {
|
|
sockets.push(this);
|
|
}
|
|
open() {
|
|
this.readyState = 1;
|
|
this.onopen?.();
|
|
}
|
|
send(raw: string) {
|
|
this.sent.push(JSON.parse(raw));
|
|
}
|
|
close() {
|
|
this.readyState = 3;
|
|
this.onclose?.({ code: 1000 });
|
|
}
|
|
receive(message: any) {
|
|
this.onmessage?.({ data: JSON.stringify(message) });
|
|
}
|
|
}
|
|
const addEventListener = (name: string, handler: (event?: any) => void) => events.set(name, handler);
|
|
runInNewContext(compile("main.ts"), {
|
|
exports: {},
|
|
require: (name: string) => {
|
|
if (name === "./style.css") return {};
|
|
if (name === "./ReconnectPolicy") return { shouldReconnectAfterClose };
|
|
if (name === "./SessionId") return { SessionId };
|
|
if (name === "./SessionIdDisplay") return { SessionIdDisplay };
|
|
throw new Error(`Unexpected import: ${name}`);
|
|
},
|
|
console: { log() {}, error() {}, warn() {} },
|
|
crypto: { randomUUID },
|
|
URLSearchParams,
|
|
WebSocket: FakeSocket,
|
|
window: { location: { hash: "" }, addEventListener },
|
|
document: { body: { dataset: {} }, getElementById: () => null, addEventListener },
|
|
parent: { postMessage: (message: any) => messages.push(JSON.parse(JSON.stringify(message))) },
|
|
setTimeout: (callback: () => void) => {
|
|
timers.set(++timerId, callback);
|
|
return timerId;
|
|
},
|
|
clearTimeout: (id: number) => timers.delete(id),
|
|
setInterval: () => ++timerId,
|
|
clearInterval: () => {},
|
|
PENPOT_MCP_WEBSOCKET_URL: "ws://localhost:4402",
|
|
});
|
|
const message = (data: any) => events.get("message")!({ data });
|
|
const connect = () => {
|
|
message({ type: "start-server" });
|
|
sockets.at(-1)!.open();
|
|
return messages.filter((message) => message.type === "connection-metadata-request").at(-1)!.sessionId;
|
|
};
|
|
return { messages, sockets, events, timers, message, connect };
|
|
}
|
|
|
|
test("sends file metadata before marking a plugin connection ready", async () => {
|
|
const ui = pluginUi();
|
|
const sessionId = ui.connect();
|
|
const initialization = {
|
|
type: "initialize",
|
|
penpotUserSessionId: "tab-1",
|
|
session: { sessionId, fileId: "file-1", fileName: "Design" },
|
|
};
|
|
await ui.message(initialization);
|
|
const registeredSession = ui.sockets[0].sent[0].session;
|
|
assert.equal(registeredSession.sessionId, "ftdmhltgwk");
|
|
assert.notEqual(registeredSession.sessionId, sessionId);
|
|
assert.deepEqual(registeredSession, { ...initialization.session, sessionId: registeredSession.sessionId });
|
|
assert.equal(
|
|
ui.messages.some((message) => message.status === "connected"),
|
|
false
|
|
);
|
|
ui.sockets[0].receive({ type: "initialized" });
|
|
assert.equal(
|
|
ui.messages.filter((message) => message.type === "update-connection-status").at(-1).status,
|
|
"connected"
|
|
);
|
|
});
|
|
|
|
test("reconnects with the same session ID and ignores metadata for the old socket", async () => {
|
|
const ui = pluginUi();
|
|
const firstId = ui.connect();
|
|
await ui.message({
|
|
type: "initialize",
|
|
penpotUserSessionId: "tab-1",
|
|
session: { sessionId: firstId, fileId: "file-1", fileName: "Design" },
|
|
});
|
|
const sessionId = ui.sockets[0].sent[0].session.sessionId;
|
|
ui.sockets[0].close();
|
|
const reconnect = [...ui.timers.values()][0];
|
|
reconnect();
|
|
ui.sockets[1].open();
|
|
const secondId = ui.messages.filter((message) => message.type === "connection-metadata-request").at(-1).sessionId;
|
|
assert.notEqual(firstId, secondId);
|
|
await ui.message({
|
|
type: "initialize",
|
|
penpotUserSessionId: "tab-1",
|
|
session: { sessionId: firstId, fileId: "old", fileName: "Old file" },
|
|
});
|
|
assert.deepEqual(ui.sockets[1].sent, []);
|
|
await ui.message({
|
|
type: "initialize",
|
|
penpotUserSessionId: "tab-1",
|
|
session: { sessionId: secondId, fileId: "file-1", fileName: "Design" },
|
|
});
|
|
assert.equal(ui.sockets[1].sent[0].session.sessionId, sessionId);
|
|
assert.notEqual(ui.sockets[1].sent[0].session.sessionId, secondId);
|
|
});
|
|
|
|
test("resume and freeze events cannot send liveness messages before initialization", () => {
|
|
const ui = pluginUi();
|
|
ui.connect();
|
|
ui.events.get("resume")!();
|
|
ui.events.get("freeze")!();
|
|
assert.deepEqual(ui.sockets[0].sent, []);
|
|
ui.sockets[0].receive({ type: "initialized" });
|
|
ui.events.get("resume")!();
|
|
assert.deepEqual(ui.sockets[0].sent, [{ type: "heartbeat" }]);
|
|
});
|
|
|
|
test("answers heartbeat requests once initialized, without forwarding them to the plugin", () => {
|
|
const ui = pluginUi();
|
|
ui.connect();
|
|
ui.sockets[0].receive({ type: "heartbeat-request" });
|
|
assert.deepEqual(ui.sockets[0].sent, []);
|
|
ui.sockets[0].receive({ type: "initialized" });
|
|
ui.sockets[0].receive({ type: "heartbeat-request" });
|
|
assert.deepEqual(ui.sockets[0].sent, [{ type: "heartbeat" }]);
|
|
assert.equal(ui.messages.filter((message) => message.type === "heartbeat-request").length, 0);
|
|
});
|
|
|
|
test("plugin reads the current file for each connection metadata request", () => {
|
|
let onMessage!: (message: any) => void;
|
|
const messages: any[] = [];
|
|
const penpot = {
|
|
currentFile: { id: "file-1", name: "First design" },
|
|
currentUser: { sessionId: "tab-1" },
|
|
theme: "dark",
|
|
ui: {
|
|
open() {},
|
|
onMessage: (handler: (message: any) => void) => {
|
|
onMessage = handler;
|
|
},
|
|
sendMessage: (message: any) => messages.push(JSON.parse(JSON.stringify(message))),
|
|
},
|
|
on() {},
|
|
};
|
|
runInNewContext(compile("plugin.ts"), {
|
|
exports: {},
|
|
penpot,
|
|
penpotMgmt: {
|
|
get workspace() {
|
|
return { status: "ready", fileId: penpot.currentFile.id, fileName: penpot.currentFile.name };
|
|
},
|
|
on() {},
|
|
},
|
|
mcp: undefined,
|
|
require: pluginRequire,
|
|
});
|
|
onMessage({ type: "connection-metadata-request", sessionId: "first" });
|
|
penpot.currentFile = { id: "file-2", name: "Second design" };
|
|
onMessage({ type: "connection-metadata-request", sessionId: "second" });
|
|
assert.deepEqual(messages, [
|
|
{
|
|
type: "initialize",
|
|
penpotUserSessionId: "tab-1",
|
|
session: { sessionId: "first", fileId: "file-1", fileName: "First design", workspaceState: "ready" },
|
|
},
|
|
{
|
|
type: "initialize",
|
|
penpotUserSessionId: "tab-1",
|
|
session: { sessionId: "second", fileId: "file-2", fileName: "Second design", workspaceState: "ready" },
|
|
},
|
|
]);
|
|
});
|
|
|
|
test("plugin without penpotMgmt derives metadata from the open file", () => {
|
|
let onMessage!: (message: any) => void;
|
|
const messages: any[] = [];
|
|
const penpot = {
|
|
currentFile: { id: "file-1", name: "Design" } as { id: string; name: string } | null,
|
|
currentUser: { sessionId: "tab-1" },
|
|
theme: "dark",
|
|
ui: {
|
|
open() {},
|
|
onMessage: (handler: (message: any) => void) => {
|
|
onMessage = handler;
|
|
},
|
|
sendMessage: (message: any) => messages.push(JSON.parse(JSON.stringify(message))),
|
|
},
|
|
on() {},
|
|
};
|
|
runInNewContext(compile("plugin.ts"), {
|
|
exports: {},
|
|
penpot,
|
|
mcp: undefined,
|
|
require: pluginRequire,
|
|
});
|
|
onMessage({ type: "connection-metadata-request", sessionId: "first" });
|
|
penpot.currentFile = null;
|
|
onMessage({ type: "context-request" });
|
|
assert.deepEqual(messages, [
|
|
{
|
|
type: "initialize",
|
|
penpotUserSessionId: "tab-1",
|
|
session: { sessionId: "first", fileId: "file-1", fileName: "Design", workspaceState: "ready" },
|
|
},
|
|
{
|
|
type: "context-update",
|
|
context: { fileId: null, fileName: null, workspaceState: "none" },
|
|
},
|
|
]);
|
|
});
|
|
|
|
test("integrated connections reuse the short session ID across reconnects", async () => {
|
|
const ui = pluginUi();
|
|
ui.message({ type: "mcp-mode", integratedRemoteMcp: true });
|
|
const initialize = async (requestId: string) => {
|
|
await ui.message({
|
|
type: "initialize",
|
|
penpotUserSessionId: "tab-1",
|
|
session: { sessionId: requestId, fileId: "file-1", fileName: "Design" },
|
|
});
|
|
};
|
|
await initialize(ui.connect());
|
|
assert.equal(ui.sockets[0].sent[0].session.sessionId, "ftdmhltgwk");
|
|
ui.sockets[0].receive({ type: "initialized" });
|
|
assert.equal(
|
|
ui.messages.filter((message) => message.type === "update-connection-status").at(-1).sessionId,
|
|
"ftdmhltgwk"
|
|
);
|
|
ui.sockets[0].close();
|
|
[...ui.timers.values()][0]();
|
|
ui.sockets[1].open();
|
|
const requestId = ui.messages.filter((message) => message.type === "connection-metadata-request").at(-1).sessionId;
|
|
await initialize(requestId);
|
|
assert.equal(ui.sockets[1].sent[0].session.sessionId, "ftdmhltgwk");
|
|
});
|
|
|
|
test("disconnect cancels metadata initialization and ignores a late acknowledgement", async () => {
|
|
const ui = pluginUi();
|
|
ui.message({ type: "mcp-mode", integratedRemoteMcp: true });
|
|
const requestId = ui.connect();
|
|
const pending = ui.message({
|
|
type: "initialize",
|
|
penpotUserSessionId: "tab-1",
|
|
session: { sessionId: requestId, fileId: "file-1", fileName: "Design" },
|
|
});
|
|
ui.message({ type: "stop-server" });
|
|
await pending;
|
|
assert.deepEqual(ui.sockets[0].sent, []);
|
|
ui.sockets[0].receive({ type: "initialized" });
|
|
assert.ok(!ui.messages.some((message) => message.status === "connected"));
|
|
assert.equal(ui.messages.at(-1).status, "disconnected");
|
|
assert.equal(ui.timers.size, 0);
|
|
});
|
|
|
|
test("integrated plugin stays idle until explicitly connected, including while the UI loads", () => {
|
|
let onMessage!: (message: any) => void;
|
|
const events = new Map<string, () => void>();
|
|
const messages: any[] = [];
|
|
const statuses: string[] = [];
|
|
runInNewContext(compile("plugin.ts"), {
|
|
exports: {},
|
|
console: { log() {} },
|
|
PENPOT_MCP_VERSION: "0.0.0",
|
|
penpot: {
|
|
theme: "dark",
|
|
version: "0.0.0",
|
|
ui: {
|
|
open() {},
|
|
onMessage: (handler: (message: any) => void) => {
|
|
onMessage = handler;
|
|
},
|
|
sendMessage: (message: any) => messages.push(message),
|
|
},
|
|
on() {},
|
|
},
|
|
penpotMgmt: { on() {} },
|
|
mcp: {
|
|
getToken: () => "test-token",
|
|
getServerUrl: () => "ws://localhost:4402",
|
|
isConnectionRequested: () => false,
|
|
setMcpStatus: (status: string) => statuses.push(status),
|
|
on: (event: string, callback: () => void) => events.set(event, callback),
|
|
},
|
|
require: pluginRequire,
|
|
});
|
|
assert.ok(!statuses.includes("connecting"));
|
|
events.get("connect")!();
|
|
events.get("disconnect")!();
|
|
onMessage({ type: "ui-initialized" });
|
|
assert.ok(!messages.some((message) => message.type === "start-server"));
|
|
events.get("connect")!();
|
|
assert.equal(messages.at(-1).type, "start-server");
|
|
});
|
|
|
|
test("standalone session ID survives reconnects and plugin restarts in the same app instance", async () => {
|
|
const ui = pluginUi();
|
|
const firstRequestId = ui.connect();
|
|
await ui.message({
|
|
type: "initialize",
|
|
penpotUserSessionId: "tab-1",
|
|
session: { sessionId: firstRequestId, fileId: "file-1", fileName: "Design" },
|
|
});
|
|
const sessionId = ui.sockets[0].sent[0].session.sessionId;
|
|
ui.message({ type: "stop-server" });
|
|
assert.equal(ui.messages.at(-1).sessionId, null);
|
|
const secondRequestId = ui.connect();
|
|
await ui.message({
|
|
type: "initialize",
|
|
penpotUserSessionId: "tab-1",
|
|
session: { sessionId: secondRequestId, fileId: "file-1", fileName: "Design" },
|
|
});
|
|
assert.equal(ui.sockets[1].sent[0].session.sessionId, sessionId);
|
|
ui.sockets[1].receive({ type: "initialized" });
|
|
assert.equal(
|
|
ui.messages.filter((message) => message.type === "update-connection-status").at(-1).sessionId,
|
|
sessionId
|
|
);
|
|
|
|
const next = pluginUi();
|
|
const nextRequestId = next.connect();
|
|
await next.message({
|
|
type: "initialize",
|
|
penpotUserSessionId: "tab-1",
|
|
session: { sessionId: nextRequestId, fileId: "file-1", fileName: "Design" },
|
|
});
|
|
assert.equal(next.sockets[0].sent[0].session.sessionId, sessionId);
|
|
});
|
|
|
|
test("file metadata does not affect the short session ID", async () => {
|
|
for (const fileId of [null, "file-1", "file-2"]) {
|
|
const ui = pluginUi();
|
|
const requestId = ui.connect();
|
|
await ui.message({
|
|
type: "initialize",
|
|
penpotUserSessionId: "tab-1",
|
|
session: { sessionId: requestId, fileId, fileName: fileId },
|
|
});
|
|
assert.equal(ui.sockets[0].sent[0].session.sessionId, "ftdmhltgwk");
|
|
}
|
|
});
|
|
|
|
test("dashboard sessions use the short session ID and publish context changes without reconnecting", async () => {
|
|
const ui = pluginUi();
|
|
const requestId = ui.connect();
|
|
await ui.message({
|
|
type: "initialize",
|
|
penpotUserSessionId: "tab-1",
|
|
session: { sessionId: requestId, fileId: null, fileName: null, workspaceState: "none" },
|
|
});
|
|
assert.equal(ui.sockets[0].sent[0].session.sessionId, "ftdmhltgwk");
|
|
ui.sockets[0].receive({ type: "initialized" });
|
|
assert.equal(ui.messages.at(-1).type, "context-request");
|
|
for (const context of [
|
|
{ fileId: "file-1", fileName: "Design", workspaceState: "ready" },
|
|
{ fileId: null, fileName: null, workspaceState: "none" },
|
|
]) {
|
|
await ui.message({ type: "context-update", context });
|
|
assert.deepEqual(ui.sockets[0].sent.at(-1), { type: "context-update", context });
|
|
}
|
|
assert.equal(ui.sockets.length, 1);
|
|
});
|