From 8882f18db45da4231b20dd253f0bdb409d07083b Mon Sep 17 00:00:00 2001 From: Dominik Jain Date: Thu, 26 Feb 2026 17:22:24 +0100 Subject: [PATCH] :ambulance: Fix multi-user mode MCP connections Previously, only the latest streamable HTTP connection was operational --- mcp/packages/server/src/PenpotMcpServer.ts | 98 +++++++++------------- 1 file changed, 40 insertions(+), 58 deletions(-) diff --git a/mcp/packages/server/src/PenpotMcpServer.ts b/mcp/packages/server/src/PenpotMcpServer.ts index 3aa741dd83..6dab371ff4 100644 --- a/mcp/packages/server/src/PenpotMcpServer.ts +++ b/mcp/packages/server/src/PenpotMcpServer.ts @@ -23,7 +23,6 @@ export interface SessionContext { export class PenpotMcpServer { private readonly logger = createLogger("PenpotMcpServer"); - private readonly server: McpServer; private readonly tools: Map>; public readonly configLoader: ConfigurationLoader; private app: any; @@ -36,10 +35,7 @@ export class PenpotMcpServer { */ private readonly sessionContext = new AsyncLocalStorage(); - private readonly transports = { - streamable: {} as Record, - sse: {} as Record, - }; + private readonly sseTransports: Record = {}; public readonly host: string; public readonly port: number; @@ -56,21 +52,11 @@ export class PenpotMcpServer { this.configLoader = new ConfigurationLoader(process.cwd()); this.apiDocs = new ApiDocs(); - this.server = new McpServer( - { - name: "penpot-mcp-server", - version: "1.0.0", - }, - { - instructions: this.getInitialInstructions(), - } - ); - this.tools = new Map>(); this.pluginBridge = new PluginBridge(this, this.webSocketPort); this.replServer = new ReplServer(this.pluginBridge, this.replPort); - this.registerTools(); + this.initTools(); } /** @@ -119,35 +105,44 @@ export class PenpotMcpServer { return this.sessionContext.getStore(); } - private registerTools(): void { - // Create relevant tool instances (depending on file system access) + private initTools(): void { const toolInstances: Tool[] = [ new ExecuteCodeTool(this), new HighLevelOverviewTool(this), new PenpotApiInfoTool(this, this.apiDocs), - new ExportShapeTool(this), // tool adapts to file system access internally + new ExportShapeTool(this), ]; if (this.isFileSystemAccessEnabled()) { toolInstances.push(new ImportImageTool(this)); } for (const tool of toolInstances) { - const toolName = tool.getToolName(); - this.tools.set(toolName, tool); + this.logger.info(`Registering tool: ${tool.getToolName()}`); + this.tools.set(tool.getToolName(), tool); + } + } - // Register each tool with McpServer - this.logger.info(`Registering tool: ${toolName}`); - this.server.registerTool( - toolName, + /** + * Creates a fresh {@link McpServer} instance with all tools registered. + */ + private createMcpServer(): McpServer { + const server = new McpServer( + { name: "penpot-mcp-server", version: "1.0.0" }, + { instructions: this.getInitialInstructions() } + ); + + for (const tool of this.tools.values()) { + server.registerTool( + tool.getToolName(), { description: tool.getToolDescription(), inputSchema: tool.getInputSchema(), }, - async (args) => { - return tool.execute(args); - } + async (args) => tool.execute(args) ); } + + return server; } private setupHttpEndpoints(): void { @@ -156,51 +151,38 @@ export class PenpotMcpServer { */ this.app.all("/mcp", async (req: any, res: any) => { const userToken = req.query.userToken as string | undefined; + this.logger.info(`Received /mcp request with userToken: ${userToken}`); await this.sessionContext.run({ userToken }, async () => { - const { randomUUID } = await import("node:crypto"); - - const sessionId = req.headers["mcp-session-id"] as string | undefined; - let transport: StreamableHTTPServerTransport; - - if (sessionId && this.transports.streamable[sessionId]) { - transport = this.transports.streamable[sessionId]; - } else { - transport = new StreamableHTTPServerTransport({ - sessionIdGenerator: () => randomUUID(), - onsessioninitialized: (id: string) => { - this.transports.streamable[id] = transport; - }, - }); - - transport.onclose = () => { - if (transport.sessionId) { - delete this.transports.streamable[transport.sessionId]; - } - }; - - await this.server.connect(transport); - } - + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: undefined, + }); + const server = this.createMcpServer(); + await server.connect(transport); await transport.handleRequest(req, res, req.body); + res.on("close", () => { + transport.close(); + server.close(); + }); }); }); /** - * Legacy SSE connection endpoint + * Legacy SSE connection endpoint. */ this.app.get("/sse", async (req: any, res: any) => { const userToken = req.query.userToken as string | undefined; await this.sessionContext.run({ userToken }, async () => { const transport = new SSEServerTransport("/messages", res); - this.transports.sse[transport.sessionId] = { transport, userToken }; + this.sseTransports[transport.sessionId] = { transport, userToken }; + const server = this.createMcpServer(); + await server.connect(transport); res.on("close", () => { - delete this.transports.sse[transport.sessionId]; + delete this.sseTransports[transport.sessionId]; + server.close(); }); - - await this.server.connect(transport); }); }); @@ -209,7 +191,7 @@ export class PenpotMcpServer { */ this.app.post("/messages", async (req: any, res: any) => { const sessionId = req.query.sessionId as string; - const session = this.transports.sse[sessionId]; + const session = this.sseTransports[sessionId]; if (session) { await this.sessionContext.run({ userToken: session.userToken }, async () => {