fix: abort in-flight request handlers on connection close

Previously, request handlers would continue running after the transport
disconnected, wasting resources and preventing cleanup of long-running
operations. Protocol._onclose() now aborts all active request handler
AbortControllers with a ConnectionClosed error.

Also fixes InMemoryTransport.close() firing onclose twice on the
initiating side due to peer recursion.

Fixes #611

Co-authored-by: Aljosa Asanovic <aljosa.a@gmail.com>

Author: felixweinberger

.changeset/abort-handlers-on-close.md

@@ -0,0 +1,5 @@
+---
+'@modelcontextprotocol/core': patch
+---
+
+Abort in-flight request handlers when the connection closes. Previously, request handlers would continue running after the transport disconnected, wasting resources and preventing proper cleanup. Also fixes `InMemoryTransport.close()` firing `onclose` twice on the initiating side.

packages/core/src/shared/protocol.ts

@@ -730,6 +730,9 @@ export abstract class Protocol<ContextT extends BaseContext> {
         this._taskProgressTokens.clear();
         this._pendingDebouncedNotifications.clear();
 
+        const requestHandlerAbortControllers = this._requestHandlerAbortControllers;
+        this._requestHandlerAbortControllers = new Map();
+
         const error = new SdkError(SdkErrorCode.ConnectionClosed, 'Connection closed');
 
         this._transport = undefined;
@@ -738,6 +741,10 @@ export abstract class Protocol<ContextT extends BaseContext> {
         for (const handler of responseHandlers.values()) {
             handler(error);
         }
+
+        for (const controller of requestHandlerAbortControllers.values()) {
+            controller.abort(error);
+        }
     }
 
     private _onerror(error: Error): void {

packages/core/src/util/inMemory.ts

@@ -13,6 +13,7 @@ interface QueuedMessage {
 export class InMemoryTransport implements Transport {
     private _otherTransport?: InMemoryTransport;
     private _messageQueue: QueuedMessage[] = [];
+    private _closed = false;
 
     onclose?: () => void;
     onerror?: (error: Error) => void;
@@ -39,6 +40,9 @@ export class InMemoryTransport implements Transport {
     }
 
     async close(): Promise<void> {
+        if (this._closed) return;
+        this._closed = true;
+
         const other = this._otherTransport;
         this._otherTransport = undefined;
         await other?.close();

packages/core/test/inMemory.test.ts

@@ -99,6 +99,42 @@ describe('InMemoryTransport', () => {
         await expect(clientTransport.send({ jsonrpc: '2.0', method: 'test', id: 1 })).rejects.toThrow('Not connected');
     });
 
+    test('should fire onclose exactly once per transport', async () => {
+        let clientCloseCount = 0;
+        let serverCloseCount = 0;
+
+        clientTransport.onclose = () => clientCloseCount++;
+        serverTransport.onclose = () => serverCloseCount++;
+
+        await clientTransport.close();
+
+        expect(clientCloseCount).toBe(1);
+        expect(serverCloseCount).toBe(1);
+    });
+
+    test('should handle double close idempotently', async () => {
+        let clientCloseCount = 0;
+        clientTransport.onclose = () => clientCloseCount++;
+
+        await clientTransport.close();
+        await clientTransport.close();
+
+        expect(clientCloseCount).toBe(1);
+    });
+
+    test('should handle concurrent close from both sides', async () => {
+        let clientCloseCount = 0;
+        let serverCloseCount = 0;
+
+        clientTransport.onclose = () => clientCloseCount++;
+        serverTransport.onclose = () => serverCloseCount++;
+
+        await Promise.all([clientTransport.close(), serverTransport.close()]);
+
+        expect(clientCloseCount).toBe(1);
+        expect(serverCloseCount).toBe(1);
+    });
+
     test('should queue messages sent before start', async () => {
         const message: JSONRPCMessage = {
             jsonrpc: '2.0',

packages/core/test/shared/protocol.test.ts

@@ -209,6 +209,36 @@ describe('protocol tests', () => {
         expect(oncloseMock).toHaveBeenCalled();
     });
 
+    test('should abort in-flight request handlers when the connection is closed', async () => {
+        await protocol.connect(transport);
+
+        let abortReason: unknown;
+        let handlerStarted = false;
+        const handlerDone = new Promise<void>(resolve => {
+            protocol.setRequestHandler('ping', async (_request, ctx) => {
+                handlerStarted = true;
+                await new Promise<void>(resolveInner => {
+                    ctx.mcpReq.signal.addEventListener('abort', () => {
+                        abortReason = ctx.mcpReq.signal.reason;
+                        resolveInner();
+                    });
+                });
+                resolve();
+                return {};
+            });
+        });
+
+        transport.onmessage?.({ jsonrpc: '2.0', id: 1, method: 'ping', params: {} });
+
+        await vi.waitFor(() => expect(handlerStarted).toBe(true));
+
+        await transport.close();
+        await handlerDone;
+
+        expect(abortReason).toBeInstanceOf(SdkError);
+        expect((abortReason as SdkError).code).toBe(SdkErrorCode.ConnectionClosed);
+    });
+
     test('should not overwrite existing hooks when connecting transports', async () => {
         const oncloseMock = vi.fn();
         const onerrorMock = vi.fn();