basant307/AI_Governance_Project

1/**2 * @license3 * Copyright 2025 Google LLC4 * SPDX-License-Identifier: Apache-2.05 */6 7import { promises as fs } from 'node:fs';8import { join } from 'node:path';9 10const cleanupFunctions: Array<(() => void) | (() => Promise<void>)> = [];11 12export function registerCleanup(fn: (() => void) | (() => Promise<void>)) {13  cleanupFunctions.push(fn);14}15 16/**17 * Per-cleanup ceiling. Caps any single hung cleanup (slow disk on18 * `chatRecording.flush`, MCP disconnect on a dead socket, telemetry HTTP19 * stall) so it can't starve the rest of the cleanup chain.20 */21const PER_CLEANUP_TIMEOUT_MS = 2_000;22 23/**24 * Wall-clock ceiling for the whole cleanup pass. Pre-async-jsonl, sync25 * fs writes were inherently bounded by their syscall return; with the26 * write queue moved off-thread, an unbounded `await flush()` could now27 * hang exit indefinitely. This ceiling guarantees the process always28 * exits within a bounded time, even if a cleanup never resolves.29 */30const OVERALL_CLEANUP_TIMEOUT_MS = 5_000;31 32/**33 * Awaits `promise`, but resolves to `undefined` if `ms` elapses first.34 * Rejection collapses to the same undefined resolution — caller treats35 * cleanup errors as best-effort. Timer is unrefed so it can't keep the36 * event loop alive on its own.37 */38function withTimeout<T>(promise: Promise<T>, ms: number): Promise<T | void> {39  return new Promise((resolve) => {40    const timer = setTimeout(() => resolve(undefined), ms);41    timer.unref?.();42    promise.then(43      (value) => {44        clearTimeout(timer);45        resolve(value);46      },47      () => {48        clearTimeout(timer);49        resolve(undefined);50      },51    );52  });53}54 55export interface RunExitCleanupOptions {56  /** TEST ONLY — override per-cleanup-function timeout (default 2s). */57  _testPerFnTimeoutMs?: number;58  /** TEST ONLY — override overall wall-clock timeout (default 5s). */59  _testOverallTimeoutMs?: number;60}61 62export async function runExitCleanup(63  options: RunExitCleanupOptions = {},64): Promise<void> {65  const perFn = options._testPerFnTimeoutMs ?? PER_CLEANUP_TIMEOUT_MS;66  const overall = options._testOverallTimeoutMs ?? OVERALL_CLEANUP_TIMEOUT_MS;67 68  const drain = (async () => {69    for (const fn of cleanupFunctions) {70      try {71        await withTimeout(Promise.resolve().then(fn), perFn);72      } catch (_) {73        // Ignore errors during cleanup.74      }75    }76  })();77 78  // clearTimeout when drain wins; unref keeps the handle from blocking exit.79  let wallClockTimer: NodeJS.Timeout | undefined;80  const wallClock = new Promise<void>((resolve) => {81    wallClockTimer = setTimeout(() => resolve(), overall);82    wallClockTimer.unref?.();83  });84 85  try {86    await Promise.race([drain, wallClock]);87  } finally {88    if (wallClockTimer) clearTimeout(wallClockTimer);89    cleanupFunctions.length = 0; // Clear the array90  }91}92 93/**94 * Test-only: clear the registered cleanup functions array. Module-private95 * state otherwise leaks across vitest cases — the previous test isolation96 * via `global['cleanupFunctions']` was a no-op (the array isn't on global)97 * and only happened to work because `runExitCleanup` itself clears at the98 * end. Naming follows the `_reset*ForTest` convention from99 * d6485964c (paths, jsonl-utils, ripGrep).100 */101export function _resetCleanupFunctionsForTest(): void {102  cleanupFunctions.length = 0;103}104 105export async function cleanupCheckpoints() {106  const { Storage } = await import('@qwen-code/qwen-code-core');107  const storage = new Storage(process.cwd());108  const tempDir = storage.getProjectTempDir();109  const checkpointsDir = join(tempDir, 'checkpoints');110  try {111    await fs.rm(checkpointsDir, { recursive: true, force: true });112  } catch {113    // Ignore errors if the directory doesn't exist or fails to delete.114  }115}116