fix(client): preserve pre-set transport handlers across the version-negotiation probe window (#2455)

Author: felixweinberger

.changeset/probe-window-handler-restore.md

@@ -0,0 +1,5 @@
+---
+'@modelcontextprotocol/client': patch
+---
+
+Version negotiation no longer discards transport handlers the caller set before `connect()`. The probe window now saves any pre-set `onmessage`/`onerror`/`onclose`, forwards error and close events to them while the probe is in flight, and restores them when the window closes — so `Protocol.connect()` chains them exactly as it does on a plain connect. Previously, connecting with `versionNegotiation` silently cleared pre-set handlers (e.g. an `onerror` used to detect session-expiry auth failures), leaving them permanently detached for the life of the connection.

packages/client/src/client/versionNegotiation.ts

@@ -171,12 +171,30 @@ type RawProbeReply =
  * starts the transport; `release()` detaches them and arms a one-shot `start()`
  * pass-through so the subsequent Protocol connect (which always starts its
  * transport) takes over the already-started channel without a double-start error.
+ *
+ * Handlers a caller pre-set on the transport before `connect()` are saved on
+ * open and restored on detach, so `Protocol.connect()` finds and chains them
+ * exactly as it would on a plain (non-negotiating) connect. Error and close
+ * events are forwarded to the saved handlers during the window, so negotiation
+ * does not change what a pre-attached observer sees of the transport
+ * lifecycle. (Inbound messages are deliberately NOT forwarded — the window's
+ * drop-guard is a module invariant, and the probe reply has no plain-connect
+ * counterpart; a pre-set `onmessage` is restored at detach and chained by
+ * `Protocol.connect()` like the others.)
  */
 class ProbeWindow {
     private _pending: { id: string; resolve: (reply: RawProbeReply) => void } | undefined;
     private _probeCounter = 0;
+    private readonly _savedOnMessage: Transport['onmessage'];
+    private readonly _savedOnError: Transport['onerror'];
+    private readonly _savedOnClose: Transport['onclose'];
+    private _closeDelivered = false;
 
-    private constructor(private readonly _transport: Transport) {}
+    private constructor(private readonly _transport: Transport) {
+        this._savedOnMessage = _transport.onmessage;
+        this._savedOnError = _transport.onerror;
+        this._savedOnClose = _transport.onclose;
+    }
 
     static async open(transport: Transport): Promise<ProbeWindow> {
         const window = new ProbeWindow(transport);
@@ -197,18 +215,29 @@ class ProbeWindow {
             }
             // Probe-window guard: drop everything else with zero bytes written back (see module doc).
         };
-        transport.onerror = () => {
+        transport.onerror = error => {
             // Out-of-band transport errors are not necessarily fatal; the probe
             // resolves via a send failure, the close signal, or the timeout.
+            window._savedOnError?.(error);
         };
         transport.onclose = () => {
             const pending = window._pending;
             if (pending !== undefined) {
                 window._pending = undefined;
                 pending.resolve({ kind: 'closed' });
             }
+            // Forward exactly once: after a mid-window close is delivered, the
+            // restored handler must not re-deliver it when cleanup paths call
+            // `transport.close()` again.
+            window._closeDelivered = true;
+            window._savedOnClose?.();
         };
-        await transport.start();
+        try {
+            await transport.start();
+        } catch (error) {
+            window.detach();
+            throw error;
+        }
         return window;
     }
 
@@ -235,12 +264,12 @@ class ProbeWindow {
         });
     }
 
-    /** Detach the window's handlers, leaving the transport's own `start` untouched. */
+    /** Detach the window's handlers, restoring any the caller pre-set, leaving the transport's own `start` untouched. */
     detach(): void {
         this._pending = undefined;
-        this._transport.onmessage = undefined;
-        this._transport.onerror = undefined;
-        this._transport.onclose = undefined;
+        this._transport.onmessage = this._savedOnMessage;
+        this._transport.onerror = this._savedOnError;
+        this._transport.onclose = this._closeDelivered ? undefined : this._savedOnClose;
     }
 
     /** Detach the handlers and arm the one-shot `start()` pass-through for the `Protocol.connect()` handover. */

packages/client/test/client/versionNegotiation.test.ts

@@ -782,3 +782,116 @@ describe('probe send-error classification', () => {
         expect(requests(transport.sent).some(r => r.method === 'initialize')).toBe(false);
     });
 });
