[ide-mode] Write port to file in ide-server (#5811)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

integration-tests/ide-client.test.ts

@@ -0,0 +1,139 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as fs from 'node:fs';
+import * as os from 'node:os';
+import * as path from 'node:path';
+import * as net from 'node:net';
+import { IdeClient } from '../packages/core/src/ide/ide-client.js';
+import { getIdeProcessId } from '../packages/core/src/ide/process-utils.js';
+import { spawn, ChildProcess } from 'child_process';
+
+describe('IdeClient', () => {
+  it('reads port from file and connects', async () => {
+    const port = 12345;
+    const pid = await getIdeProcessId();
+    const portFile = path.join(os.tmpdir(), `gemini-ide-server-${pid}.json`);
+    fs.writeFileSync(portFile, JSON.stringify({ port }));
+
+    const ideClient = IdeClient.getInstance();
+    await ideClient.connect();
+
+    expect(ideClient.getConnectionStatus().status).not.toBe('disconnected');
+
+    fs.unlinkSync(portFile);
+  });
+});
+
+const getFreePort = (): Promise<number> => {
+  return new Promise((resolve, reject) => {
+    const server = net.createServer();
+    server.unref();
+    server.on('error', reject);
+    server.listen(0, () => {
+      const port = (server.address() as net.AddressInfo).port;
+      server.close(() => {
+        resolve(port);
+      });
+    });
+  });
+};
+
+describe('IdeClient fallback connection logic', () => {
+  let server: net.Server;
+  let envPort: number;
+  let pid: number;
+  let portFile: string;
+
+  beforeEach(async () => {
+    pid = await getIdeProcessId();
+    portFile = path.join(os.tmpdir(), `gemini-ide-server-${pid}.json`);
+    envPort = await getFreePort();
+    server = net.createServer().listen(envPort);
+    process.env['GEMINI_CLI_IDE_SERVER_PORT'] = String(envPort);
+    process.env['GEMINI_CLI_IDE_WORKSPACE_PATH'] = process.cwd();
+    // Reset instance
+    IdeClient.instance = undefined;
+  });
+
+  afterEach(() => {
+    server.close();
+    delete process.env['GEMINI_CLI_IDE_SERVER_PORT'];
+    delete process.env['GEMINI_CLI_IDE_WORKSPACE_PATH'];
+    if (fs.existsSync(portFile)) {
+      fs.unlinkSync(portFile);
+    }
+  });
+
+  it('connects using env var when port file does not exist', async () => {
+    // Ensure port file doesn't exist
+    if (fs.existsSync(portFile)) {
+      fs.unlinkSync(portFile);
+    }
+
+    const ideClient = IdeClient.getInstance();
+    await ideClient.connect();
+
+    expect(ideClient.getConnectionStatus().status).toBe('connected');
+  });
+
+  it('falls back to env var when connection with port from file fails', async () => {
+    const filePort = await getFreePort();
+    // Write port file with a port that is not listening
+    fs.writeFileSync(portFile, JSON.stringify({ port: filePort }));
+
+    const ideClient = IdeClient.getInstance();
+    await ideClient.connect();
+
+    expect(ideClient.getConnectionStatus().status).toBe('connected');
+  });
+});
+
+describe('getIdeProcessId', () => {
+  let child: ChildProcess;
+
+  afterEach(() => {
+    if (child) {
+      child.kill();
+    }
+  });
+
+  it('should return the pid of the parent process', async () => {
+    // We need to spawn a child process that will run the test
+    // so that we can check that getIdeProcessId returns the pid of the parent
+    const parentPid = process.pid;
+    const output = await new Promise<string>((resolve, reject) => {
+      child = spawn(
+        'node',
+        [
+          '-e',
+          `
+        const { getIdeProcessId } = require('../packages/core/src/ide/process-utils.js');
+        getIdeProcessId().then(pid => console.log(pid));
+      `,
+        ],
+        {
+          stdio: ['pipe', 'pipe', 'pipe'],
+        },
+      );
+
+      let out = '';
+      child.stdout?.on('data', (data) => {
+        out += data.toString();
+      });
+
+      child.on('close', (code) => {
+        if (code === 0) {
+          resolve(out.trim());
+        } else {
+          reject(new Error(`Child process exited with code ${code}`));
+        }
+      });
+    });
+
+    expect(parseInt(output, 10)).toBe(parentPid);
+  }, 10000);
+});

packages/core/src/ide/ide-client.ts

@@ -5,7 +5,6 @@
  */
 
 import * as fs from 'node:fs';
-import * as path from 'node:path';
 import { detectIde, DetectedIde, getIdeInfo } from '../ide/detect-ide.js';
 import {
   ideContext,
@@ -15,8 +14,11 @@ import {
   CloseDiffResponseSchema,
   DiffUpdateResult,
 } from '../ide/ideContext.js';
+import { getIdeProcessId } from './process-utils.js';
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
+import * as os from 'node:os';
+import * as path from 'node:path';
 
 const logger = {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -95,12 +97,27 @@ export class IdeClient {
       return;
     }
 
-    const port = this.getPortFromEnv();
-    if (!port) {
-      return;
+    const portFromFile = await this.getPortFromFile();
+    if (portFromFile) {
+      const connected = await this.establishConnection(portFromFile);
+      if (connected) {
+        return;
+      }
     }
 
-    await this.establishConnection(port);
+    const portFromEnv = this.getPortFromEnv();
+    if (portFromEnv) {
+      const connected = await this.establishConnection(portFromEnv);
+      if (connected) {
+        return;
+      }
+    }
+
+    this.setState(
+      IDEConnectionStatus.Disconnected,
+      `Failed to connect to IDE companion extension for ${this.currentIdeDisplayName}. Please ensure the extension is running and try restarting your terminal. To install the extension, run /ide install.`,
+      true,
+    );
   }
 
   /**
@@ -264,16 +281,26 @@ export class IdeClient {
   private getPortFromEnv(): string | undefined {
     const port = process.env['GEMINI_CLI_IDE_SERVER_PORT'];
     if (!port) {
-      this.setState(
-        IDEConnectionStatus.Disconnected,
-        `Failed to connect to IDE companion extension for ${this.currentIdeDisplayName}. Please ensure the extension is running and try restarting your terminal. To install the extension, run /ide install.`,
-        true,
-      );
       return undefined;
     }
     return port;
   }
 
+  private async getPortFromFile(): Promise<string | undefined> {
+    try {
+      const ideProcessId = await getIdeProcessId();
+      const portFile = path.join(
+        os.tmpdir(),
+        `gemini-ide-server-${ideProcessId}.json`,
+      );
+      const portFileContents = await fs.promises.readFile(portFile, 'utf8');
+      const port = JSON.parse(portFileContents).port;
+      return port.toString();
+    } catch (_) {
+      return undefined;
+    }
+  }
+
   private registerClientHandlers() {
     if (!this.client) {
       return;
@@ -328,7 +355,7 @@ export class IdeClient {
     );
   }
 
-  private async establishConnection(port: string) {
+  private async establishConnection(port: string): Promise<boolean> {
     let transport: StreamableHTTPClientTransport | undefined;
     try {
       this.client = new Client({
@@ -342,12 +369,8 @@ export class IdeClient {
       await this.client.connect(transport);
       this.registerClientHandlers();
       this.setState(IDEConnectionStatus.Connected);
+      return true;
     } catch (_error) {
-      this.setState(
-        IDEConnectionStatus.Disconnected,
-        `Failed to connect to IDE companion extension for ${this.currentIdeDisplayName}. Please ensure the extension is running and try restarting your terminal. To install the extension, run /ide install.`,
-        true,
-      );
       if (transport) {
         try {
           await transport.close();
@@ -355,6 +378,7 @@ export class IdeClient {
           logger.debug('Failed to close transport:', closeError);
         }
       }
+      return false;
     }
   }
 }

packages/core/src/ide/process-utils.ts

@@ -0,0 +1,62 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { exec } from 'child_process';
+import { promisify } from 'util';
+import os from 'os';
+
+const execAsync = promisify(exec);
+
+/**
+ * Traverses up the process tree from the current process to find the top-level ancestor process ID.
+ * This is useful for identifying the main application process that spawned the current script,
+ * such as the main VS Code window process.
+ *
+ * @returns A promise that resolves to the numeric PID of the top-level process.
+ * @throws Will throw an error if the underlying shell commands fail unexpectedly.
+ */
+export async function getIdeProcessId(): Promise<number> {
+  const platform = os.platform();
+  let currentPid = process.pid;
+
+  // Loop upwards through the process tree, with a depth limit to prevent infinite loops.
+  const MAX_TRAVERSAL_DEPTH = 32;
+  for (let i = 0; i < MAX_TRAVERSAL_DEPTH; i++) {
+    let parentPid: number;
+
+    try {
+      // Use wmic for Windows
+      if (platform === 'win32') {
+        const command = `wmic process where "ProcessId=${currentPid}" get ParentProcessId /value`;
+        const { stdout } = await execAsync(command);
+        const match = stdout.match(/ParentProcessId=(\d+)/);
+        parentPid = match ? parseInt(match[1], 10) : 0; // Top of the tree is 0
+      }
+      // Use ps for macOS, Linux, and other Unix-like systems
+      else {
+        const command = `ps -o ppid= -p ${currentPid}`;
+        const { stdout } = await execAsync(command);
+        const ppid = parseInt(stdout.trim(), 10);
+        parentPid = isNaN(ppid) ? 1 : ppid; // Top of the tree is 1
+      }
+    } catch (_) {
+      // This can happen if a process in the chain dies during execution.
+      // We'll break the loop and return the last valid PID we found.
+      break;
+    }
+
+    // Define the root PID for the current OS
+    const rootPid = platform === 'win32' ? 0 : 1;
+
+    // If the parent is the root process or invalid, we've found our target.
+    if (parentPid === rootPid || parentPid <= 0) {
+      break;
+    }
+    // Move one level up the tree for the next iteration.
+    currentPid = parentPid;
+  }
+  return currentPid;
+}

packages/vscode-ide-companion/src/ide-server.ts

@@ -12,6 +12,9 @@ import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/
 import express, { type Request, type Response } from 'express';
 import { randomUUID } from 'node:crypto';
 import { type Server as HTTPServer } from 'node:http';
+import * as path from 'node:path';
+import * as fs from 'node:fs/promises';
+import * as os from 'node:os';
 import { z } from 'zod';
 import { DiffManager } from './diff-manager.js';
 import { OpenFilesManager } from './open-files-manager.js';
@@ -46,11 +49,16 @@ export class IDEServer {
   private server: HTTPServer | undefined;
   private context: vscode.ExtensionContext | undefined;
   private log: (message: string) => void;
+  private portFile: string;
   diffManager: DiffManager;
 
   constructor(log: (message: string) => void, diffManager: DiffManager) {
     this.log = log;
     this.diffManager = diffManager;
+    this.portFile = path.join(
+      os.tmpdir(),
+      `gemini-ide-server-${process.ppid}.json`,
+    );
   }
 
   async start(context: vscode.ExtensionContext) {
@@ -197,6 +205,10 @@ export class IDEServer {
           port.toString(),
         );
         this.log(`IDE server listening on port ${port}`);
+        fs.writeFile(this.portFile, JSON.stringify({ port })).catch((err) => {
+          this.log(`Failed to write port to file: ${err}`);
+        });
+        this.log(this.portFile);
       }
     });
   }
@@ -219,6 +231,11 @@ export class IDEServer {
     if (this.context) {
       this.context.environmentVariableCollection.clear();
     }
+    try {
+      await fs.unlink(this.portFile);
+    } catch (_err) {
+      // Ignore errors if the file doesn't exist.
+    }
   }
 }