From 01003ff3f93fae1622931edf6e54c064c57a6890 Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Mon, 23 Mar 2026 17:53:26 +0000 Subject: [PATCH 1/4] 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 --- .changeset/abort-handlers-on-close.md | 5 +++ packages/core/src/shared/protocol.ts | 7 +++++ packages/core/src/util/inMemory.ts | 4 +++ packages/core/test/inMemory.test.ts | 36 ++++++++++++++++++++++ packages/core/test/shared/protocol.test.ts | 30 ++++++++++++++++++ 5 files changed, 82 insertions(+) create mode 100644 .changeset/abort-handlers-on-close.md diff --git a/.changeset/abort-handlers-on-close.md b/.changeset/abort-handlers-on-close.md new file mode 100644 index 000000000..b6bc65e65 --- /dev/null +++ b/.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. diff --git a/packages/core/src/shared/protocol.ts b/packages/core/src/shared/protocol.ts index b82731582..1c8ddc579 100644 --- a/packages/core/src/shared/protocol.ts +++ b/packages/core/src/shared/protocol.ts @@ -730,6 +730,9 @@ export abstract class Protocol { 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 { for (const handler of responseHandlers.values()) { handler(error); } + + for (const controller of requestHandlerAbortControllers.values()) { + controller.abort(error); + } } private _onerror(error: Error): void { diff --git a/packages/core/src/util/inMemory.ts b/packages/core/src/util/inMemory.ts index a02bcafc9..8efce7441 100644 --- a/packages/core/src/util/inMemory.ts +++ b/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 { + if (this._closed) return; + this._closed = true; + const other = this._otherTransport; this._otherTransport = undefined; await other?.close(); diff --git a/packages/core/test/inMemory.test.ts b/packages/core/test/inMemory.test.ts index 72f28240b..1bd360180 100644 --- a/packages/core/test/inMemory.test.ts +++ b/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', diff --git a/packages/core/test/shared/protocol.test.ts b/packages/core/test/shared/protocol.test.ts index 8675c1e03..cc32f11c9 100644 --- a/packages/core/test/shared/protocol.test.ts +++ b/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(resolve => { + protocol.setRequestHandler('ping', async (_request, ctx) => { + handlerStarted = true; + await new Promise(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(); From 3b88756b194b5e7925cb30bb47a24a38eee9a3c0 Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Wed, 25 Mar 2026 14:17:41 +0000 Subject: [PATCH 2/4] address review: clear timeouts on close, ensure onclose fires if peer throws - Clear _timeoutInfo and _requestResolvers in _onclose() to prevent leaked setTimeout handles from keeping the event loop alive - Wrap peer close in try-finally so InMemoryTransport.onclose fires even if the peer's onclose callback throws --- packages/core/src/shared/protocol.ts | 6 ++++++ packages/core/src/util/inMemory.ts | 7 +++++-- packages/core/test/inMemory.test.ts | 11 +++++++++++ 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/packages/core/src/shared/protocol.ts b/packages/core/src/shared/protocol.ts index 1c8ddc579..768af707f 100644 --- a/packages/core/src/shared/protocol.ts +++ b/packages/core/src/shared/protocol.ts @@ -729,6 +729,12 @@ export abstract class Protocol { this._progressHandlers.clear(); this._taskProgressTokens.clear(); this._pendingDebouncedNotifications.clear(); + this._requestResolvers.clear(); + + for (const info of this._timeoutInfo.values()) { + clearTimeout(info.timeoutId); + } + this._timeoutInfo.clear(); const requestHandlerAbortControllers = this._requestHandlerAbortControllers; this._requestHandlerAbortControllers = new Map(); diff --git a/packages/core/src/util/inMemory.ts b/packages/core/src/util/inMemory.ts index 8efce7441..a91126084 100644 --- a/packages/core/src/util/inMemory.ts +++ b/packages/core/src/util/inMemory.ts @@ -45,8 +45,11 @@ export class InMemoryTransport implements Transport { const other = this._otherTransport; this._otherTransport = undefined; - await other?.close(); - this.onclose?.(); + try { + await other?.close(); + } finally { + this.onclose?.(); + } } /** diff --git a/packages/core/test/inMemory.test.ts b/packages/core/test/inMemory.test.ts index 1bd360180..f28f7eb86 100644 --- a/packages/core/test/inMemory.test.ts +++ b/packages/core/test/inMemory.test.ts @@ -135,6 +135,17 @@ describe('InMemoryTransport', () => { expect(serverCloseCount).toBe(1); }); + test('should fire onclose even if peer onclose throws', async () => { + let clientCloseCount = 0; + clientTransport.onclose = () => clientCloseCount++; + serverTransport.onclose = () => { + throw new Error('boom'); + }; + + await expect(clientTransport.close()).rejects.toThrow('boom'); + expect(clientCloseCount).toBe(1); + }); + test('should queue messages sent before start', async () => { const message: JSONRPCMessage = { jsonrpc: '2.0', From a51268493c11dffeedd9bbee841407c6365ebca3 Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Wed, 25 Mar 2026 14:49:45 +0000 Subject: [PATCH 3/4] address review: ensure cleanup runs even if user onclose throws Wrap the user-provided onclose callback in try/finally so response handler rejection and abort controller firing happen even if the callback throws. Consistent with the pattern in InMemoryTransport. --- packages/core/src/shared/protocol.ts | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/packages/core/src/shared/protocol.ts b/packages/core/src/shared/protocol.ts index 768af707f..a050b012c 100644 --- a/packages/core/src/shared/protocol.ts +++ b/packages/core/src/shared/protocol.ts @@ -742,14 +742,17 @@ export abstract class Protocol { const error = new SdkError(SdkErrorCode.ConnectionClosed, 'Connection closed'); this._transport = undefined; - this.onclose?.(); - for (const handler of responseHandlers.values()) { - handler(error); - } + try { + this.onclose?.(); + } finally { + for (const handler of responseHandlers.values()) { + handler(error); + } - for (const controller of requestHandlerAbortControllers.values()) { - controller.abort(error); + for (const controller of requestHandlerAbortControllers.values()) { + controller.abort(error); + } } } From 01193f816636ada1a369ca08ada6880abe7afeab Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Wed, 25 Mar 2026 20:45:31 +0000 Subject: [PATCH 4/4] fix: move _requestResolvers cleanup to TaskManager.onClose after refactor The TaskManager extraction moved _requestResolvers to the new class but the merge left a stale reference in Protocol._onclose. Relocate the cleanup to TaskManager.onClose alongside _taskProgressTokens. --- packages/core/src/shared/protocol.ts | 1 - packages/core/src/shared/taskManager.ts | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/core/src/shared/protocol.ts b/packages/core/src/shared/protocol.ts index c09e52d8e..836ec3a19 100644 --- a/packages/core/src/shared/protocol.ts +++ b/packages/core/src/shared/protocol.ts @@ -493,7 +493,6 @@ export abstract class Protocol { this._progressHandlers.clear(); this._taskManager.onClose(); this._pendingDebouncedNotifications.clear(); - this._requestResolvers.clear(); for (const info of this._timeoutInfo.values()) { clearTimeout(info.timeoutId); diff --git a/packages/core/src/shared/taskManager.ts b/packages/core/src/shared/taskManager.ts index 28460d1d9..2d6f4eeaa 100644 --- a/packages/core/src/shared/taskManager.ts +++ b/packages/core/src/shared/taskManager.ts @@ -801,6 +801,7 @@ export class TaskManager { onClose(): void { this._taskProgressTokens.clear(); + this._requestResolvers.clear(); } // -- Private helpers --