Fix IDE Companion Connection in Proxy Environments (#6308)

Co-authored-by: Jacob Richman <jacob314@gmail.com>

.vscode/launch.json

@@ -67,6 +67,27 @@
       "console": "integratedTerminal",
       "internalConsoleOptions": "neverOpen",
       "skipFiles": ["<node_internals>/**"]
+    },
+    {
+      "name": "Debug Integration Test File",
+      "type": "node",
+      "request": "launch",
+      "runtimeExecutable": "npx",
+      "runtimeArgs": [
+        "vitest",
+        "run",
+        "--root",
+        "./integration-tests",
+        "--inspect-brk=9229",
+        "${file}"
+      ],
+      "cwd": "${workspaceFolder}",
+      "console": "integratedTerminal",
+      "internalConsoleOptions": "neverOpen",
+      "skipFiles": ["<node_internals>/**"],
+      "env": {
+        "GEMINI_SANDBOX": "false"
+      }
     }
   ],
   "inputs": [

integration-tests/ide-client.test.ts

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, beforeEach, afterEach } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
 import * as fs from 'node:fs';
 import * as os from 'node:os';
 import * as path from 'node:path';
@@ -35,8 +35,6 @@ describe.skip('IdeClient', () => {
     fs.unlinkSync(portFile);
     await server.stop();
     delete process.env['GEMINI_CLI_IDE_WORKSPACE_PATH'];
-    // Reset instance
-    IdeClient.instance = undefined;
   });
 });
 
@@ -156,3 +154,46 @@ describe.skip('getIdeProcessId', () => {
     expect(parseInt(output, 10)).toBe(parentPid);
   }, 10000);
 });
+
+describe('IdeClient with proxy', () => {
+  let mcpServer: TestMcpServer;
+  let proxyServer: net.Server;
+  let mcpServerPort: number;
+  let proxyServerPort: number;
+
+  beforeEach(async () => {
+    mcpServer = new TestMcpServer();
+    mcpServerPort = await mcpServer.start();
+
+    proxyServer = net.createServer().listen();
+    proxyServerPort = (proxyServer.address() as net.AddressInfo).port;
+
+    vi.stubEnv('GEMINI_CLI_IDE_SERVER_PORT', String(mcpServerPort));
+    vi.stubEnv('TERM_PROGRAM', 'vscode');
+    vi.stubEnv('GEMINI_CLI_IDE_WORKSPACE_PATH', process.cwd());
+
+    // Reset instance
+    IdeClient.instance = undefined;
+  });
+
+  afterEach(async () => {
+    IdeClient.getInstance().disconnect();
+    await mcpServer.stop();
+    proxyServer.close();
+    vi.unstubAllEnvs();
+  });
+
+  it('should connect to IDE server when HTTP_PROXY, HTTPS_PROXY and NO_PROXY are set', async () => {
+    vi.stubEnv('HTTP_PROXY', `http://localhost:${proxyServerPort}`);
+    vi.stubEnv('HTTPS_PROXY', `http://localhost:${proxyServerPort}`);
+    vi.stubEnv('NO_PROXY', 'example.com,127.0.0.1,::1');
+
+    const ideClient = IdeClient.getInstance();
+    await ideClient.connect();
+
+    expect(ideClient.getConnectionStatus()).toEqual({
+      status: 'connected',
+      details: undefined,
+    });
+  });
+});

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

@@ -20,6 +20,7 @@ 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';
+import { EnvHttpProxyAgent } from 'undici';
 
 const logger = {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -325,6 +326,29 @@ export class IdeClient {
     }
   }
 
+  private createProxyAwareFetch() {
+    // ignore proxy for 'localhost' by deafult to allow connecting to the ide mcp server
+    const existingNoProxy = process.env['NO_PROXY'] || '';
+    const agent = new EnvHttpProxyAgent({
+      noProxy: [existingNoProxy, 'localhost'].filter(Boolean).join(','),
+    });
+    const undiciPromise = import('undici');
+    return async (url: string | URL, init?: RequestInit): Promise<Response> => {
+      const { fetch: fetchFn } = await undiciPromise;
+      const fetchOptions: RequestInit & { dispatcher?: unknown } = {
+        ...init,
+        dispatcher: agent,
+      };
+      const options = fetchOptions as unknown as import('undici').RequestInit;
+      const response = await fetchFn(url, options);
+      return new Response(response.body as ReadableStream<unknown> | null, {
+        status: response.status,
+        statusText: response.statusText,
+        headers: response.headers,
+      });
+    };
+  }
+
   private registerClientHandlers() {
     if (!this.client) {
       return;
@@ -389,6 +413,9 @@ export class IdeClient {
       });
       transport = new StreamableHTTPClientTransport(
         new URL(`http://${getIdeServerHost()}:${port}/mcp`),
+        {
+          fetch: this.createProxyAwareFetch(),
+        },
       );
       await this.client.connect(transport);
       this.registerClientHandlers();