Skip to content

Commit 111c735

Browse files
galz10jackwotherspoonchrstnb
authored andcommitted
refactor(core): Improve environment variable handling in shell execution (google-gemini#14742)
Co-authored-by: Jack Wotherspoon <jackwoth@google.com> Co-authored-by: christine betts <chrstn@uw.edu>
1 parent 38a320c commit 111c735

File tree

2 files changed

+230
-3
lines changed

2 files changed

+230
-3
lines changed

packages/core/src/services/shellExecutionService.test.ts

Lines changed: 176 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,15 @@
44
* SPDX-License-Identifier: Apache-2.0
55
*/
66

7-
import { vi, describe, it, expect, beforeEach, type Mock } from 'vitest';
7+
import {
8+
vi,
9+
describe,
10+
it,
11+
expect,
12+
beforeEach,
13+
afterEach,
14+
type Mock,
15+
} from 'vitest';
816
import EventEmitter from 'node:events';
917
import type { Readable } from 'node:stream';
1018
import { type ChildProcess } from 'node:child_process';
@@ -1284,3 +1292,170 @@ describe('ShellExecutionService execution method selection', () => {
12841292
expect(result.executionMethod).toBe('child_process');
12851293
});
12861294
});
1295+
1296+
describe('ShellExecutionService environment variables', () => {
1297+
let mockPtyProcess: EventEmitter & {
1298+
pid: number;
1299+
kill: Mock;
1300+
onData: Mock;
1301+
onExit: Mock;
1302+
write: Mock;
1303+
resize: Mock;
1304+
};
1305+
let mockChildProcess: EventEmitter & Partial<ChildProcess>;
1306+
1307+
beforeEach(() => {
1308+
vi.clearAllMocks();
1309+
vi.resetModules(); // Reset modules to ensure process.env changes are fresh
1310+
1311+
// Mock for pty
1312+
mockPtyProcess = new EventEmitter() as EventEmitter & {
1313+
pid: number;
1314+
kill: Mock;
1315+
onData: Mock;
1316+
onExit: Mock;
1317+
write: Mock;
1318+
resize: Mock;
1319+
};
1320+
mockPtyProcess.pid = 12345;
1321+
mockPtyProcess.kill = vi.fn();
1322+
mockPtyProcess.onData = vi.fn();
1323+
mockPtyProcess.onExit = vi.fn();
1324+
mockPtyProcess.write = vi.fn();
1325+
mockPtyProcess.resize = vi.fn();
1326+
1327+
mockPtySpawn.mockReturnValue(mockPtyProcess);
1328+
mockGetPty.mockResolvedValue({
1329+
module: { spawn: mockPtySpawn },
1330+
name: 'mock-pty',
1331+
});
1332+
1333+
// Mock for child_process
1334+
mockChildProcess = new EventEmitter() as EventEmitter &
1335+
Partial<ChildProcess>;
1336+
mockChildProcess.stdout = new EventEmitter() as Readable;
1337+
mockChildProcess.stderr = new EventEmitter() as Readable;
1338+
mockChildProcess.kill = vi.fn();
1339+
Object.defineProperty(mockChildProcess, 'pid', {
1340+
value: 54321,
1341+
configurable: true,
1342+
});
1343+
mockCpSpawn.mockReturnValue(mockChildProcess);
1344+
1345+
// Default exit behavior for mocks
1346+
mockPtyProcess.onExit.mockImplementationOnce(({ exitCode, signal }) => {
1347+
// Small delay to allow async ops to complete
1348+
setTimeout(() => mockPtyProcess.emit('exit', { exitCode, signal }), 0);
1349+
});
1350+
mockChildProcess.on('exit', (code, signal) => {
1351+
// Small delay to allow async ops to complete
1352+
setTimeout(() => mockChildProcess.emit('close', code, signal), 0);
1353+
});
1354+
});
1355+
1356+
afterEach(() => {
1357+
// Clean up process.env after each test
1358+
vi.unstubAllEnvs();
1359+
});
1360+
1361+
it('should use a sanitized environment when in a GitHub run', async () => {
1362+
// Mock the environment to simulate a GitHub Actions run
1363+
vi.stubEnv('GITHUB_SHA', 'test-sha');
1364+
vi.stubEnv('MY_SENSITIVE_VAR', 'secret-value'); // This should be stripped out
1365+
vi.stubEnv('PATH', '/test/path'); // An essential var that should be kept
1366+
vi.stubEnv('GEMINI_CLI_TEST_VAR', 'test-value'); // A test var that should be kept
1367+
1368+
vi.resetModules();
1369+
const { ShellExecutionService } = await import(
1370+
'./shellExecutionService.js'
1371+
);
1372+
1373+
// Test pty path
1374+
await ShellExecutionService.execute(
1375+
'test-pty-command',
1376+
'/',
1377+
vi.fn(),
1378+
new AbortController().signal,
1379+
true,
1380+
shellExecutionConfig,
1381+
);
1382+
1383+
const ptyEnv = mockPtySpawn.mock.calls[0][2].env;
1384+
expect(ptyEnv).not.toHaveProperty('MY_SENSITIVE_VAR');
1385+
expect(ptyEnv).toHaveProperty('PATH', '/test/path');
1386+
expect(ptyEnv).toHaveProperty('GEMINI_CLI_TEST_VAR', 'test-value');
1387+
1388+
// Ensure pty process exits for next test
1389+
mockPtyProcess.onExit.mock.calls[0][0]({ exitCode: 0, signal: null });
1390+
await new Promise(process.nextTick);
1391+
1392+
// Test child_process path
1393+
mockGetPty.mockResolvedValue(null); // Force fallback
1394+
await ShellExecutionService.execute(
1395+
'test-cp-command',
1396+
'/',
1397+
vi.fn(),
1398+
new AbortController().signal,
1399+
true,
1400+
shellExecutionConfig,
1401+
);
1402+
1403+
const cpEnv = mockCpSpawn.mock.calls[0][2].env;
1404+
expect(cpEnv).not.toHaveProperty('MY_SENSITIVE_VAR');
1405+
expect(cpEnv).toHaveProperty('PATH', '/test/path');
1406+
expect(cpEnv).toHaveProperty('GEMINI_CLI_TEST_VAR', 'test-value');
1407+
1408+
// Ensure child_process exits
1409+
mockChildProcess.emit('exit', 0, null);
1410+
mockChildProcess.emit('close', 0, null);
1411+
await new Promise(process.nextTick);
1412+
});
1413+
1414+
it('should include the full process.env when not in a GitHub run', async () => {
1415+
vi.stubEnv('MY_TEST_VAR', 'test-value');
1416+
vi.stubEnv('GITHUB_SHA', '');
1417+
vi.stubEnv('SURFACE', '');
1418+
vi.resetModules();
1419+
const { ShellExecutionService } = await import(
1420+
'./shellExecutionService.js'
1421+
);
1422+
1423+
// Test pty path
1424+
await ShellExecutionService.execute(
1425+
'test-pty-command-no-github',
1426+
'/',
1427+
vi.fn(),
1428+
new AbortController().signal,
1429+
true,
1430+
shellExecutionConfig,
1431+
);
1432+
expect(mockPtySpawn).toHaveBeenCalled();
1433+
const ptyEnv = mockPtySpawn.mock.calls[0][2].env;
1434+
expect(ptyEnv).toHaveProperty('MY_TEST_VAR', 'test-value');
1435+
expect(ptyEnv).toHaveProperty('GEMINI_CLI', '1');
1436+
1437+
// Ensure pty process exits
1438+
mockPtyProcess.onExit.mock.calls[0][0]({ exitCode: 0, signal: null });
1439+
await new Promise(process.nextTick);
1440+
1441+
// Test child_process path (forcing fallback by making pty unavailable)
1442+
mockGetPty.mockResolvedValue(null);
1443+
await ShellExecutionService.execute(
1444+
'test-cp-command-no-github',
1445+
'/',
1446+
vi.fn(),
1447+
new AbortController().signal,
1448+
true, // Still tries pty, but it will fall back
1449+
shellExecutionConfig,
1450+
);
1451+
expect(mockCpSpawn).toHaveBeenCalled();
1452+
const cpEnv = mockCpSpawn.mock.calls[0][2].env;
1453+
expect(cpEnv).toHaveProperty('MY_TEST_VAR', 'test-value');
1454+
expect(cpEnv).toHaveProperty('GEMINI_CLI', '1');
1455+
1456+
// Ensure child_process exits
1457+
mockChildProcess.emit('exit', 0, null);
1458+
mockChildProcess.emit('close', 0, null);
1459+
await new Promise(process.nextTick);
1460+
});
1461+
});

