From 0bc4c25910b55a867fffb00cefd3766dd973274c Mon Sep 17 00:00:00 2001
From: echozyr2001 <echo.zyr.2001@gmail.com>
Date: Mon, 21 Jul 2025 13:57:19 +0800
Subject: [PATCH] fix: resolve initialization protocol violation

---
 src/custom-stdio.ts | 54 ++++++++++++++++++++++++++++++++++++++++-----
 src/index.ts        | 11 +++++++++
 src/server.ts       |  3 ++-
 3 files changed, 61 insertions(+), 7 deletions(-)

diff --git a/src/custom-stdio.ts b/src/custom-stdio.ts
index 20730c66..c65bdb77 100644
--- a/src/custom-stdio.ts
+++ b/src/custom-stdio.ts
@@ -24,6 +24,7 @@ export class FilteredStdioServerTransport extends StdioServerTransport {
     info: typeof console.info;
   };
   private originalStdoutWrite: typeof process.stdout.write;
+  private isInitialized: boolean = false;
 
   constructor() {
     super();
@@ -49,25 +50,61 @@ export class FilteredStdioServerTransport extends StdioServerTransport {
     process.stderr.write(`[desktop-commander] Enhanced FilteredStdioServerTransport initialized\n`);
   }
 
+  /**
+   * Call this method after MCP initialization is complete to enable JSON-RPC notifications
+   */
+  public enableNotifications() {
+    this.isInitialized = true;
+    process.stderr.write(`[desktop-commander] JSON-RPC notifications enabled\n`);
+  }
+
+  /**
+   * Check if notifications are enabled
+   */
+  public get isNotificationsEnabled(): boolean {
+    return this.isInitialized;
+  }
+
   private setupConsoleRedirection() {
     console.log = (...args: any[]) => {
-      this.sendLogNotification("info", args);
+      if (this.isInitialized) {
+        this.sendLogNotification("info", args);
+      } else {
+        // During initialization, send to stderr to avoid protocol violations
+        process.stderr.write(`[LOG] ${args.join(' ')}\n`);
+      }
     };
 
     console.info = (...args: any[]) => {
-      this.sendLogNotification("info", args);
+      if (this.isInitialized) {
+        this.sendLogNotification("info", args);
+      } else {
+        process.stderr.write(`[INFO] ${args.join(' ')}\n`);
+      }
     };
 
     console.warn = (...args: any[]) => {
-      this.sendLogNotification("warning", args);
+      if (this.isInitialized) {
+        this.sendLogNotification("warning", args);
+      } else {
+        process.stderr.write(`[WARN] ${args.join(' ')}\n`);
+      }
     };
 
     console.error = (...args: any[]) => {
-      this.sendLogNotification("error", args);
+      if (this.isInitialized) {
+        this.sendLogNotification("error", args);
+      } else {
+        process.stderr.write(`[ERROR] ${args.join(' ')}\n`);
+      }
     };
 
     console.debug = (...args: any[]) => {
-      this.sendLogNotification("debug", args);
+      if (this.isInitialized) {
+        this.sendLogNotification("debug", args);
+      } else {
+        process.stderr.write(`[DEBUG] ${args.join(' ')}\n`);
+      }
     };
   }
 
@@ -87,7 +124,12 @@ export class FilteredStdioServerTransport extends StdioServerTransport {
           return this.originalStdoutWrite.call(process.stdout, buffer, encoding, callback);
         } else if (trimmed.length > 0) {
           // Non-JSON-RPC output, wrap it in a log notification
-          this.sendLogNotification("info", [buffer.replace(/\n$/, '')]);
+          if (this.isInitialized) {
+            this.sendLogNotification("info", [buffer.replace(/\n$/, '')]);
+          } else {
+            // During initialization, send to stderr
+            process.stderr.write(`[STDOUT] ${buffer}`);
+          }
           if (callback) callback();
           return true;
         }
diff --git a/src/index.ts b/src/index.ts
index fbc51cc8..009d8a9a 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -113,6 +113,17 @@ async function runServer() {
 
 
     console.error("Connecting server...");
+
+    // Set up event-driven initialization completion handler
+    server.oninitialized = () => {
+      // This callback is triggered after the client sends the "initialized" notification
+      // At this point, the MCP protocol handshake is fully complete
+      transport.enableNotifications();
+      process.stderr.write(
+        `[desktop-commander] MCP fully initialized, notifications enabled\n`
+      );
+    };
+
     await server.connect(transport);
     console.error("Server connected successfully");
   } catch (error) {
diff --git a/src/server.ts b/src/server.ts
index dcde304a..6ffd44d0 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -97,7 +97,8 @@ server.setRequestHandler(InitializeRequestSchema, async (request: InitializeRequ
                 name: clientInfo.name || 'unknown',
                 version: clientInfo.version || 'unknown'
             };
-            console.log(`Client connected: ${currentClient.name} v${currentClient.version}`);
+            // Log to stderr during initialization to avoid protocol violations
+            process.stderr.write(`[desktop-commander] Client connected: ${currentClient.name} v${currentClient.version}\n`);
         }
 
         // Return standard initialization response