diff options
| -rw-r--r-- | integration-tests/ide-client.test.ts | 139 | ||||
| -rw-r--r-- | packages/core/src/ide/ide-client.ts | 56 | ||||
| -rw-r--r-- | packages/core/src/ide/process-utils.ts | 62 | ||||
| -rw-r--r-- | packages/vscode-ide-companion/src/ide-server.ts | 17 |
4 files changed, 258 insertions, 16 deletions
diff --git a/integration-tests/ide-client.test.ts b/integration-tests/ide-client.test.ts new file mode 100644 index 00000000..186320b3 --- /dev/null +++ b/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); +}); diff --git a/packages/core/src/ide/ide-client.ts b/packages/core/src/ide/ide-client.ts index fe605eb2..94107f21 100644 --- a/packages/core/src/ide/ide-client.ts +++ b/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; } } } diff --git a/packages/core/src/ide/process-utils.ts b/packages/core/src/ide/process-utils.ts new file mode 100644 index 00000000..40e16a73 --- /dev/null +++ b/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; +} diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts index eec99cb3..ee77bdb8 100644 --- a/packages/vscode-ide-companion/src/ide-server.ts +++ b/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. + } } } |