packages/core/src/services/shellExecutionService.ts

Lines changed: 54 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,58 @@ const getFullBufferText = (terminal: pkg.Terminal): string => {
148148
return lines.join('\n');
149149
};
150150

151+
function getSanitizedEnv(): NodeJS.ProcessEnv {
152+
const isRunningInGithub =
153+
process.env['GITHUB_SHA'] || process.env['SURFACE'] === 'Github';
154+
155+
if (!isRunningInGithub) {
156+
// For local runs, we want to preserve the user's full environment.
157+
return { ...process.env };
158+
}
159+
160+
// For CI runs (GitHub), we sanitize the environment for security.
161+
const env: NodeJS.ProcessEnv = {};
162+
const essentialVars = [
163+
// Cross-platform
164+
'PATH',
165+
// Windows specific
166+
'Path',
167+
'SYSTEMROOT',
168+
'SystemRoot',
169+
'COMSPEC',
170+
'ComSpec',
171+
'PATHEXT',
172+
'WINDIR',
173+
'TEMP',
174+
'TMP',
175+
'USERPROFILE',
176+
'SYSTEMDRIVE',
177+
'SystemDrive',
178+
// Unix/Linux/macOS specific
179+
'HOME',
180+
'LANG',
181+
'SHELL',
182+
'TMPDIR',
183+
'USER',
184+
'LOGNAME',
185+
];
186+
187+
for (const key of essentialVars) {
188+
if (process.env[key] !== undefined) {
189+
env[key] = process.env[key];
190+
}
191+
}
192+
193+
// Always carry over test-specific variables for our own integration tests.
194+
for (const key in process.env) {
195+
if (key.startsWith('GEMINI_CLI_TEST')) {
196+
env[key] = process.env[key];
197+
}
198+
}
199+
200+
return env;
201+
}
202+
151203
/**
152204
* A centralized service for executing shell commands with robust process
153205
* management, cross-platform compatibility, and streaming output capabilities.
@@ -249,7 +301,7 @@ export class ShellExecutionService {
249301
shell: false,
250302
detached: !isWindows,
251303
env: {
252-
...process.env,
304+
...getSanitizedEnv(),
253305
GEMINI_CLI: '1',
254306
TERM: 'xterm-256color',
255307
PAGER: 'cat',
@@ -463,7 +515,7 @@ export class ShellExecutionService {
463515
cols,
464516
rows,
465517
env: {
466-
...process.env,
518+
...getSanitizedEnv(),
467519
GEMINI_CLI: '1',
468520
TERM: 'xterm-256color',
469521
PAGER: shellExecutionConfig.pager ?? 'cat',

0 commit comments

Comments
 (0)