From 01003ff3f93fae1622931edf6e54c064c57a6890 Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Mon, 23 Mar 2026 17:53:26 +0000 Subject: [PATCH 1/6] 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/6] 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/6] 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/6] 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 -- From 9ef07287652c49e83b8217a83431b94e42f2300b Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Thu, 26 Mar 2026 13:44:39 +0000 Subject: [PATCH 5/6] address review: guard connect() onclose, identity-check abort controller delete - Wrap preserved transport onclose in try/finally so _onclose() runs even if the transport-level handler throws - Add identity check before deleting from _requestHandlerAbortControllers in the .finally() cleanup to avoid deleting a reconnected request's controller when IDs collide across close/reconnect --- packages/core/src/shared/protocol.ts | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/packages/core/src/shared/protocol.ts b/packages/core/src/shared/protocol.ts index 836ec3a19..d6daf0172 100644 --- a/packages/core/src/shared/protocol.ts +++ b/packages/core/src/shared/protocol.ts @@ -457,8 +457,11 @@ export abstract class Protocol { this._transport = transport; const _onclose = this.transport?.onclose; this._transport.onclose = () => { - _onclose?.(); - this._onclose(); + try { + _onclose?.(); + } finally { + this._onclose(); + } }; const _onerror = this.transport?.onerror; @@ -657,7 +660,9 @@ export abstract class Protocol { ) .catch(error => this._onerror(new Error(`Failed to send response: ${error}`))) .finally(() => { - this._requestHandlerAbortControllers.delete(request.id); + if (this._requestHandlerAbortControllers.get(request.id) === abortController) { + this._requestHandlerAbortControllers.delete(request.id); + } }); } From 6dd5629a044781d6e31c7d600f71336dbd5f81d1 Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Thu, 26 Mar 2026 14:05:55 +0000 Subject: [PATCH 6/6] address review: remove NullTaskManager.onClose no-op override The override prevented the base class _requestResolvers cleanup from running. Base TaskManager.onClose just clears two maps, which is cheap and correct for the null variant too. --- packages/core/src/shared/taskManager.ts | 4 ---- 1 file changed, 4 deletions(-) diff --git a/packages/core/src/shared/taskManager.ts b/packages/core/src/shared/taskManager.ts index 2d6f4eeaa..7ca7f7b4a 100644 --- a/packages/core/src/shared/taskManager.ts +++ b/packages/core/src/shared/taskManager.ts @@ -894,8 +894,4 @@ export class NullTaskManager extends TaskManager { ): Promise<{ queued: boolean; jsonrpcNotification?: JSONRPCNotification }> { return { queued: false, jsonrpcNotification: { ...notification, jsonrpc: '2.0' } }; } - - override onClose(): void { - // No-op - } }