gemini-cli/packages/cli/src/utils/cleanup.ts

/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import { promises as fs } from 'node:fs';
import { join } from 'node:path';
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);
}

export function registerSyncCleanup(fn: () => void) {
  syncCleanupFunctions.push(fn);
}

/**
 * Resets the internal cleanup state for testing purposes.
 * This allows tests to run in isolation without vi.resetModules().
 */
export function resetCleanupForTesting() {
  cleanupFunctions.length = 0;
  syncCleanupFunctions.length = 0;
  configForTelemetry = null;
  isShuttingDown = false;
}

export function runSyncCleanup() {
  for (const fn of syncCleanupFunctions) {
    try {
      fn();
    } catch (_) {
      // Ignore errors during cleanup.
    }
  }
  syncCleanupFunctions.length = 0;
}

/**
 * Register the config instance for telemetry shutdown.
 * This must be called early in the application lifecycle.
 */
export function registerTelemetryConfig(config: Config) {
  configForTelemetry = config;
}

export async function runExitCleanup() {
  // drain stdin to prevent printing garbage on exit
  // https://github.com/google-gemini/gemini-cli/issues/1680
  await drainStdin();

  runSyncCleanup();
  for (const fn of cleanupFunctions) {
    try {
      await fn();
    } catch (_) {
      // Ignore errors during cleanup.
    }
  }
  cleanupFunctions.length = 0; // Clear the array

  if (configForTelemetry) {
    try {
      await configForTelemetry.dispose();
    } catch (_) {
      // Ignore errors during disposal
    }
  }

  // IMPORTANT: Shutdown telemetry AFTER all other cleanup functions have run
  // This ensures SessionEnd hooks and other telemetry are properly flushed
  if (configForTelemetry && isTelemetrySdkInitialized()) {
    try {
      await shutdownTelemetry(configForTelemetry);
    } catch (_) {
      // Ignore errors during telemetry shutdown
    }
  }
}

async function drainStdin() {
  if (!process.stdin?.isTTY) return;
  // Resume stdin and attach a no-op listener to drain the buffer.
  // We use removeAllListeners to ensure we don't trigger other handlers.
  process.stdin
    .resume()
    .removeAllListeners('data')
    .on('data', () => {});
  // Give it a moment to flush the OS buffer.
  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();
  const tempDir = storage.getProjectTempDir();
  const checkpointsDir = join(tempDir, 'checkpoints');
  try {
    await fs.rm(checkpointsDir, { recursive: true, force: true });
  } catch {
    // Ignore errors if the directory doesn't exist or fails to delete.
  }
}