🪵 refactor: onmessage Handler and Restructure MCP Debug Logging (#12004)

* 🪵 refactor: Simplify MCP Transport Log Messages

- Updated the logging in MCPConnection to provide clearer output by explicitly logging the method and ID of messages received and sent, improving traceability during debugging.
- This change replaces the previous JSON stringification of messages with a more structured log format, enhancing readability and understanding of the transport interactions.

* 🔧 refactor: Streamline MCPConnection Message Handling

- Removed redundant onmessage logging in MCPConnection to simplify the codebase.
- Introduced a dedicated setupTransportOnMessageHandler method to centralize message handling and improve clarity in transport interactions.
- Enhanced logging to provide clearer output for received messages, ensuring better traceability during debugging.

* 🔧 refactor: Rename setupTransportDebugHandlers to patchTransportSend

- Updated the MCPConnection class to rename the setupTransportDebugHandlers method to patchTransportSend for improved clarity.
- Adjusted the method call in the connection setup process to reflect the new naming, enhancing code readability and maintainability.

packages/api/src/mcp/connection.ts

@@ -11,7 +11,6 @@ import { WebSocketClientTransport } from '@modelcontextprotocol/sdk/client/webso
 import { ResourceListChangedNotificationSchema } from '@modelcontextprotocol/sdk/types.js';
 import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
 import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
-import type { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
 import type {
   RequestInit as UndiciRequestInit,
   RequestInfo as UndiciRequestInfo,
@@ -503,10 +502,6 @@ export class MCPConnection extends EventEmitter {
             this.emit('connectionChange', 'disconnected');
           };
 
-          transport.onmessage = (message) => {
-            logger.info(`${this.getLogPrefix()} Message received: ${JSON.stringify(message)}`);
-          };
-
           this.setupTransportErrorHandlers(transport);
           return transport;
         }
@@ -545,10 +540,6 @@ export class MCPConnection extends EventEmitter {
             this.emit('connectionChange', 'disconnected');
           };
 
-          transport.onmessage = (message: JSONRPCMessage) => {
-            logger.info(`${this.getLogPrefix()} Message received: ${JSON.stringify(message)}`);
-          };
-
           this.setupTransportErrorHandlers(transport);
           return transport;
         }
@@ -700,7 +691,7 @@ export class MCPConnection extends EventEmitter {
         }
 
         this.transport = await runOutsideTracing(() => this.constructTransport(this.options));
-        this.setupTransportDebugHandlers();
+        this.patchTransportSend();
 
         const connectTimeout = this.options.initTimeout ?? 120000;
         await runOutsideTracing(() =>
@@ -711,6 +702,7 @@ export class MCPConnection extends EventEmitter {
           ),
         );
 
+        this.setupTransportOnMessageHandler();
         this.connectionState = 'connected';
         this.emit('connectionChange', 'connected');
         this.reconnectAttempts = 0;
@@ -824,15 +816,11 @@ export class MCPConnection extends EventEmitter {
     return this.connectPromise;
   }
 
-  private setupTransportDebugHandlers(): void {
+  private patchTransportSend(): void {
     if (!this.transport) {
       return;
     }
 
-    this.transport.onmessage = (msg) => {
-      logger.debug(`${this.getLogPrefix()} Transport received: ${JSON.stringify(msg)}`);
-    };
-
     const originalSend = this.transport.send.bind(this.transport);
     this.transport.send = async (msg) => {
       if ('result' in msg && !('method' in msg) && Object.keys(msg.result ?? {}).length === 0) {
@@ -841,11 +829,31 @@ export class MCPConnection extends EventEmitter {
         }
         this.lastPingTime = Date.now();
       }
-      logger.debug(`${this.getLogPrefix()} Transport sending: ${JSON.stringify(msg)}`);
+      const method = 'method' in msg ? msg.method : undefined;
+      const id = 'id' in msg ? (msg as { id: string | number | null }).id : undefined;
+      logger.debug(
+        `${this.getLogPrefix()} Transport sending: method=${method ?? 'response'} id=${id ?? 'none'}`,
+      );
       return originalSend(msg);
     };
   }
 
+  private setupTransportOnMessageHandler(): void {
+    if (!this.transport?.onmessage) {
+      return;
+    }
+
+    const sdkHandler = this.transport.onmessage;
+    this.transport.onmessage = (msg) => {
+      const method = 'method' in msg ? msg.method : undefined;
+      const id = 'id' in msg ? (msg as { id: string | number | null }).id : undefined;
+      logger.debug(
+        `${this.getLogPrefix()} Transport received: method=${method ?? 'response'} id=${id ?? 'none'}`,
+      );
+      sdkHandler(msg);
+    };
+  }
+
   async connect(): Promise<void> {
     try {
       await this.disconnect();