Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion packages/cli/src/gemini.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -1216,6 +1216,8 @@ describe('startInteractiveUI', () => {
runExitCleanup: vi.fn(),
registerSyncCleanup: vi.fn(),
registerTelemetryConfig: vi.fn(),
setupSignalHandlers: vi.fn(),
setupTtyCheck: vi.fn(() => vi.fn()),
}));

beforeEach(() => {
Expand Down Expand Up @@ -1322,7 +1324,8 @@ describe('startInteractiveUI', () => {

// Verify all startup tasks were called
expect(getVersion).toHaveBeenCalledTimes(1);
expect(registerCleanup).toHaveBeenCalledTimes(4);
// 5 cleanups: mouseEvents, consolePatcher, lineWrapping, instance.unmount, and TTY check
expect(registerCleanup).toHaveBeenCalledTimes(5);

// Verify cleanup handler is registered with unmount function
const cleanupFn = vi.mocked(registerCleanup).mock.calls[0][0];
Expand Down
11 changes: 7 additions & 4 deletions packages/cli/src/gemini.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@ import {
registerSyncCleanup,
runExitCleanup,
registerTelemetryConfig,
setupSignalHandlers,
setupTtyCheck,
} from './utils/cleanup.js';
import {
cleanupToolOutputFiles,
Expand Down Expand Up @@ -319,6 +321,8 @@ export async function startInteractiveUI(
});

registerCleanup(() => instance.unmount());

registerCleanup(setupTtyCheck());
}

export async function main() {
Expand All @@ -340,6 +344,8 @@ export async function main() {

setupUnhandledRejectionHandler();

setupSignalHandlers();

const slashCommandConflictHandler = new SlashCommandConflictHandler();
slashCommandConflictHandler.start();
registerCleanup(() => slashCommandConflictHandler.stop());
Expand Down Expand Up @@ -646,10 +652,7 @@ export async function main() {
process.stdin.setRawMode(true);

// This cleanup isn't strictly needed but may help in certain situations.
process.on('SIGTERM', () => {
process.stdin.setRawMode(wasRaw);
});
process.on('SIGINT', () => {
registerSyncCleanup(() => {
process.stdin.setRawMode(wasRaw);
});
}
Expand Down
162 changes: 161 additions & 1 deletion packages/cli/src/utils/cleanup.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/

import { vi, describe, it, expect, beforeEach } from 'vitest';
import { vi, describe, it, expect, beforeEach, afterEach } from 'vitest';
import { promises as fs } from 'node:fs';
import * as path from 'node:path';

Expand All @@ -15,6 +15,7 @@ vi.mock('@google/gemini-cli-core', () => ({
})),
shutdownTelemetry: vi.fn(),
isTelemetrySdkInitialized: vi.fn().mockReturnValue(false),
ExitCodes: { SUCCESS: 0 },
}));

vi.mock('node:fs', () => ({
Expand All @@ -30,6 +31,8 @@ import {
runSyncCleanup,
cleanupCheckpoints,
resetCleanupForTesting,
setupSignalHandlers,
setupTtyCheck,
} from './cleanup.js';

describe('cleanup', () => {
Expand Down Expand Up @@ -123,3 +126,160 @@ describe('cleanup', () => {
});
});
});

describe('signal and TTY handling', () => {
let processOnHandlers: Map<
string,
Array<(...args: unknown[]) => void | Promise<void>>
>;

beforeEach(() => {
processOnHandlers = new Map();
resetCleanupForTesting();

vi.spyOn(process, 'on').mockImplementation(
(event: string | symbol, handler: (...args: unknown[]) => void) => {
if (typeof event === 'string') {
const handlers = processOnHandlers.get(event) || [];
handlers.push(handler);
processOnHandlers.set(event, handlers);
}
return process;
},
);

vi.spyOn(process, 'exit').mockImplementation((() => {
// Don't actually exit
}) as typeof process.exit);
});

afterEach(() => {
vi.restoreAllMocks();
processOnHandlers.clear();
});

describe('setupSignalHandlers', () => {
it('should register handlers for SIGHUP, SIGTERM, and SIGINT', () => {
setupSignalHandlers();

expect(processOnHandlers.has('SIGHUP')).toBe(true);
expect(processOnHandlers.has('SIGTERM')).toBe(true);
expect(processOnHandlers.has('SIGINT')).toBe(true);
});

it('should gracefully shutdown when SIGHUP is received', async () => {
setupSignalHandlers();

const sighupHandlers = processOnHandlers.get('SIGHUP') || [];
expect(sighupHandlers.length).toBeGreaterThan(0);

await sighupHandlers[0]?.();

expect(process.exit).toHaveBeenCalledWith(0);
});

it('should register SIGTERM handler that can trigger shutdown', () => {
setupSignalHandlers();

const sigtermHandlers = processOnHandlers.get('SIGTERM') || [];
expect(sigtermHandlers.length).toBeGreaterThan(0);
expect(typeof sigtermHandlers[0]).toBe('function');
});
});

describe('setupTtyCheck', () => {
let originalStdinIsTTY: boolean | undefined;
let originalStdoutIsTTY: boolean | undefined;

beforeEach(() => {
originalStdinIsTTY = process.stdin.isTTY;
originalStdoutIsTTY = process.stdout.isTTY;
vi.useFakeTimers();
});

afterEach(() => {
vi.useRealTimers();
Object.defineProperty(process.stdin, 'isTTY', {
value: originalStdinIsTTY,
writable: true,
configurable: true,
});
Object.defineProperty(process.stdout, 'isTTY', {
value: originalStdoutIsTTY,
writable: true,
configurable: true,
});
});

it('should return a cleanup function', () => {
const cleanup = setupTtyCheck();
expect(typeof cleanup).toBe('function');
cleanup();
});

it('should not exit when both stdin and stdout are TTY', async () => {
Object.defineProperty(process.stdin, 'isTTY', {
value: true,
writable: true,
configurable: true,
});
Object.defineProperty(process.stdout, 'isTTY', {
value: true,
writable: true,
configurable: true,
});

const cleanup = setupTtyCheck();
await vi.advanceTimersByTimeAsync(5000);
expect(process.exit).not.toHaveBeenCalled();
cleanup();
});

it('should exit when both stdin and stdout are not TTY', async () => {
Object.defineProperty(process.stdin, 'isTTY', {
value: false,
writable: true,
configurable: true,
});
Object.defineProperty(process.stdout, 'isTTY', {
value: false,
writable: true,
configurable: true,
});

const cleanup = setupTtyCheck();
await vi.advanceTimersByTimeAsync(5000);
expect(process.exit).toHaveBeenCalledWith(0);
cleanup();
});

it('should not check when SANDBOX env is set', async () => {
const originalSandbox = process.env['SANDBOX'];
process.env['SANDBOX'] = 'true';

Object.defineProperty(process.stdin, 'isTTY', {
value: false,
writable: true,
configurable: true,
});
Object.defineProperty(process.stdout, 'isTTY', {
value: false,
writable: true,
configurable: true,
});

const cleanup = setupTtyCheck();
await vi.advanceTimersByTimeAsync(5000);
expect(process.exit).not.toHaveBeenCalled();
cleanup();
process.env['SANDBOX'] = originalSandbox;
});

it('cleanup function should stop the interval', () => {
const cleanup = setupTtyCheck();
cleanup();
vi.advanceTimersByTime(10000);
expect(process.exit).not.toHaveBeenCalled();
});
});
});
62 changes: 62 additions & 0 deletions packages/cli/src/utils/cleanup.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,14 @@ import {
Storage,
shutdownTelemetry,
isTelemetrySdkInitialized,
ExitCodes,
} from '@google/gemini-cli-core';
import type { Config } from '@google/gemini-cli-core';

const cleanupFunctions: Array<(() => void) | (() => Promise<void>)> = [];
const syncCleanupFunctions: Array<() => void> = [];
let configForTelemetry: Config | null = null;
let isShuttingDown = false;

export function registerCleanup(fn: (() => void) | (() => Promise<void>)) {
cleanupFunctions.push(fn);
Expand All @@ -33,6 +35,7 @@ export function resetCleanupForTesting() {
cleanupFunctions.length = 0;
syncCleanupFunctions.length = 0;
configForTelemetry = null;
isShuttingDown = false;
}

export function runSyncCleanup() {
Expand Down Expand Up @@ -100,6 +103,65 @@ async function drainStdin() {
await new Promise((resolve) => setTimeout(resolve, 50));
}

/**
* Gracefully shuts down the process, ensuring cleanup runs exactly once.
* Guards against concurrent shutdown from signals (SIGHUP, SIGTERM, SIGINT)
* and TTY loss detection racing each other.
*
* @see https://github.com/google-gemini/gemini-cli/issues/15874
*/
async function gracefulShutdown(_reason: string) {
if (isShuttingDown) {
return;
}
isShuttingDown = true;

await runExitCleanup();
process.exit(ExitCodes.SUCCESS);
}

export function setupSignalHandlers() {
process.on('SIGHUP', () => gracefulShutdown('SIGHUP'));
process.on('SIGTERM', () => gracefulShutdown('SIGTERM'));
process.on('SIGINT', () => gracefulShutdown('SIGINT'));
}

export function setupTtyCheck(): () => void {
let intervalId: ReturnType<typeof setInterval> | null = null;
let isCheckingTty = false;

intervalId = setInterval(async () => {
if (isCheckingTty || isShuttingDown) {
return;
}

if (process.env['SANDBOX']) {
return;
}

if (!process.stdin.isTTY && !process.stdout.isTTY) {
isCheckingTty = true;

if (intervalId) {
clearInterval(intervalId);
intervalId = null;
}

await gracefulShutdown('TTY loss');
}
}, 5000);

// Don't keep the process alive just for this interval
intervalId.unref();

return () => {
if (intervalId) {
clearInterval(intervalId);
intervalId = null;
}
};
}

export async function cleanupCheckpoints() {
const storage = new Storage(process.cwd());
await storage.initialize();
Expand Down
Loading