+
+/* ------------------------------------------------------------------------- *
+ * Probe window handler preservation: handlers pre-set on the transport
+ * before connect() survive negotiation exactly as they survive a plain
+ * connect — Protocol.connect() must find and chain them after the window.
+ * ------------------------------------------------------------------------- */
+
+describe('probe window preserves pre-set transport handlers', () => {
+    test('pre-set onerror/onclose are restored after a modern negotiation and reachable through the Protocol chain', async () => {
+        const transport = new ScriptedTransport(modernServerScript());
+        const seenErrors: Error[] = [];
+        let closed = 0;
+        transport.onerror = error => {
+            seenErrors.push(error);
+        };
+        transport.onclose = () => {
+            closed++;
+        };
+
+        const client = new Client({ name: 'c', version: '0' }, { versionNegotiation: { mode: 'auto' } });
+        await client.connect(transport);
+
+        // The window restored the handler, so Protocol.connect chained it:
+        // post-connect transport errors still reach the pre-set observer.
+        const boom = new Error('post-connect transport error');
+        transport.onerror?.(boom);
+        expect(seenErrors).toContain(boom);
+
+        await client.close();
+        expect(closed).toBeGreaterThan(0);
+    });
+
+    test('pre-set onerror survives the legacy fallback path too', async () => {
+        const transport = new ScriptedTransport(legacyServerScript);
+        const seenErrors: Error[] = [];
+        transport.onerror = error => {
+            seenErrors.push(error);
+        };
+
+        const client = new Client({ name: 'c', version: '0' }, { versionNegotiation: { mode: 'auto' } });
+        await client.connect(transport);
+        expect(client.getNegotiatedProtocolVersion()).toBe('2025-11-25');
+
+        const boom = new Error('post-fallback transport error');
+        transport.onerror?.(boom);
+        expect(seenErrors).toContain(boom);
+
+        await client.close();
+    });
+
+    test('failed negotiation (pin mode, no fallback) restores handlers via the detach path — onclose fires exactly once', async () => {
+        const transport = new ScriptedTransport(legacyServerScript);
+        const presetOnError = (_error: Error) => {};
+        let closes = 0;
+        transport.onerror = presetOnError;
+        transport.onclose = () => {
+            closes++;
+        };
+
+        const client = new Client({ name: 'c', version: '0' }, { versionNegotiation: { mode: { pin: MODERN } } });
+        await expect(client.connect(transport)).rejects.toThrow();
+
+        // detach() restored the pre-set handlers, and Client's cleanup close
+        // delivered the close event exactly once.
+        expect(transport.onerror).toBe(presetOnError);
+        expect(closes).toBe(1);
+    });
+
+    test('a mid-probe transport close reaches the pre-set onclose exactly once (no re-delivery from cleanup close)', async () => {
+        const script: Script = (message, t) => {
+            if (!isJSONRPCRequest(message)) return;
+            if (message.method === 'server/discover') {
+                t.onclose?.();
+                return;
+            }
+            legacyServerScript(message, t);
+        };
+        const transport = new ScriptedTransport(script);
+        let closes = 0;
+        transport.onclose = () => {
+            closes++;
+        };
+
+        const client = new Client({ name: 'c', version: '0' }, { versionNegotiation: { mode: 'auto' } });
+        await client.connect(transport).catch(() => undefined);
+
+        expect(closes).toBe(1);
+    });
+
+    test('transport errors DURING the probe window are forwarded to the pre-set handler', async () => {
+        const duringProbe = new Error('mid-probe transport error');
+        const script: Script = (message, t) => {
+            if (!isJSONRPCRequest(message)) return;
+            if (message.method === 'server/discover') {
+                t.onerror?.(duringProbe);
+                t.reply({ jsonrpc: '2.0', id: message.id, error: { code: -32_601, message: 'Method not found' } });
+                return;
+            }
+            legacyServerScript(message, t);
+        };
+        const transport = new ScriptedTransport(script);
+        const seenErrors: Error[] = [];
+        transport.onerror = error => {
+            seenErrors.push(error);
+        };
+
+        const client = new Client({ name: 'c', version: '0' }, { versionNegotiation: { mode: 'auto' } });
+        await client.connect(transport);
+
+        expect(seenErrors).toContain(duringProbe);
+        await client.close();
+    });
+});