Skip to content

Commit a71d401

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

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';
@@ -1226,3 +1234,170 @@ describe('ShellExecutionService execution method selection', () => {
12261234
expect(result.executionMethod).toBe('child_process');
12271235
});
12281236
});
1237+
1238+
describe('ShellExecutionService environment variables', () => {
1239+
let mockPtyProcess: EventEmitter & {
1240+
pid: number;
1241+
kill: Mock;
1242+
onData: Mock;
1243+
onExit: Mock;
1244+
write: Mock;
1245+
resize: Mock;
1246+
};
1247+
let mockChildProcess: EventEmitter & Partial<ChildProcess>;
1248+
1249+
beforeEach(() => {
1250+
vi.clearAllMocks();
1251+
vi.resetModules(); // Reset modules to ensure process.env changes are fresh
1252+
1253+
// Mock for pty
1254+
mockPtyProcess = new EventEmitter() as EventEmitter & {
1255+
pid: number;
1256+
kill: Mock;
1257+
onData: Mock;
1258+
onExit: Mock;
1259+
write: Mock;
1260+
resize: Mock;
1261+
};
1262+
mockPtyProcess.pid = 12345;
1263+
mockPtyProcess.kill = vi.fn();
1264+
mockPtyProcess.onData = vi.fn();
1265+
mockPtyProcess.onExit = vi.fn();
1266+
mockPtyProcess.write = vi.fn();
1267+
mockPtyProcess.resize = vi.fn();
1268+
1269+
mockPtySpawn.mockReturnValue(mockPtyProcess);
1270+
mockGetPty.mockResolvedValue({
1271+
module: { spawn: mockPtySpawn },
1272+
name: 'mock-pty',
1273+
});
1274+
1275+
// Mock for child_process
1276+
mockChildProcess = new EventEmitter() as EventEmitter &
1277+
Partial<ChildProcess>;
1278+
mockChildProcess.stdout = new EventEmitter() as Readable;
1279+
mockChildProcess.stderr = new EventEmitter() as Readable;
1280+
mockChildProcess.kill = vi.fn();
1281+
Object.defineProperty(mockChildProcess, 'pid', {
1282+
value: 54321,
1283+
configurable: true,
1284+
});
1285+
mockCpSpawn.mockReturnValue(mockChildProcess);
1286+
1287+
// Default exit behavior for mocks
1288+
mockPtyProcess.onExit.mockImplementationOnce(({ exitCode, signal }) => {
1289+
// Small delay to allow async ops to complete
1290+
setTimeout(() => mockPtyProcess.emit('exit', { exitCode, signal }), 0);
1291+
});
1292+
mockChildProcess.on('exit', (code, signal) => {
1293+
// Small delay to allow async ops to complete
1294+
setTimeout(() => mockChildProcess.emit('close', code, signal), 0);
1295+
});
1296+
});
1297+
1298+
afterEach(() => {
1299+
// Clean up process.env after each test
1300+
vi.unstubAllEnvs();
1301+
});
1302+
1303+
it('should use a sanitized environment when in a GitHub run', async () => {
1304+
// Mock the environment to simulate a GitHub Actions run
1305+
vi.stubEnv('GITHUB_SHA', 'test-sha');
1306+
vi.stubEnv('MY_SENSITIVE_VAR', 'secret-value'); // This should be stripped out
1307+
vi.stubEnv('PATH', '/test/path'); // An essential var that should be kept
1308+
vi.stubEnv('GEMINI_CLI_TEST_VAR', 'test-value'); // A test var that should be kept
1309+
1310+
vi.resetModules();
1311+
const { ShellExecutionService } = await import(
1312+
'./shellExecutionService.js'
1313+
);
1314+
1315+
// Test pty path
1316+
await ShellExecutionService.execute(
1317+
'test-pty-command',
1318+
'/',
1319+
vi.fn(),
1320+
new AbortController().signal,
1321+
true,
1322+
shellExecutionConfig,
1323+
);
1324+
1325+
const ptyEnv = mockPtySpawn.mock.calls[0][2].env;
1326+
expect(ptyEnv).not.toHaveProperty('MY_SENSITIVE_VAR');
1327+
expect(ptyEnv).toHaveProperty('PATH', '/test/path');
1328+
expect(ptyEnv).toHaveProperty('GEMINI_CLI_TEST_VAR', 'test-value');
1329+
1330+
// Ensure pty process exits for next test
1331+
mockPtyProcess.onExit.mock.calls[0][0]({ exitCode: 0, signal: null });
1332+
await new Promise(process.nextTick);
1333+
1334+
// Test child_process path
1335+
mockGetPty.mockResolvedValue(null); // Force fallback
1336+
await ShellExecutionService.execute(
1337+
'test-cp-command',
1338+
'/',
1339+
vi.fn(),
1340+
new AbortController().signal,
1341+
true,
1342+
shellExecutionConfig,
1343+
);
1344+
1345+
const cpEnv = mockCpSpawn.mock.calls[0][2].env;
1346+
expect(cpEnv).not.toHaveProperty('MY_SENSITIVE_VAR');
1347+
expect(cpEnv).toHaveProperty('PATH', '/test/path');
1348+
expect(cpEnv).toHaveProperty('GEMINI_CLI_TEST_VAR', 'test-value');
1349+
1350+
// Ensure child_process exits
1351+
mockChildProcess.emit('exit', 0, null);
1352+
mockChildProcess.emit('close', 0, null);
1353+
await new Promise(process.nextTick);
1354+
});
1355+
1356+
it('should include the full process.env when not in a GitHub run', async () => {
1357+
vi.stubEnv('MY_TEST_VAR', 'test-value');
1358+
vi.stubEnv('GITHUB_SHA', '');
1359+
vi.stubEnv('SURFACE', '');
1360+
vi.resetModules();
1361+
const { ShellExecutionService } = await import(
1362+
'./shellExecutionService.js'
1363+
);
1364+
1365+
// Test pty path
1366+
await ShellExecutionService.execute(
1367+
'test-pty-command-no-github',
1368+
'/',
1369+
vi.fn(),
1370+
new AbortController().signal,
1371+
true,
1372+
shellExecutionConfig,
1373+
);
1374+
expect(mockPtySpawn).toHaveBeenCalled();
1375+
const ptyEnv = mockPtySpawn.mock.calls[0][2].env;
1376+
expect(ptyEnv).toHaveProperty('MY_TEST_VAR', 'test-value');
1377+
expect(ptyEnv).toHaveProperty('GEMINI_CLI', '1');
1378+
1379+
// Ensure pty process exits
1380+
mockPtyProcess.onExit.mock.calls[0][0]({ exitCode: 0, signal: null });
1381+
await new Promise(process.nextTick);
1382+
1383+
// Test child_process path (forcing fallback by making pty unavailable)
1384+
mockGetPty.mockResolvedValue(null);
1385+
await ShellExecutionService.execute(
1386+
'test-cp-command-no-github',
1387+
'/',
1388+
vi.fn(),
1389+
new AbortController().signal,
1390+
true, // Still tries pty, but it will fall back
1391+
shellExecutionConfig,
1392+
);
1393+
expect(mockCpSpawn).toHaveBeenCalled();
1394+
const cpEnv = mockCpSpawn.mock.calls[0][2].env;
1395+
expect(cpEnv).toHaveProperty('MY_TEST_VAR', 'test-value');
1396+
expect(cpEnv).toHaveProperty('GEMINI_CLI', '1');
1397+
1398+
// Ensure child_process exits
1399+
mockChildProcess.emit('exit', 0, null);
1400+
mockChildProcess.emit('close', 0, null);
1401+
await new Promise(process.nextTick);
1402+
});
1403+
});

