/** * Graceful Shutdown — E-2 Critical Fix * * Handles SIGTERM / SIGINT to drain in-flight requests before exit. * Critical for Docker containers and Kubernetes pods where hard kills * can drop active SSE streams. * * Usage: * import { initGracefulShutdown } from "@/lib/gracefulShutdown"; * initGracefulShutdown(); * * @module lib/gracefulShutdown */ /** Grace period before forced exit (default 30s, configurable) */ const SHUTDOWN_TIMEOUT_MS = parseInt(process.env.SHUTDOWN_TIMEOUT_MS || "30000", 10); declare global { var __omnirouteShutdown: | { init: boolean; shuttingDown: boolean; activeRequests: number } | undefined; } function getShutdownState() { if (!globalThis.__omnirouteShutdown) { globalThis.__omnirouteShutdown = { init: false, shuttingDown: false, activeRequests: 0 }; } return globalThis.__omnirouteShutdown; } /** * Check if the server is currently shutting down. * Route handlers can use this to reject new requests. */ export function isDraining(): boolean { return getShutdownState().shuttingDown; } /** * Track a new in-flight request. Call `done()` when it completes. * Returns a done callback. */ export function trackRequest(): () => void { const state = getShutdownState(); state.activeRequests++; let called = false; return () => { if (!called) { called = true; state.activeRequests--; } }; } /** * Get current active request count (for monitoring/health endpoints). */ export function getActiveRequestCount(): number { return getShutdownState().activeRequests; } /** * Wait for all in-flight requests to complete, with timeout. */ async function waitForDrain(): Promise { const state = getShutdownState(); const start = Date.now(); const CHECK_INTERVAL_MS = 250; return new Promise((resolve) => { const check = () => { if (state.activeRequests <= 0) { console.log("[Shutdown] All in-flight requests drained."); resolve(); return; } if (Date.now() - start > SHUTDOWN_TIMEOUT_MS) { console.warn( `[Shutdown] Timeout after ${SHUTDOWN_TIMEOUT_MS}ms with ${state.activeRequests} active requests. Forcing exit.` ); resolve(); return; } console.log(`[Shutdown] Waiting for ${state.activeRequests} in-flight request(s)...`); setTimeout(check, CHECK_INTERVAL_MS); }; check(); }); } /** * Perform cleanup: close DB connections, flush logs. */ async function cleanup(): Promise { try { const [{ closeAuditDb }, { closeDbInstance }, { flushSpendBatchWriter }, { closeLogRotation }] = await Promise.all([ import("@omniroute/open-sse/mcp-server/audit.ts"), import("@/lib/db/core"), import("@/lib/spend/batchWriter"), import("@/lib/logRotation"), ]); const flushResult = await flushSpendBatchWriter(); if (flushResult.flushedEntries > 0) { console.log( `[Shutdown] Spend batch writer flushed ${flushResult.flushedEntries} pending entry(ies).` ); } if (closeAuditDb()) { console.log("[Shutdown] MCP audit database checkpointed and closed."); } if (closeDbInstance()) { console.log("[Shutdown] SQLite database checkpointed and closed."); } closeLogRotation(); console.log("[Shutdown] Log rotation timer stopped."); } catch (err) { console.error("[Shutdown] Error during cleanup:", (err as Error).message); } } /** * Initialize graceful shutdown handlers. * Should be called once during server startup. */ export function initGracefulShutdown(): void { const state = getShutdownState(); if (state.init) return; state.init = true; const shutdown = async (signal: string) => { if (state.shuttingDown) return; state.shuttingDown = true; console.log(`\n[Shutdown] Received ${signal}. Draining ${state.activeRequests} request(s)...`); await waitForDrain(); await cleanup(); console.log("[Shutdown] Bye."); process.exit(0); }; process.on("SIGTERM", () => shutdown("SIGTERM")); process.on("SIGINT", () => shutdown("SIGINT")); console.log("[Shutdown] Graceful shutdown handlers registered."); }