diff --git a/src/__tests__/graceful-shutdown-integration.test.ts b/src/__tests__/graceful-shutdown-integration.test.ts new file mode 100644 index 0000000..17de1a4 --- /dev/null +++ b/src/__tests__/graceful-shutdown-integration.test.ts @@ -0,0 +1,119 @@ +import { describe, it, expect } from 'vitest'; +import { spawn, ChildProcess } from 'child_process'; +import { createServer } from 'net'; +import fs from 'fs'; +import path from 'path'; +import os from 'os'; + +describe('Graceful shutdown', () => { + async function getEphemeralPort(): Promise { + return new Promise((resolve, reject) => { + const srv = createServer(); + srv.unref(); + srv.on('error', reject); + srv.listen(0, '127.0.0.1', () => { + const addr = srv.address(); + if (addr && typeof addr === 'object') { + const port = addr.port; + srv.close(() => resolve(port)); + } else { + srv.close(); + reject(new Error('Failed to acquire ephemeral port')); + } + }); + }); + } + + async function waitForServerReady(baseUrl: string, timeoutMs = 20_000): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + try { + const response = await fetch(`${baseUrl}/healthz`); + if (response.ok) return; + } catch { + // not ready yet + } + await new Promise((r) => setTimeout(r, 500)); + } + throw new Error('Server did not become ready in time'); + } + + async function startServer(): Promise<{ proc: ChildProcess; dbPath: string; port: number }> { + const port = await getEphemeralPort(); + const dbPath = path.join( + os.tmpdir(), + `graceful_shutdown_${Date.now()}_${Math.random().toString(36).slice(2, 11)}.db` + ); + // Spawn tsx directly (no `pnpm dev:backend` wrapper) so SIGTERM/SIGINT + // reach the Node process unmediated by pnpm / concurrently. + const proc = spawn( + 'node', + ['--import', 'tsx/esm', 'src/index.ts', '--transport=http'], + { + env: { + ...process.env, + DSN: `sqlite://${dbPath}`, + PORT: port.toString(), + NODE_ENV: 'test', + }, + stdio: 'pipe', + } + ); + proc.stderr?.on('data', (d) => process.stderr.write(`[server] ${d}`)); + return { proc, dbPath, port }; + } + + async function waitForExit( + proc: ChildProcess, + timeoutMs: number + ): Promise<{ code: number | null; signal: NodeJS.Signals | null; timedOut: boolean }> { + return new Promise((resolve) => { + const timer = setTimeout(() => { + resolve({ code: null, signal: null, timedOut: true }); + }, timeoutMs); + proc.once('exit', (code, signal) => { + clearTimeout(timer); + resolve({ code, signal, timedOut: false }); + }); + }); + } + + async function cleanup(proc: ChildProcess, dbPath: string): Promise { + // `proc.killed` flips true as soon as we send any signal, even before + // the child has exited — check `exitCode`/`signalCode` to know whether + // the process is actually gone, and force-kill otherwise to avoid leaks. + if (proc.exitCode === null && proc.signalCode === null) { + proc.kill('SIGKILL'); + await waitForExit(proc, 5_000); + } + if (fs.existsSync(dbPath)) fs.unlinkSync(dbPath); + } + + it('exits with code 0 within a few seconds of SIGTERM in HTTP mode', async () => { + const { proc, dbPath, port } = await startServer(); + try { + await waitForServerReady(`http://localhost:${port}`); + proc.kill('SIGTERM'); + const result = await waitForExit(proc, 10_000); + expect(result.timedOut).toBe(false); + expect(result.code).toBe(0); + expect(result.signal).toBeNull(); + } finally { + await cleanup(proc, dbPath); + } + }, 60_000); + + it('exits with code 0 within a few seconds of SIGINT in HTTP mode', async () => { + const { proc, dbPath, port } = await startServer(); + try { + await waitForServerReady(`http://localhost:${port}`); + proc.kill('SIGINT'); + const result = await waitForExit(proc, 10_000); + expect(result.timedOut).toBe(false); + expect(result.code).toBe(0); + expect(result.signal).toBeNull(); + } finally { + await cleanup(proc, dbPath); + } + }, 60_000); +}); diff --git a/src/server.ts b/src/server.ts index 4319e3a..24cb3a6 100644 --- a/src/server.ts +++ b/src/server.ts @@ -158,8 +158,41 @@ See documentation for more details on configuring database connections. ); console.error(generateStartupTable(sourceDisplayInfos)); - // Clean up config watcher when the process is exiting (covers both transports) - process.on("exit", () => { stopConfigWatcher?.(); }); + let isShuttingDown = false; + const installShutdownHandlers = (closeTransport: () => Promise) => { + const shutdown = async (signal: string) => { + if (isShuttingDown) return; + isShuttingDown = true; + console.error(`Received ${signal}, shutting down...`); + + const forceExit = setTimeout(() => { + console.error("Graceful shutdown timed out after 25s, forcing exit"); + process.exit(1); + }, 25_000); + forceExit.unref(); + + let hadError = false; + const step = async (label: string, fn: () => unknown | Promise) => { + try { + await fn(); + } catch (err) { + hadError = true; + console.error(`Error during shutdown (${label}):`, err); + } + }; + + await step("close transport", closeTransport); + await step("disconnect connectors", () => connectorManager.disconnect()); + await step("stop config watcher", () => stopConfigWatcher?.()); + + clearTimeout(forceExit); + process.exit(hadError ? 1 : 0); + }; + + process.on("SIGINT", () => shutdown("SIGINT")); + process.on("SIGTERM", () => shutdown("SIGTERM")); + return shutdown; + }; // Set up transport-specific server if (transportData.type === "http") { @@ -258,7 +291,7 @@ See documentation for more details on configuring database connections. } // Start the HTTP server - app.listen(port, '0.0.0.0', () => { + const httpServer = app.listen(port, '0.0.0.0', () => { // In development mode, suggest using the Vite dev server for hot reloading if (process.env.NODE_ENV === 'development') { console.error('Development mode detected!'); @@ -270,6 +303,13 @@ See documentation for more details on configuring database connections. } console.error(`MCP server endpoint at http://localhost:${port}/mcp`); }); + + installShutdownHandlers( + () => + new Promise((resolve, reject) => + httpServer.close((err) => (err ? reject(err) : resolve())) + ) + ); } else { // STDIO transport: Pure MCP-over-stdio, no HTTP server const server = createServer(); @@ -277,24 +317,12 @@ See documentation for more details on configuring database connections. await server.connect(transport); console.error("MCP server running on stdio"); - let isShuttingDown = false; - const shutdown = async () => { - if (isShuttingDown) return; - isShuttingDown = true; - console.error("Shutting down..."); - await transport.close(); - await connectorManager.disconnect(); - process.exit(0); - }; - - // Listen for SIGINT/SIGTERM to gracefully shut down - process.on("SIGINT", shutdown); - process.on("SIGTERM", shutdown); + const shutdown = installShutdownHandlers(() => transport.close()); // Exit when stdin closes (parent process terminated). // On Windows, SIGINT/SIGTERM are not reliably sent when the parent // process exits — detecting stdin EOF is the portable way to handle this. - process.stdin.on("end", shutdown); + process.stdin.on("end", () => shutdown("stdin EOF")); } } catch (err) { console.error("Fatal error:", err);