packages/core/src/services/shellExecutionService.ts

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

131+
function getSanitizedEnv(): NodeJS.ProcessEnv {
132+
const isRunningInGithub =
133+
process.env['GITHUB_SHA'] || process.env['SURFACE'] === 'Github';
134+
135+
if (!isRunningInGithub) {
136+
// For local runs, we want to preserve the user's full environment.
137+
return { ...process.env };
138+
}
139+
140+
// For CI runs (GitHub), we sanitize the environment for security.
141+
const env: NodeJS.ProcessEnv = {};
142+
const essentialVars = [
143+
// Cross-platform
144+
'PATH',
145+
// Windows specific
146+
'Path',
147+
'SYSTEMROOT',
148+
'SystemRoot',
149+
'COMSPEC',
150+
'ComSpec',
151+
'PATHEXT',
152+
'WINDIR',
153+
'TEMP',
154+
'TMP',
155+
'USERPROFILE',
156+
'SYSTEMDRIVE',
157+
'SystemDrive',
158+
// Unix/Linux/macOS specific
159+
'HOME',
160+
'LANG',
161+
'SHELL',
162+
'TMPDIR',
163+
'USER',
164+
'LOGNAME',
165+
];
166+
167+
for (const key of essentialVars) {
168+
if (process.env[key] !== undefined) {
169+
env[key] = process.env[key];
170+
}
171+
}
172+
173+
// Always carry over test-specific variables for our own integration tests.
174+
for (const key in process.env) {
175+
if (key.startsWith('GEMINI_CLI_TEST')) {
176+
env[key] = process.env[key];
177+
}
178+
}
179+
180+
return env;
181+
}
182+
131183
/**
132184
* A centralized service for executing shell commands with robust process
133185
* management, cross-platform compatibility, and streaming output capabilities.
@@ -229,7 +281,7 @@ export class ShellExecutionService {
229281
shell: false,
230282
detached: !isWindows,
231283
env: {
232-
...process.env,
284+
...getSanitizedEnv(),
233285
GEMINI_CLI: '1',
234286
TERM: 'xterm-256color',
235287
PAGER: 'cat',
@@ -443,7 +495,7 @@ export class ShellExecutionService {
443495
cols,
444496
rows,
445497
env: {
446-
...process.env,
498+
...getSanitizedEnv(),
447499
GEMINI_CLI: '1',
448500
TERM: 'xterm-256color',
449501
PAGER: shellExecutionConfig.pager ?? 'cat',

0 commit comments

Comments
 (0)