@office-ai/aioncli-core 0.2.2 → 0.8.1

package/dist/src/ide/ide-client.js

@@ -5,14 +5,19 @@
  */
 import * as fs from 'node:fs';
 import { isSubpath } from '../utils/paths.js';
-import { detectIde, DetectedIde, getIdeInfo } from '../ide/detect-ide.js';
-import { ideContext, IdeContextNotificationSchema, IdeDiffAcceptedNotificationSchema, IdeDiffClosedNotificationSchema, CloseDiffResponseSchema, } from '../ide/ideContext.js';
-import { getIdeProcessId } from './process-utils.js';
+import { detectIde } from '../ide/detect-ide.js';
+import { ideContextStore } from './ideContext.js';
+import { IdeContextNotificationSchema, IdeDiffAcceptedNotificationSchema, IdeDiffClosedNotificationSchema, IdeDiffRejectedNotificationSchema, } from './types.js';
+import { getIdeProcessInfo } from './process-utils.js';
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
+import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
+import { CallToolResultSchema } from '@modelcontextprotocol/sdk/types.js';
 import * as os from 'node:os';
 import * as path from 'node:path';
 import { EnvHttpProxyAgent } from 'undici';
+import { ListToolsResultSchema } from '@modelcontextprotocol/sdk/types.js';
+import { IDE_REQUEST_TIMEOUT_MS } from './constants.js';
 const logger = {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     debug: (...args) => console.debug('[DEBUG] [IDEClient]', ...args),
@@ -39,27 +44,38 @@ function getRealPath(path) {
  * Manages the connection to and interaction with the IDE server.
  */
 export class IdeClient {
-    static instance;
+    static instancePromise = null;
     client = undefined;
     state = {
         status: IDEConnectionStatus.Disconnected,
         details: 'IDE integration is currently disabled. To enable it, run /ide enable.',
     };
     currentIde;
-    currentIdeDisplayName;
+    ideProcessInfo;
+    connectionConfig;
+    authToken;
     diffResponses = new Map();
     statusListeners = new Set();
-    constructor() {
-        this.currentIde = detectIde();
-        if (this.currentIde) {
-            this.currentIdeDisplayName = getIdeInfo(this.currentIde).displayName;
-        }
-    }
+    trustChangeListeners = new Set();
+    availableTools = [];
+    /**
+     * A mutex to ensure that only one diff view is open in the IDE at a time.
+     * This prevents race conditions and UI issues in IDEs like VSCode that
+     * can't handle multiple diff views being opened simultaneously.
+     */
+    diffMutex = Promise.resolve();
+    constructor() { }
     static getInstance() {
-        if (!IdeClient.instance) {
-            IdeClient.instance = new IdeClient();
+        if (!IdeClient.instancePromise) {
+            IdeClient.instancePromise = (async () => {
+                const client = new IdeClient();
+                client.ideProcessInfo = await getIdeProcessInfo();
+                client.connectionConfig = await client.getConnectionConfigFromFile();
+                client.currentIde = detectIde(client.ideProcessInfo, client.connectionConfig?.ideInfo);
+                return client;
+            })();
         }
-        return IdeClient.instance;
+        return IdeClient.instancePromise;
     }
     addStatusChangeListener(listener) {
         this.statusListeners.add(listener);
@@ -67,92 +83,195 @@ export class IdeClient {
     removeStatusChangeListener(listener) {
         this.statusListeners.delete(listener);
     }
+    addTrustChangeListener(listener) {
+        this.trustChangeListeners.add(listener);
+    }
+    removeTrustChangeListener(listener) {
+        this.trustChangeListeners.delete(listener);
+    }
     async connect() {
-        if (!this.currentIde || !this.currentIdeDisplayName) {
-            this.setState(IDEConnectionStatus.Disconnected, `IDE integration is not supported in your current environment. To use this feature, run Gemini CLI in one of these supported IDEs: ${Object.values(DetectedIde)
-                .map((ide) => getIdeInfo(ide).displayName)
-                .join(', ')}`, false);
+        if (!this.currentIde) {
+            this.setState(IDEConnectionStatus.Disconnected, `IDE integration is not supported in your current environment. To use this feature, run Gemini CLI in one of these supported IDEs: VS Code or VS Code forks`, false);
             return;
         }
         this.setState(IDEConnectionStatus.Connecting);
-        const ideInfoFromFile = await this.getIdeInfoFromFile();
-        const workspacePath = ideInfoFromFile.workspacePath ??
+        this.connectionConfig = await this.getConnectionConfigFromFile();
+        if (this.connectionConfig?.authToken) {
+            this.authToken = this.connectionConfig.authToken;
+        }
+        const workspacePath = this.connectionConfig?.workspacePath ??
             process.env['GEMINI_CLI_IDE_WORKSPACE_PATH'];
-        const { isValid, error } = IdeClient.validateWorkspacePath(workspacePath, this.currentIdeDisplayName, process.cwd());
+        const { isValid, error } = IdeClient.validateWorkspacePath(workspacePath, process.cwd());
         if (!isValid) {
             this.setState(IDEConnectionStatus.Disconnected, error, true);
             return;
         }
-        const portFromFile = ideInfoFromFile.port;
-        if (portFromFile) {
-            const connected = await this.establishConnection(portFromFile);
-            if (connected) {
-                return;
+        if (this.connectionConfig) {
+            if (this.connectionConfig.port) {
+                const connected = await this.establishHttpConnection(this.connectionConfig.port);
+                if (connected) {
+                    return;
+                }
+            }
+            if (this.connectionConfig.stdio) {
+                const connected = await this.establishStdioConnection(this.connectionConfig.stdio);
+                if (connected) {
+                    return;
+                }
             }
         }
         const portFromEnv = this.getPortFromEnv();
         if (portFromEnv) {
-            const connected = await this.establishConnection(portFromEnv);
+            const connected = await this.establishHttpConnection(portFromEnv);
+            if (connected) {
+                return;
+            }
+        }
+        const stdioConfigFromEnv = this.getStdioConfigFromEnv();
+        if (stdioConfigFromEnv) {
+            const connected = await this.establishStdioConnection(stdioConfigFromEnv);
             if (connected) {
                 return;
             }
         }
-        this.setState(IDEConnectionStatus.Disconnected, `Failed to connect to IDE companion extension for ${this.currentIdeDisplayName}. Please ensure the extension is running. To install the extension, run /ide install.`, true);
+        this.setState(IDEConnectionStatus.Disconnected, `Failed to connect to IDE companion extension in ${this.currentIde.displayName}. Please ensure the extension is running. To install the extension, run /ide install.`, true);
     }
     /**
-     * A diff is accepted with any modifications if the user performs one of the
-     * following actions:
-     * - Clicks the checkbox icon in the IDE to accept
-     * - Runs `command+shift+p` > "Gemini CLI: Accept Diff in IDE" to accept
-     * - Selects "accept" in the CLI UI
-     * - Saves the file via `ctrl/command+s`
+     * Opens a diff view in the IDE, allowing the user to review and accept or
+     * reject changes.
      *
-     * A diff is rejected if the user performs one of the following actions:
-     * - Clicks the "x" icon in the IDE
-     * - Runs "Gemini CLI: Close Diff in IDE"
-     * - Selects "no" in the CLI UI
-     * - Closes the file
+     * This method sends a request to the IDE to display a diff between the
+     * current content of a file and the new content provided. It then waits for
+     * a notification from the IDE indicating that the user has either accepted
+     * (potentially with manual edits) or rejected the diff.
+     *
+     * A mutex ensures that only one diff view can be open at a time to prevent
+     * race conditions.
+     *
+     * @param filePath The absolute path to the file to be diffed.
+     * @param newContent The proposed new content for the file.
+     * @returns A promise that resolves with a `DiffUpdateResult`, indicating
+     *   whether the diff was 'accepted' or 'rejected' and including the final
+     *   content if accepted.
      */
     async openDiff(filePath, newContent) {
-        return new Promise((resolve, reject) => {
+        const release = await this.acquireMutex();
+        const promise = new Promise((resolve, reject) => {
+            if (!this.client) {
+                // The promise will be rejected, and the finally block below will release the mutex.
+                return reject(new Error('IDE client is not connected.'));
+            }
             this.diffResponses.set(filePath, resolve);
             this.client
-                ?.callTool({
-                name: `openDiff`,
-                arguments: {
-                    filePath,
-                    newContent,
+                .request({
+                method: 'tools/call',
+                params: {
+                    name: `openDiff`,
+                    arguments: {
+                        filePath,
+                        newContent,
+                    },
                 },
+            }, CallToolResultSchema, { timeout: IDE_REQUEST_TIMEOUT_MS })
+                .then((parsedResultData) => {
+                if (parsedResultData.isError) {
+                    const textPart = parsedResultData.content.find((part) => part.type === 'text');
+                    const errorMessage = textPart?.text ?? `Tool 'openDiff' reported an error.`;
+                    logger.debug(`Request for openDiff ${filePath} failed with isError:`, errorMessage);
+                    this.diffResponses.delete(filePath);
+                    reject(new Error(errorMessage));
+                }
             })
                 .catch((err) => {
-                logger.debug(`callTool for ${filePath} failed:`, err);
+                logger.debug(`Request for openDiff ${filePath} failed:`, err);
+                this.diffResponses.delete(filePath);
                 reject(err);
             });
         });
+        // Ensure the mutex is released only after the diff interaction is complete.
+        promise.finally(release);
+        return promise;
     }
-    async closeDiff(filePath) {
+    /**
+     * Acquires a lock to ensure sequential execution of critical sections.
+     *
+     * This method implements a promise-based mutex. It works by chaining promises.
+     * Each call to `acquireMutex` gets the current `diffMutex` promise. It then
+     * creates a *new* promise (`newMutex`) that will be resolved when the caller
+     * invokes the returned `release` function. The `diffMutex` is immediately
+     * updated to this `newMutex`.
+     *
+     * The method returns a promise that resolves with the `release` function only
+     * *after* the *previous* `diffMutex` promise has resolved. This creates a
+     * queue where each subsequent operation must wait for the previous one to release
+     * the lock.
+     *
+     * @returns A promise that resolves to a function that must be called to
+     *   release the lock.
+     */
+    acquireMutex() {
+        let release;
+        const newMutex = new Promise((resolve) => {
+            release = resolve;
+        });
+        const oldMutex = this.diffMutex;
+        this.diffMutex = newMutex;
+        return oldMutex.then(() => release);
+    }
+    async closeDiff(filePath, options) {
         try {
-            const result = await this.client?.callTool({
-                name: `closeDiff`,
-                arguments: {
-                    filePath,
+            if (!this.client) {
+                return undefined;
+            }
+            const resultData = await this.client.request({
+                method: 'tools/call',
+                params: {
+                    name: `closeDiff`,
+                    arguments: {
+                        filePath,
+                        suppressNotification: options?.suppressNotification,
+                    },
                 },
-            });
-            if (result) {
-                const parsed = CloseDiffResponseSchema.parse(result);
-                return parsed.content;
+            }, CallToolResultSchema, { timeout: IDE_REQUEST_TIMEOUT_MS });
+            if (!resultData) {
+                return undefined;
+            }
+            if (resultData.isError) {
+                const textPart = resultData.content.find((part) => part.type === 'text');
+                const errorMessage = textPart?.text ?? `Tool 'closeDiff' reported an error.`;
+                logger.debug(`Request for closeDiff ${filePath} failed with isError:`, errorMessage);
+                return undefined;
+            }
+            const textPart = resultData.content.find((part) => part.type === 'text');
+            if (textPart?.text) {
+                try {
+                    const parsedJson = JSON.parse(textPart.text);
+                    if (parsedJson && typeof parsedJson.content === 'string') {
+                        return parsedJson.content;
+                    }
+                    if (parsedJson && parsedJson.content === null) {
+                        return undefined;
+                    }
+                }
+                catch (_e) {
+                    logger.debug(`Invalid JSON in closeDiff response for ${filePath}:`, textPart.text);
+                }
             }
         }
         catch (err) {
-            logger.debug(`callTool for ${filePath} failed:`, err);
+            logger.debug(`Request for closeDiff ${filePath} failed:`, err);
         }
-        return;
+        return undefined;
     }
     // Closes the diff. Instead of waiting for a notification,
     // manually resolves the diff resolver as the desired outcome.
     async resolveDiffFromCli(filePath, outcome) {
-        const content = await this.closeDiff(filePath);
         const resolver = this.diffResponses.get(filePath);
+        const content = await this.closeDiff(filePath, {
+            // Suppress notification to avoid race where closing the diff rejects the
+            // request.
+            suppressNotification: true,
+        });
         if (resolver) {
             if (outcome === 'accepted') {
                 resolver({ status: 'accepted', content });
@@ -181,7 +300,42 @@ export class IdeClient {
         return this.state;
     }
     getDetectedIdeDisplayName() {
-        return this.currentIdeDisplayName;
+        return this.currentIde?.displayName;
+    }
+    isDiffingEnabled() {
+        return (!!this.client &&
+            this.state.status === IDEConnectionStatus.Connected &&
+            this.availableTools.includes('openDiff') &&
+            this.availableTools.includes('closeDiff'));
+    }
+    async discoverTools() {
+        if (!this.client) {
+            return;
+        }
+        try {
+            logger.debug('Discovering tools from IDE...');
+            const response = await this.client.request({ method: 'tools/list', params: {} }, ListToolsResultSchema);
+            // Map the array of tool objects to an array of tool names (strings)
+            this.availableTools = response.tools.map((tool) => tool.name);
+            if (this.availableTools.length > 0) {
+                logger.debug(`Discovered ${this.availableTools.length} tools from IDE: ${this.availableTools.join(', ')}`);
+            }
+            else {
+                logger.debug('IDE supports tool discovery, but no tools are available.');
+            }
+        }
+        catch (error) {
+            // It's okay if this fails, the IDE might not support it.
+            // Don't log an error if the method is not found, which is a common case.
+            if (error instanceof Error &&
+                !error.message?.includes('Method not found')) {
+                logger.error(`Error discovering tools from IDE: ${error.message}`);
+            }
+            else {
+                logger.debug('IDE does not support tool discovery.');
+            }
+            this.availableTools = [];
+        }
     }
     setState(status, details, logToConsole = false) {
         const isAlreadyDisconnected = this.state.status === IDEConnectionStatus.Disconnected &&
@@ -205,20 +359,20 @@ export class IdeClient {
             }
         }
         if (status === IDEConnectionStatus.Disconnected) {
-            ideContext.clearIdeContext();
+            ideContextStore.clear();
         }
     }
-    static validateWorkspacePath(ideWorkspacePath, currentIdeDisplayName, cwd) {
+    static validateWorkspacePath(ideWorkspacePath, cwd) {
         if (ideWorkspacePath === undefined) {
             return {
                 isValid: false,
-                error: `Failed to connect to IDE companion extension for ${currentIdeDisplayName}. Please ensure the extension is running. To install the extension, run /ide install.`,
+                error: `Failed to connect to IDE companion extension. Please ensure the extension is running. To install the extension, run /ide install.`,
             };
         }
         if (ideWorkspacePath === '') {
             return {
                 isValid: false,
-                error: `To use this feature, please open a workspace folder in ${currentIdeDisplayName} and try again.`,
+                error: `To use this feature, please open a workspace folder in your IDE and try again.`,
             };
         }
         const ideWorkspacePaths = ideWorkspacePath.split(path.delimiter);
@@ -230,7 +384,7 @@ export class IdeClient {
         if (!isWithinWorkspace) {
             return {
                 isValid: false,
-                error: `Directory mismatch. Gemini CLI is running in a different location than the open workspace in ${currentIdeDisplayName}. Please run the CLI from one of the following directories: ${ideWorkspacePaths.join(', ')}`,
+                error: `Directory mismatch. Gemini CLI is running in a different location than the open workspace in the IDE. Please run the CLI from one of the following directories: ${ideWorkspacePaths.join(', ')}`,
             };
         }
         return { isValid: true };
@@ -242,20 +396,102 @@ export class IdeClient {
         }
         return port;
     }
-    async getIdeInfoFromFile() {
+    getStdioConfigFromEnv() {
+        const command = process.env['GEMINI_CLI_IDE_SERVER_STDIO_COMMAND'];
+        if (!command) {
+            return undefined;
+        }
+        const argsStr = process.env['GEMINI_CLI_IDE_SERVER_STDIO_ARGS'];
+        let args = [];
+        if (argsStr) {
+            try {
+                const parsedArgs = JSON.parse(argsStr);
+                if (Array.isArray(parsedArgs)) {
+                    args = parsedArgs;
+                }
+                else {
+                    logger.error('GEMINI_CLI_IDE_SERVER_STDIO_ARGS must be a JSON array string.');
+                }
+            }
+            catch (e) {
+                logger.error('Failed to parse GEMINI_CLI_IDE_SERVER_STDIO_ARGS:', e);
+            }
+        }
+        return { command, args };
+    }
+    async getConnectionConfigFromFile() {
+        if (!this.ideProcessInfo) {
+            return undefined;
+        }
+        // For backwards compatability
         try {
-            const ideProcessId = await getIdeProcessId();
-            const portFile = path.join(os.tmpdir(), `gemini-ide-server-${ideProcessId}.json`);
+            const portFile = path.join(os.tmpdir(), `gemini-ide-server-${this.ideProcessInfo.pid}.json`);
             const portFileContents = await fs.promises.readFile(portFile, 'utf8');
-            const ideInfo = JSON.parse(portFileContents);
-            return {
-                port: ideInfo?.port?.toString(),
-                workspacePath: ideInfo?.workspacePath,
-            };
+            return JSON.parse(portFileContents);
         }
         catch (_) {
-            return {};
+            // For newer extension versions, the file name matches the pattern
+            // /^gemini-ide-server-${pid}-\d+\.json$/. If multiple IDE
+            // windows are open, multiple files matching the pattern are expected to
+            // exist.
+        }
+        const portFileDir = path.join(os.tmpdir(), 'gemini', 'ide');
+        let portFiles;
+        try {
+            portFiles = await fs.promises.readdir(portFileDir);
+        }
+        catch (e) {
+            logger.debug('Failed to read IDE connection directory:', e);
+            return undefined;
+        }
+        if (!portFiles) {
+            return undefined;
+        }
+        const fileRegex = new RegExp(`^gemini-ide-server-${this.ideProcessInfo.pid}-\\d+\\.json$`);
+        const matchingFiles = portFiles
+            .filter((file) => fileRegex.test(file))
+            .sort();
+        if (matchingFiles.length === 0) {
+            return undefined;
+        }
+        let fileContents;
+        try {
+            fileContents = await Promise.all(matchingFiles.map((file) => fs.promises.readFile(path.join(portFileDir, file), 'utf8')));
+        }
+        catch (e) {
+            logger.debug('Failed to read IDE connection config file(s):', e);
+            return undefined;
         }
+        const parsedContents = fileContents.map((content) => {
+            try {
+                return JSON.parse(content);
+            }
+            catch (e) {
+                logger.debug('Failed to parse JSON from config file: ', e);
+                return undefined;
+            }
+        });
+        const validWorkspaces = parsedContents.filter((content) => {
+            if (!content) {
+                return false;
+            }
+            const { isValid } = IdeClient.validateWorkspacePath(content.workspacePath, process.cwd());
+            return isValid;
+        });
+        if (validWorkspaces.length === 0) {
+            return undefined;
+        }
+        if (validWorkspaces.length === 1) {
+            return validWorkspaces[0];
+        }
+        const portFromEnv = this.getPortFromEnv();
+        if (portFromEnv) {
+            const matchingPort = validWorkspaces.find((content) => String(content.port) === portFromEnv);
+            if (matchingPort) {
+                return matchingPort;
+            }
+        }
+        return validWorkspaces[0];
     }
     createProxyAwareFetch() {
         // ignore proxy for 'localhost' by deafult to allow connecting to the ide mcp server
@@ -275,7 +511,7 @@ export class IdeClient {
             return new Response(response.body, {
                 status: response.status,
                 statusText: response.statusText,
-                headers: new Headers(response.headers),
+                headers: [...response.headers.entries()],
             });
         };
     }
@@ -284,13 +520,20 @@ export class IdeClient {
             return;
         }
         this.client.setNotificationHandler(IdeContextNotificationSchema, (notification) => {
-            ideContext.setIdeContext(notification.params);
+            ideContextStore.set(notification.params);
+            const isTrusted = notification.params.workspaceState?.isTrusted;
+            if (isTrusted !== undefined) {
+                for (const listener of this.trustChangeListeners) {
+                    listener(isTrusted);
+                }
+            }
         });
         this.client.onerror = (_error) => {
-            this.setState(IDEConnectionStatus.Disconnected, `IDE connection error. The connection was lost unexpectedly. Please try reconnecting by running /ide enable`, true);
+            const errorMessage = _error instanceof Error ? _error.message : `_error`;
+            this.setState(IDEConnectionStatus.Disconnected, `IDE connection error. The connection was lost unexpectedly. Please try reconnecting by running /ide enable\n${errorMessage}`, true);
         };
         this.client.onclose = () => {
-            this.setState(IDEConnectionStatus.Disconnected, `IDE connection error. The connection was lost unexpectedly. Please try reconnecting by running /ide enable`, true);
+            this.setState(IDEConnectionStatus.Disconnected, `IDE connection closed. To reconnect, run /ide enable.`, true);
         };
         this.client.setNotificationHandler(IdeDiffAcceptedNotificationSchema, (notification) => {
             const { filePath, content } = notification.params;
@@ -303,6 +546,19 @@ export class IdeClient {
                 logger.debug(`No resolver found for ${filePath}`);
             }
         });
+        this.client.setNotificationHandler(IdeDiffRejectedNotificationSchema, (notification) => {
+            const { filePath } = notification.params;
+            const resolver = this.diffResponses.get(filePath);
+            if (resolver) {
+                resolver({ status: 'rejected', content: undefined });
+                this.diffResponses.delete(filePath);
+            }
+            else {
+                logger.debug(`No resolver found for ${filePath}`);
+            }
+        });
+        // For backwards compatability. Newer extension versions will only send
+        // IdeDiffRejectedNotificationSchema.
         this.client.setNotificationHandler(IdeDiffClosedNotificationSchema, (notification) => {
             const { filePath } = notification.params;
             const resolver = this.diffResponses.get(filePath);
@@ -315,9 +571,10 @@ export class IdeClient {
             }
         });
     }
-    async establishConnection(port) {
+    async establishHttpConnection(port) {
         let transport;
         try {
+            logger.debug('Attempting to connect to IDE via HTTP SSE');
             this.client = new Client({
                 name: 'streamable-http-client',
                 // TODO(#3487): use the CLI version here.
@@ -325,9 +582,46 @@ export class IdeClient {
             });
             transport = new StreamableHTTPClientTransport(new URL(`http://${getIdeServerHost()}:${port}/mcp`), {
                 fetch: this.createProxyAwareFetch(),
+                requestInit: {
+                    headers: this.authToken
+                        ? { Authorization: `Bearer ${this.authToken}` }
+                        : {},
+                },
+            });
+            await this.client.connect(transport);
+            this.registerClientHandlers();
+            await this.discoverTools();
+            this.setState(IDEConnectionStatus.Connected);
+            return true;
+        }
+        catch (_error) {
+            if (transport) {
+                try {
+                    await transport.close();
+                }
+                catch (closeError) {
+                    logger.debug('Failed to close transport:', closeError);
+                }
+            }
+            return false;
+        }
+    }
+    async establishStdioConnection({ command, args, }) {
+        let transport;
+        try {
+            logger.debug('Attempting to connect to IDE via stdio');
+            this.client = new Client({
+                name: 'stdio-client',
+                // TODO(#3487): use the CLI version here.
+                version: '1.0.0',
+            });
+            transport = new StdioClientTransport({
+                command,
+                args,
             });
             await this.client.connect(transport);
             this.registerClientHandlers();
+            await this.discoverTools();
             this.setState(IDEConnectionStatus.Connected);
             return true;
         }