🐛 Fail MCP tasks at once when their plugin connection closes

A response only counts when it comes from the connection the task went
to, so a task sent over a connection that then closes can never finish.
It used to wait for its full timeout. The plugin bridge now fails such
tasks as soon as their connection closes.

Connections also log their close code and reason, to tell a page
closing its socket from a browser or network failure.
This commit is contained in:
alonso.torres 2026-10-30 09:11:49 +00:00
parent 0fa1077b9a
commit ab10d8ca93
2 changed files with 30 additions and 4 deletions

View File

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

View File

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