From dd847b0e0dbd3977cdb58ea39782429b12fb3de6 Mon Sep 17 00:00:00 2001 From: Nik <84yk8btb9f@privaterelay.appleid.com> Date: Thu, 17 Sep 2026 08:02:15 +0300 Subject: [PATCH] fix: prevent unbounded onclose handler chain on transport reconnect (#2607) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Protocol.connect() wrapped transport.onclose in a new closure on every call. When connect() is invoked multiple times on the same transport (e.g. session-resuming reconnects), each call nests a new closure around the previous one. Every transport close then walks the entire accumulated chain — a memory leak proportional to the number of reconnects. The fix detects when the transport is the same as the previously connected one and skips re-wrapping onclose/onerror/onmessage, so the existing wrapper (which already chains the prior handlers and calls this._onclose()) is reused instead of nested again. --- .changeset/unbounded-onclose-chain.md | 18 ++++++ packages/core-internal/src/shared/protocol.ts | 62 +++++++++++-------- .../test/shared/protocol.test.ts | 26 ++++++++ 3 files changed, 80 insertions(+), 26 deletions(-) create mode 100644 .changeset/unbounded-onclose-chain.md diff --git a/.changeset/unbounded-onclose-chain.md b/.changeset/unbounded-onclose-chain.md new file mode 100644 index 0000000000..72173e4733 --- /dev/null +++ b/.changeset/unbounded-onclose-chain.md @@ -0,0 +1,18 @@ +--- +'@modelcontextprotocol/core-internal': patch +--- + +Fix an unbounded `onclose` handler chain in `Protocol.connect()`. + +Each call to `connect()` on the same transport wrapped the existing +`transport.onclose` in a new closure that called the previous handler plus +`this._onclose()`. On transports that reconnect in-place (e.g. session-resuming +Streamable HTTP, or any caller that re-invokes `connect()` on the same transport +instance), the chain grew by one closure per reconnect. Every transport close +then walked the entire accumulated chain — a linear-memory leak proportional to +the number of reconnects. + +`connect()` now detects that the transport is the same as the previously +connected one and skips re-wrapping `onclose`/`onerror`/`onmessage`, so the +existing wrapper (which already chains the prior handlers and calls +`this._onclose()`) is reused instead of nested again. diff --git a/packages/core-internal/src/shared/protocol.ts b/packages/core-internal/src/shared/protocol.ts index 637be389aa..d63d658944 100644 --- a/packages/core-internal/src/shared/protocol.ts +++ b/packages/core-internal/src/shared/protocol.ts @@ -783,35 +783,45 @@ export abstract class Protocol { * The caller assumes ownership of the {@linkcode Transport}, replacing any callbacks that have already been set, and expects that it is the only user of the {@linkcode Transport} instance going forward. */ async connect(transport: Transport): Promise { + const isReconnectingSameTransport = this._transport === transport; + this._transport = transport; - const _onclose = this.transport?.onclose; - this._transport.onclose = () => { - try { - _onclose?.(); - } finally { - this._onclose(); - } - }; - const _onerror = this.transport?.onerror; - this._transport.onerror = (error: Error) => { - _onerror?.(error); - this._onerror(error); - }; + // Only wrap the transport's callbacks when connecting to a NEW transport. + // If this is a reconnect on the same transport, the callbacks were + // already wrapped by the previous connect() call — re-wrapping would + // create an unbounded chain of closures that grows with every reconnect, + // each close event walking the entire chain (a memory leak, #2607). + if (!isReconnectingSameTransport) { + const _onclose = this.transport?.onclose; + this._transport.onclose = () => { + try { + _onclose?.(); + } finally { + this._onclose(); + } + }; - const _onmessage = this._transport?.onmessage; - this._transport.onmessage = (message, extra) => { - _onmessage?.(message, extra); - if (isJSONRPCResultResponse(message) || isJSONRPCErrorResponse(message)) { - this._onresponse(message); - } else if (isJSONRPCRequest(message)) { - this._onrequest(message, extra); - } else if (isJSONRPCNotification(message)) { - this._onnotification(message, extra); - } else { - this._onerror(new Error(`Unknown message type: ${JSON.stringify(message)}`)); - } - }; + const _onerror = this.transport?.onerror; + this._transport.onerror = (error: Error) => { + _onerror?.(error); + this._onerror(error); + }; + + const _onmessage = this._transport?.onmessage; + this._transport.onmessage = (message, extra) => { + _onmessage?.(message, extra); + if (isJSONRPCResultResponse(message) || isJSONRPCErrorResponse(message)) { + this._onresponse(message); + } else if (isJSONRPCRequest(message)) { + this._onrequest(message, extra); + } else if (isJSONRPCNotification(message)) { + this._onnotification(message, extra); + } else { + this._onerror(new Error(`Unknown message type: ${JSON.stringify(message)}`)); + } + }; + } // Pass supported protocol versions to transport for header validation transport.setSupportedProtocolVersions?.(this._supportedProtocolVersions); diff --git a/packages/core-internal/test/shared/protocol.test.ts b/packages/core-internal/test/shared/protocol.test.ts index 2fb0f64813..79286d979c 100644 --- a/packages/core-internal/test/shared/protocol.test.ts +++ b/packages/core-internal/test/shared/protocol.test.ts @@ -209,6 +209,32 @@ describe('protocol tests', () => { expect(onmessageMock).toHaveBeenCalled(); }); + test('should not create unbounded onclose chain on reconnect to same transport (#2607)', async () => { + // On reconnect to the SAME transport, connect() must not re-wrap + // onclose/onerror/onmessage. Otherwise each reconnect nests a new + // closure around the previous one, and every transport close walks + // the full accumulated chain — a memory leak proportional to reconnects. + await protocol.connect(transport); + + // Capture the wrapper that the first connect() installed. + const firstOncloseWrapper = transport.onclose; + expect(firstOncloseWrapper).toBeDefined(); + + // Reconnect to the SAME transport instance. + await protocol.connect(transport); + + // The onclose handler must be the SAME closure, not a new wrapper + // that calls the old one. If it were re-wrapped, each close would walk + // a chain of N closures after N reconnects. + expect(transport.onclose).toBe(firstOncloseWrapper); + + // Verify the single wrapper still fires onclose exactly once. + const oncloseMock = vi.fn(); + protocol.onclose = oncloseMock; + transport.onclose?.(); + expect(oncloseMock).toHaveBeenCalledTimes(1); + }); + describe('_meta preservation with onprogress', () => { test('should preserve existing _meta when adding progressToken', async () => { await protocol.connect(transport);