From b30e7cc036ef026a47192dece52c47139342d080 Mon Sep 17 00:00:00 2001
From: Steven Enamakel <31011319+senamakel@users.noreply.github.com>
Date: Sat, 9 May 2026 13:48:16 -0700
Subject: [PATCH] Fix deregistration bug in transport.off() (#1409)

---
 app/src/lib/mcp/transport.test.ts | 23 ++++++++++++++++-------
 app/src/lib/mcp/transport.ts      | 20 +++++++++++++++++++-
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/app/src/lib/mcp/transport.test.ts b/app/src/lib/mcp/transport.test.ts
index af5777f0c..601091233 100644
--- a/app/src/lib/mcp/transport.test.ts
+++ b/app/src/lib/mcp/transport.test.ts
@@ -108,13 +108,22 @@ describe('SocketIOMCPTransportImpl — on / off', () => {
     expect(handler).toHaveBeenCalledWith({ data: 42 });
   });
 
-  // TODO: transport.off() passes the original handler reference to socket.off() instead of
-  // the wrapped handler registered via socket.on(). This means the wrapped handler may still
-  // fire after off() is called — the deregistration is a no-op in practice. Fix the bug in
-  // transport.ts first, then replace this todo with a behavioural assertion.
-  it.todo(
-    'off removes the handler (pending fix: off() passes original instead of wrapped handler)'
-  );
+  it('off removes the handler by resolving the wrapped handler', () => {
+    const socket = makeSocket();
+    const transport = new SocketIOMCPTransportImpl(socket as never);
+    const handler = vi.fn();
+
+    transport.on('tool_call', handler);
+    // Trigger via the raw socket to ensure it works
+    socket.trigger('mcp:tool_call', { data: 42 });
+    expect(handler).toHaveBeenCalledWith({ data: 42 });
+    expect(handler).toHaveBeenCalledTimes(1);
+
+    transport.off('tool_call', handler);
+    // Trigger again, the handler should NOT be called
+    socket.trigger('mcp:tool_call', { data: 84 });
+    expect(handler).toHaveBeenCalledTimes(1); // Still 1
+  });
 });
 
 describe('SocketIOMCPTransportImpl — request / response routing', () => {
diff --git a/app/src/lib/mcp/transport.ts b/app/src/lib/mcp/transport.ts
index 58eef3921..b5653a377 100644
--- a/app/src/lib/mcp/transport.ts
+++ b/app/src/lib/mcp/transport.ts
@@ -11,6 +11,7 @@ import type { MCPRequest, MCPResponse, SocketIOMCPTransport } from './types';
 export class SocketIOMCPTransportImpl implements SocketIOMCPTransport {
   private socket: Socket | null | undefined;
   private requestHandlers = new Map<string | number, (response: MCPResponse) => void>();
+  private eventHandlers = new Map<string, Map<Function, Function>>();
   private readonly eventPrefix = 'mcp:';
   private responseHandler = (response: MCPResponse): void => {
     mcpLog(
@@ -86,12 +87,29 @@ export class SocketIOMCPTransportImpl implements SocketIOMCPTransport {
       mcpLog('Received event', createSafeLogData({ event: fullEvent }, data));
       handler(data);
     };
+
+    let handlersForEvent = this.eventHandlers.get(fullEvent);
+    if (!handlersForEvent) {
+      handlersForEvent = new Map();
+      this.eventHandlers.set(fullEvent, handlersForEvent);
+    }
+    handlersForEvent.set(handler, wrappedHandler);
+
     this.socket.on(fullEvent, wrappedHandler);
   }
 
   off(event: string, handler: (data: unknown) => void): void {
     if (!this.socket) return;
-    this.socket.off(`${this.eventPrefix}${event}`, handler);
+    const fullEvent = `${this.eventPrefix}${event}`;
+    const handlersForEvent = this.eventHandlers.get(fullEvent);
+    const wrappedHandler = handlersForEvent?.get(handler);
+
+    if (wrappedHandler) {
+      this.socket.off(fullEvent, wrappedHandler as any);
+      handlersForEvent?.delete(handler);
+    } else {
+      this.socket.off(fullEvent, handler);
+    }
   }
 
   async request(request: MCPRequest, timeoutMs = 30000): Promise<MCPResponse> {