mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-11 14:40:52 -07:00
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>
This commit is contained in:
Gal Zahavi
@@ -4,7 +4,15 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { vi, describe, it, expect, beforeEach, type Mock } from 'vitest';
|
||||
import {
|
||||
vi,
|
||||
describe,
|
||||
it,
|
||||
expect,
|
||||
beforeEach,
|
||||
afterEach,
|
||||
type Mock,
|
||||
} from 'vitest';
|
||||
import EventEmitter from 'node:events';
|
||||
import type { Readable } from 'node:stream';
|
||||
import { type ChildProcess } from 'node:child_process';
|
||||
@@ -1284,3 +1292,170 @@ describe('ShellExecutionService execution method selection', () => {
|
||||
expect(result.executionMethod).toBe('child_process');
|
||||
});
|
||||
});
|
||||
|
||||
describe('ShellExecutionService environment variables', () => {
|
||||
let mockPtyProcess: EventEmitter & {
|
||||
pid: number;
|
||||
kill: Mock;
|
||||
onData: Mock;
|
||||
onExit: Mock;
|
||||
write: Mock;
|
||||
resize: Mock;
|
||||
};
|
||||
let mockChildProcess: EventEmitter & Partial<ChildProcess>;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.resetModules(); // Reset modules to ensure process.env changes are fresh
|
||||
|
||||
// Mock for pty
|
||||
mockPtyProcess = new EventEmitter() as EventEmitter & {
|
||||
pid: number;
|
||||
kill: Mock;
|
||||
onData: Mock;
|
||||
onExit: Mock;
|
||||
write: Mock;
|
||||
resize: Mock;
|
||||
};
|
||||
mockPtyProcess.pid = 12345;
|
||||
mockPtyProcess.kill = vi.fn();
|
||||
mockPtyProcess.onData = vi.fn();
|
||||
mockPtyProcess.onExit = vi.fn();
|
||||
mockPtyProcess.write = vi.fn();
|
||||
mockPtyProcess.resize = vi.fn();
|
||||
|
||||
mockPtySpawn.mockReturnValue(mockPtyProcess);
|
||||
mockGetPty.mockResolvedValue({
|
||||
module: { spawn: mockPtySpawn },
|
||||
name: 'mock-pty',
|
||||
});
|
||||
|
||||
// Mock for child_process
|
||||
mockChildProcess = new EventEmitter() as EventEmitter &
|
||||
Partial<ChildProcess>;
|
||||
mockChildProcess.stdout = new EventEmitter() as Readable;
|
||||
mockChildProcess.stderr = new EventEmitter() as Readable;
|
||||
mockChildProcess.kill = vi.fn();
|
||||
Object.defineProperty(mockChildProcess, 'pid', {
|
||||
value: 54321,
|
||||
configurable: true,
|
||||
});
|
||||
mockCpSpawn.mockReturnValue(mockChildProcess);
|
||||
|
||||
// Default exit behavior for mocks
|
||||
mockPtyProcess.onExit.mockImplementationOnce(({ exitCode, signal }) => {
|
||||
// Small delay to allow async ops to complete
|
||||
setTimeout(() => mockPtyProcess.emit('exit', { exitCode, signal }), 0);
|
||||
});
|
||||
mockChildProcess.on('exit', (code, signal) => {
|
||||
// Small delay to allow async ops to complete
|
||||
setTimeout(() => mockChildProcess.emit('close', code, signal), 0);
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
// Clean up process.env after each test
|
||||
vi.unstubAllEnvs();
|
||||
});
|
||||
|
||||
it('should use a sanitized environment when in a GitHub run', async () => {
|
||||
// Mock the environment to simulate a GitHub Actions run
|
||||
vi.stubEnv('GITHUB_SHA', 'test-sha');
|
||||
vi.stubEnv('MY_SENSITIVE_VAR', 'secret-value'); // This should be stripped out
|
||||
vi.stubEnv('PATH', '/test/path'); // An essential var that should be kept
|
||||
vi.stubEnv('GEMINI_CLI_TEST_VAR', 'test-value'); // A test var that should be kept
|
||||
|
||||
vi.resetModules();
|
||||
const { ShellExecutionService } = await import(
|
||||
'./shellExecutionService.js'
|
||||
);
|
||||
|
||||
// Test pty path
|
||||
await ShellExecutionService.execute(
|
||||
'test-pty-command',
|
||||
'/',
|
||||
vi.fn(),
|
||||
new AbortController().signal,
|
||||
true,
|
||||
shellExecutionConfig,
|
||||
);
|
||||
|
||||
const ptyEnv = mockPtySpawn.mock.calls[0][2].env;
|
||||
expect(ptyEnv).not.toHaveProperty('MY_SENSITIVE_VAR');
|
||||
expect(ptyEnv).toHaveProperty('PATH', '/test/path');
|
||||
expect(ptyEnv).toHaveProperty('GEMINI_CLI_TEST_VAR', 'test-value');
|
||||
|
||||
// Ensure pty process exits for next test
|
||||
mockPtyProcess.onExit.mock.calls[0][0]({ exitCode: 0, signal: null });
|
||||
await new Promise(process.nextTick);
|
||||
|
||||
// Test child_process path
|
||||
mockGetPty.mockResolvedValue(null); // Force fallback
|
||||
await ShellExecutionService.execute(
|
||||
'test-cp-command',
|
||||
'/',
|
||||
vi.fn(),
|
||||
new AbortController().signal,
|
||||
true,
|
||||
shellExecutionConfig,
|
||||
);
|
||||
|
||||
const cpEnv = mockCpSpawn.mock.calls[0][2].env;
|
||||
expect(cpEnv).not.toHaveProperty('MY_SENSITIVE_VAR');
|
||||
expect(cpEnv).toHaveProperty('PATH', '/test/path');
|
||||
expect(cpEnv).toHaveProperty('GEMINI_CLI_TEST_VAR', 'test-value');
|
||||
|
||||
// Ensure child_process exits
|
||||
mockChildProcess.emit('exit', 0, null);
|
||||
mockChildProcess.emit('close', 0, null);
|
||||
await new Promise(process.nextTick);
|
||||
});
|
||||
|
||||
it('should include the full process.env when not in a GitHub run', async () => {
|
||||
vi.stubEnv('MY_TEST_VAR', 'test-value');
|
||||
vi.stubEnv('GITHUB_SHA', '');
|
||||
vi.stubEnv('SURFACE', '');
|
||||
vi.resetModules();
|
||||
const { ShellExecutionService } = await import(
|
||||
'./shellExecutionService.js'
|
||||
);
|
||||
|
||||
// Test pty path
|
||||
await ShellExecutionService.execute(
|
||||
'test-pty-command-no-github',
|
||||
'/',
|
||||
vi.fn(),
|
||||
new AbortController().signal,
|
||||
true,
|
||||
shellExecutionConfig,
|
||||
);
|
||||
expect(mockPtySpawn).toHaveBeenCalled();
|
||||
const ptyEnv = mockPtySpawn.mock.calls[0][2].env;
|
||||
expect(ptyEnv).toHaveProperty('MY_TEST_VAR', 'test-value');
|
||||
expect(ptyEnv).toHaveProperty('GEMINI_CLI', '1');
|
||||
|
||||
// Ensure pty process exits
|
||||
mockPtyProcess.onExit.mock.calls[0][0]({ exitCode: 0, signal: null });
|
||||
await new Promise(process.nextTick);
|
||||
|
||||
// Test child_process path (forcing fallback by making pty unavailable)
|
||||
mockGetPty.mockResolvedValue(null);
|
||||
await ShellExecutionService.execute(
|
||||
'test-cp-command-no-github',
|
||||
'/',
|
||||
vi.fn(),
|
||||
new AbortController().signal,
|
||||
true, // Still tries pty, but it will fall back
|
||||
shellExecutionConfig,
|
||||
);
|
||||
expect(mockCpSpawn).toHaveBeenCalled();
|
||||
const cpEnv = mockCpSpawn.mock.calls[0][2].env;
|
||||
expect(cpEnv).toHaveProperty('MY_TEST_VAR', 'test-value');
|
||||
expect(cpEnv).toHaveProperty('GEMINI_CLI', '1');
|
||||
|
||||
// Ensure child_process exits
|
||||
mockChildProcess.emit('exit', 0, null);
|
||||
mockChildProcess.emit('close', 0, null);
|
||||
await new Promise(process.nextTick);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -148,6 +148,58 @@ const getFullBufferText = (terminal: pkg.Terminal): string => {
|
||||
return lines.join('\n');
|
||||
};
|
||||
|
||||
function getSanitizedEnv(): NodeJS.ProcessEnv {
|
||||
const isRunningInGithub =
|
||||
process.env['GITHUB_SHA'] || process.env['SURFACE'] === 'Github';
|
||||
|
||||
if (!isRunningInGithub) {
|
||||
// For local runs, we want to preserve the user's full environment.
|
||||
return { ...process.env };
|
||||
}
|
||||
|
||||
// For CI runs (GitHub), we sanitize the environment for security.
|
||||
const env: NodeJS.ProcessEnv = {};
|
||||
const essentialVars = [
|
||||
// Cross-platform
|
||||
'PATH',
|
||||
// Windows specific
|
||||
'Path',
|
||||
'SYSTEMROOT',
|
||||
'SystemRoot',
|
||||
'COMSPEC',
|
||||
'ComSpec',
|
||||
'PATHEXT',
|
||||
'WINDIR',
|
||||
'TEMP',
|
||||
'TMP',
|
||||
'USERPROFILE',
|
||||
'SYSTEMDRIVE',
|
||||
'SystemDrive',
|
||||
// Unix/Linux/macOS specific
|
||||
'HOME',
|
||||
'LANG',
|
||||
'SHELL',
|
||||
'TMPDIR',
|
||||
'USER',
|
||||
'LOGNAME',
|
||||
];
|
||||
|
||||
for (const key of essentialVars) {
|
||||
if (process.env[key] !== undefined) {
|
||||
env[key] = process.env[key];
|
||||
}
|
||||
}
|
||||
|
||||
// Always carry over test-specific variables for our own integration tests.
|
||||
for (const key in process.env) {
|
||||
if (key.startsWith('GEMINI_CLI_TEST')) {
|
||||
env[key] = process.env[key];
|
||||
}
|
||||
}
|
||||
|
||||
return env;
|
||||
}
|
||||
|
||||
/**
|
||||
* A centralized service for executing shell commands with robust process
|
||||
* management, cross-platform compatibility, and streaming output capabilities.
|
||||
@@ -249,7 +301,7 @@ export class ShellExecutionService {
|
||||
shell: false,
|
||||
detached: !isWindows,
|
||||
env: {
|
||||
...process.env,
|
||||
...getSanitizedEnv(),
|
||||
GEMINI_CLI: '1',
|
||||
TERM: 'xterm-256color',
|
||||
PAGER: 'cat',
|
||||
@@ -463,7 +515,7 @@ export class ShellExecutionService {
|
||||
cols,
|
||||
rows,
|
||||
env: {
|
||||
...process.env,
|
||||
...getSanitizedEnv(),
|
||||
GEMINI_CLI: '1',
|
||||
TERM: 'xterm-256color',
|
||||
PAGER: shellExecutionConfig.pager ?? 'cat',
|
||||
|
||||
Reference in New Issue
Block a user