diff --git a/mcp/packages/server/src/PluginBridge.sessions.test.ts b/mcp/packages/server/src/PluginBridge.sessions.test.ts index 1a42241169..bc2d563602 100644 --- a/mcp/packages/server/src/PluginBridge.sessions.test.ts +++ b/mcp/packages/server/src/PluginBridge.sessions.test.ts @@ -169,3 +169,21 @@ test("a dashboard session remains routable while its workspace metadata changes" assert.deepEqual(bridge.getUserConnections("alice")!.getSessions(), [{ sessionId: "tab", ...context }]); } }); + +test("a task fails at once when its plugin connection closes before answering", async () => { + const plugin = new WebSocket(`ws://127.0.0.1:${port}?userToken=alice`); + sockets.push(plugin); + await once(plugin, "open"); + const initialized = once(plugin, "message"); + plugin.send(JSON.stringify({ type: "initialize", session: { sessionId: "tab", fileId: null, fileName: null } })); + await initialized; + // the socket closes once the task reaches it + plugin.on("message", () => plugin.terminate()); + const started = Date.now(); + await assert.rejects( + bridge.executePluginTask(new PluginTask("test", {}), "tab"), + /connection closed before the task finished/ + ); + // well within the 1s task timeout of these tests + assert.ok(Date.now() - started < 500); +}); diff --git a/mcp/packages/server/src/PluginBridge.ts b/mcp/packages/server/src/PluginBridge.ts index 913b28b26a..127ef070d6 100644 --- a/mcp/packages/server/src/PluginBridge.ts +++ b/mcp/packages/server/src/PluginBridge.ts @@ -187,8 +187,8 @@ export class PluginBridge implements TaskDispatchHost { } }); - ws.on("close", () => { - this.logger.info("WebSocket connection closed"); + ws.on("close", (code: number, reason: Buffer) => { + this.logger.info("WebSocket connection closed (code %d%s)", code, reason.length ? `: ${reason}` : ""); clearTimeout(initializationTimeout); this.removeConnection(ws); }); @@ -262,8 +262,9 @@ export class PluginBridge implements TaskDispatchHost { * * Clears the per-connection keep-alive interval and removes the connection from the * socket-keyed index and the user's session collection. The token-keyed index entry - * is removed only when the user's last local connection closes. Safe to call with a - * socket that is not (or no longer) registered. + * is removed only when the user's last local connection closes. Tasks sent over the + * connection and still awaiting a response fail at once. Safe to call with a socket + * that is not (or no longer) registered. * * @param ws - The WebSocket whose connection state should be removed */ @@ -277,6 +278,13 @@ export class PluginBridge implements TaskDispatchHost { this.connectionsBySocket.delete(ws); const connections = this.getUserConnections(connection.userToken)!; connections.remove(connection); + // a task sent over the closed socket can no longer be answered, since responses only count + // from the connection the task went to + for (const [taskId, taskConnection] of [...this.taskConnections]) { + if (taskConnection === connection) { + this.rejectPendingTask(taskId, new Error("The plugin connection closed before the task finished.")); + } + } if (connection.userToken !== null && connections.size === 0) { this.userPenpotConnectionsByToken.delete(connection.userToken); }