diff --git a/apps/dev-playground/package.json b/apps/dev-playground/package.json index 1dc9358c5..7af74ef0e 100644 --- a/apps/dev-playground/package.json +++ b/apps/dev-playground/package.json @@ -4,7 +4,7 @@ "main": "build/index.js", "scripts": { "start": "node build/index.mjs", - "start:local": "NODE_ENV=production node --env-file=./server/.env build/index.mjs", + "start:local": "NODE_ENV=production node --env-file=.env build/index.mjs", "dev": "NODE_ENV=development tsx watch server/index.ts", "dev:inspect": "NODE_ENV=development tsx --inspect --tsconfig ./tsconfig.json ./server", "build": "npm run build:app", diff --git a/docs/docs/api/appkit/Class.Plugin.md b/docs/docs/api/appkit/Class.Plugin.md index 34034537f..a5002a097 100644 --- a/docs/docs/api/appkit/Class.Plugin.md +++ b/docs/docs/api/appkit/Class.Plugin.md @@ -211,6 +211,11 @@ Plugin initialization phase. abortActiveOperations(): void; ``` +Cancel in-flight work (abort signals, SSE streams). Runs in the first +phase of graceful shutdown, BEFORE any plugin's `shutdown()` hook — +so it must not tear down shared resources (e.g. connection pools) +that other plugins' hooks may still need. Put teardown in `shutdown()`. + #### Returns `void` diff --git a/packages/appkit/src/core/appkit.ts b/packages/appkit/src/core/appkit.ts index 2973ee265..1a0ec6bd5 100644 --- a/packages/appkit/src/core/appkit.ts +++ b/packages/appkit/src/core/appkit.ts @@ -20,6 +20,7 @@ import { isPlainObject } from "../plugin/plugin"; import { ResourceRegistry, ResourceType } from "../registry"; import type { TelemetryConfig } from "../telemetry"; import { TelemetryManager } from "../telemetry"; +import { LifecycleManager } from "./lifecycle-manager"; import { isToolProvider, PluginContext } from "./plugin-context"; const logger = createLogger("appkit"); @@ -236,6 +237,12 @@ export class AppKit { await (serverPlugin as any).start(); } + // Core owns graceful shutdown: install the signal handlers once every + // plugin has started. Applies uniformly whether or not a server plugin + // is present — server-less apps still get their telemetry flushed and + // plugin shutdown() hooks run. + new LifecycleManager(instance.#context).installSignalHandlers(); + return handle; } diff --git a/packages/appkit/src/core/lifecycle-manager.ts b/packages/appkit/src/core/lifecycle-manager.ts new file mode 100644 index 000000000..8d1811203 --- /dev/null +++ b/packages/appkit/src/core/lifecycle-manager.ts @@ -0,0 +1,270 @@ +import type { BasePlugin } from "shared"; +import { CacheManager } from "../cache"; +import { TelemetryReporter } from "../internal-telemetry"; +import { createLogger } from "../logging/logger"; +import { TelemetryManager } from "../telemetry"; +import type { PluginContext } from "./plugin-context"; + +const logger = createLogger("lifecycle"); + +/** + * Owns the process's graceful-shutdown sequence. + * + * Created by AppKit core once every plugin has started. It is the single + * owner of the SIGTERM/SIGINT handlers and of `process.exit`, mirroring the + * core-owned startup in `AppKit._createApp`: core initializes telemetry, + * cache, and the internal-telemetry reporter, and core tears them all down + * here. Plugins participate through the generic hooks + * (`abortActiveOperations()`, `shutdown()`, and `onLifecycle("shutdown")`) — + * they do not touch process signals or the core singletons themselves. + */ +export class LifecycleManager { + /** + * Overall graceful-shutdown budget before the process is force-exited. + * + * Budget arithmetic: plugin `shutdown()` hooks run concurrently and are + * bounded by {@link PLUGIN_SHUTDOWN_TIMEOUT_MS} (10s); the lifecycle emit + * is bounded by {@link PHASE_SHUTDOWN_TIMEOUT_MS} (2s); the cache storage + * close and the telemetry flush run concurrently, each bounded by + * {@link PHASE_SHUTDOWN_TIMEOUT_MS} (2s). Worst case is + * 10s + 2s + max(2s, 2s) = 14s, leaving ~1s of margin for the remaining + * steps (aborts) before this timer force-exits. + */ + private static readonly SHUTDOWN_TIMEOUT_MS = 15_000; + /** + * Per-plugin budget for `shutdown()` hooks. Sized to cover the longest + * built-in drain (the files plugin waits up to 10s for in-flight writes). + */ + private static readonly PLUGIN_SHUTDOWN_TIMEOUT_MS = 10_000; + /** + * Budget for each non-plugin shutdown phase (the `"shutdown"` lifecycle + * emit, the cache storage close, and the telemetry flush). Keeps the + * worst-case total under {@link SHUTDOWN_TIMEOUT_MS} — see the arithmetic + * there. + */ + private static readonly PHASE_SHUTDOWN_TIMEOUT_MS = 2_000; + + /** + * Guards against re-entrant shutdown (e.g. SIGTERM followed by SIGINT). + * The flag set in `shutdown` must remain synchronous and first — any + * `await` before it would open a window for a second signal to re-enter + * the sequence. + */ + private isShuttingDown = false; + /** + * Name of the shutdown phase currently in flight, so the force-exit log + * can say where shutdown got stuck without extra bookkeeping. + */ + private shutdownPhase = "not started"; + + constructor(private readonly context: PluginContext) {} + + /** + * Install the SIGTERM/SIGINT handlers that trigger {@link shutdown}. + * + * Uses `process.once` (not `on`) so a repeated signal cannot register the + * handler twice; re-entrancy from a *different* signal is guarded by + * `isShuttingDown` inside {@link shutdown}. + */ + installSignalHandlers(): void { + process.once("SIGTERM", () => this.shutdown()); + process.once("SIGINT", () => this.shutdown()); + } + + /** + * Run the graceful-shutdown sequence and exit the process. + * + * Phases: + * 1. stop the internal-telemetry reporter + * 2. abort in-flight work on every plugin (cancellation only — teardown of + * shared resources belongs in `shutdown()` so peers can still drain) + * 3. run every plugin's `shutdown()` hook concurrently, each bounded + * 4. emit the `"shutdown"` lifecycle event, bounded + * 5. close the cache storage and flush telemetry concurrently, each bounded + * + * Exits 0 on completion (and on the force-exit backstop): a deliberate + * shutdown is not a crash. Exit 1 is reserved for an unexpected error + * thrown by the sequence itself. + */ + async shutdown(): Promise { + // Must stay synchronous and first: any await before the flag is set + // would let a second signal re-enter the shutdown sequence. + if (this.isShuttingDown) return; + this.isShuttingDown = true; + + logger.info("Starting graceful shutdown..."); + + let exitCode = 0; + + // Force exit once the overall budget is spent. Exit 0 is deliberate: + // a force-timeout still happens on a routine deploy (deliberate + // shutdown, not a crash), and orchestrators record nonzero exits on + // deploys as crashes. The error log below is the stuck-shutdown + // signal instead of the exit code. + const forceExitTimer = setTimeout(() => { + logger.error( + "Graceful shutdown did NOT complete within the %dms budget (phase in flight: %s); force-exiting with code 0.", + LifecycleManager.SHUTDOWN_TIMEOUT_MS, + this.shutdownPhase, + ); + process.exit(0); + }, LifecycleManager.SHUTDOWN_TIMEOUT_MS); + // unref so this backstop timer never by itself keeps the process alive. + // Any real pending teardown (OTEL export timer, DB pool sockets, the + // still-open HTTP listener) is a ref'd handle that holds the loop open + // until this fires; if nothing is ref'd, there is nothing left to tear + // down and exiting early is correct. + forceExitTimer.unref(); + + try { + const plugins = Array.from(this.context.getPlugins().values()); + + // 1. stop the internal-telemetry reporter (no-op if never started). + this.shutdownPhase = "stopping internal telemetry reporter"; + TelemetryReporter.getInstance()?.stop(); + + // 2. abort active operations from plugins (in-flight executions, SSE + // streams). Cancellation only — resource teardown (e.g. the + // lakebase pools, the server's socket close) belongs in plugin + // shutdown() hooks / lifecycle subscribers so other plugins can + // still drain state through them. + this.shutdownPhase = "aborting active operations"; + for (const plugin of plugins) { + if (plugin.abortActiveOperations) { + try { + plugin.abortActiveOperations(); + } catch (err) { + logger.error( + "Error aborting operations for plugin %s: %O", + plugin.name, + err, + ); + } + } + } + + // 3. run every plugin's shutdown() hook concurrently, each bounded + // by a per-plugin timeout so one hung plugin cannot stall exit. + this.shutdownPhase = "plugin shutdown() hooks"; + await Promise.all( + plugins + .filter((plugin) => typeof plugin.shutdown === "function") + .map((plugin) => this.runPluginShutdown(plugin)), + ); + + // 4. notify lifecycle subscribers, bounded so a slow subscriber + // cannot eat the remaining budget. The server plugin closes its + // remaining sockets here, after other plugins have drained. + this.shutdownPhase = "shutdown lifecycle emit"; + try { + await this.raceWithTimeout( + this.context.emitLifecycle("shutdown"), + LifecycleManager.PHASE_SHUTDOWN_TIMEOUT_MS, + "shutdown lifecycle emit", + ); + } catch (err) { + logger.error("Error emitting shutdown lifecycle event: %O", err); + } + + // 5. close the cache manager's storage (drains the persistent + // Lakebase pool; no-op for in-memory storage) and flush telemetry. + // Runs after the lifecycle emit so subscribers can still read the + // cache. The two are independent (the flush never touches the + // cache), so they run concurrently — each bounded so a stuck pool + // drain or stalled OTLP export cannot eat the remaining budget. + this.shutdownPhase = "cache storage close + telemetry flush"; + await Promise.all([this.closeCacheStorage(), this.flushTelemetry()]); + + logger.info("Graceful shutdown complete"); + } catch (err) { + // Exit 1 is reserved for an unexpected error thrown by the sequence + // itself; every per-phase failure above is already caught and logged. + logger.error("Error during graceful shutdown: %O", err); + exitCode = 1; + } + + clearTimeout(forceExitTimer); + process.exit(exitCode); + } + + /** Close the cache storage, bounded and error-isolated. */ + private async closeCacheStorage(): Promise { + let cache: CacheManager; + try { + cache = CacheManager.getInstanceSync(); + } catch { + // Cache was never initialized — nothing to close. + return; + } + try { + await this.raceWithTimeout( + cache.close(), + LifecycleManager.PHASE_SHUTDOWN_TIMEOUT_MS, + "cache storage close", + ); + } catch (err) { + logger.error("Error closing cache storage during shutdown: %O", err); + } + } + + /** Flush and shut down the telemetry SDK, bounded and error-isolated. */ + private async flushTelemetry(): Promise { + try { + await this.raceWithTimeout( + TelemetryManager.getInstance().shutdown(), + LifecycleManager.PHASE_SHUTDOWN_TIMEOUT_MS, + "telemetry flush", + ); + } catch (err) { + logger.error("Error flushing telemetry during shutdown: %O", err); + } + } + + /** + * Run a single plugin's `shutdown()` hook bounded by + * {@link LifecycleManager.PLUGIN_SHUTDOWN_TIMEOUT_MS}. Errors and timeouts + * are logged but never thrown so one misbehaving plugin cannot block + * the rest of the shutdown sequence. + */ + private async runPluginShutdown(plugin: BasePlugin): Promise { + try { + await this.raceWithTimeout( + plugin.shutdown?.(), + LifecycleManager.PLUGIN_SHUTDOWN_TIMEOUT_MS, + "shutdown()", + ); + } catch (err) { + logger.error("Error shutting down plugin %s: %O", plugin.name, err); + } + } + + /** + * Race `work` against a timeout. Rejects with a labeled error when the + * timeout wins. A no-op rejection handler is attached to the work promise + * before racing so a branch that rejects after the timeout already won + * does not surface as an unhandledRejection. + */ + private async raceWithTimeout( + work: Promise | T, + timeoutMs: number, + label: string, + ): Promise { + const promise = Promise.resolve(work); + promise.catch(() => {}); + let timer: NodeJS.Timeout | undefined; + try { + return await Promise.race([ + promise, + new Promise((_, reject) => { + timer = setTimeout( + () => reject(new Error(`${label} timed out after ${timeoutMs}ms`)), + timeoutMs, + ); + timer.unref(); + }), + ]); + } finally { + if (timer) clearTimeout(timer); + } + } +} diff --git a/packages/appkit/src/core/plugin-context.ts b/packages/appkit/src/core/plugin-context.ts index eda3f05a5..03919c971 100644 --- a/packages/appkit/src/core/plugin-context.ts +++ b/packages/appkit/src/core/plugin-context.ts @@ -24,6 +24,19 @@ interface RouteTarget { type ToolProviderPlugin = BasePlugin & ToolProvider & { asUser: (req: IAppRequest) => ToolProvider }; +/** + * Lifecycle events emitted through {@link PluginContext.emitLifecycle}. + * + * - `"setup:complete"` — emitted by AppKit core after every plugin's + * `setup()` has finished. + * - `"server:ready"` — emitted when the HTTP server is listening. + * - `"shutdown"` — emitted by the core lifecycle manager during graceful + * shutdown, AFTER all plugin `shutdown()` hooks have completed. The emit is + * bounded by a short timeout (see the lifecycle manager's shutdown budget), + * so subscribers must not start long-running async work — finish quickly or + * be cut off. (The server plugin subscribes here to force-close its + * remaining sockets once peers have drained.) + */ type LifecycleEvent = "setup:complete" | "server:ready" | "shutdown"; /** @@ -222,6 +235,10 @@ export class PluginContext { /** * Register a lifecycle hook callback. + * + * See {@link LifecycleEvent} for event semantics. In particular, + * `"shutdown"` subscribers run inside a bounded shutdown phase and must + * not start long-running async work. */ onLifecycle(event: LifecycleEvent, fn: () => void | Promise): void { let hooks = this.lifecycleHooks.get(event); @@ -237,7 +254,8 @@ export class PluginContext { * Errors in individual callbacks are logged but do not prevent * other callbacks from running. * - * @internal Called by AppKit core only. + * @internal Called by AppKit core: `setup:complete` after plugin setup, + * and `shutdown` by the lifecycle manager during graceful shutdown. */ async emitLifecycle(event: LifecycleEvent): Promise { const hooks = this.lifecycleHooks.get(event); diff --git a/packages/appkit/src/core/tests/lifecycle-manager.test.ts b/packages/appkit/src/core/tests/lifecycle-manager.test.ts new file mode 100644 index 000000000..121e7eb5c --- /dev/null +++ b/packages/appkit/src/core/tests/lifecycle-manager.test.ts @@ -0,0 +1,384 @@ +import type { BasePlugin } from "shared"; +import { + afterEach, + beforeEach, + describe, + expect, + type MockInstance, + test, + vi, +} from "vitest"; + +// Mock core singletons before importing the subject under test. +vi.mock("../../cache", () => ({ + CacheManager: { + getInstanceSync: vi.fn().mockReturnValue({ + close: vi.fn().mockResolvedValue(undefined), + }), + }, +})); + +vi.mock("../../telemetry", () => ({ + TelemetryManager: { + getInstance: vi.fn().mockReturnValue({ + shutdown: vi.fn().mockResolvedValue(undefined), + }), + getProvider: vi.fn().mockReturnValue({ + getTracer: vi.fn().mockReturnValue({ startActiveSpan: vi.fn() }), + }), + }, +})); + +vi.mock("../../internal-telemetry", () => ({ + TelemetryReporter: { + getInstance: vi.fn().mockReturnValue(null), + }, +})); + +const { mockLoggerError } = vi.hoisted(() => ({ + mockLoggerError: vi.fn(), +})); + +vi.mock("../../logging/logger", () => ({ + createLogger: () => ({ + info: vi.fn(), + debug: vi.fn(), + warn: vi.fn(), + error: mockLoggerError, + }), +})); + +import { CacheManager } from "../../cache"; +import { TelemetryReporter } from "../../internal-telemetry"; +import { TelemetryManager } from "../../telemetry"; +import { LifecycleManager } from "../lifecycle-manager"; +import { PluginContext } from "../plugin-context"; + +function contextWithPlugins(plugins: Record>) { + const ctx = new PluginContext(); + for (const [name, instance] of Object.entries(plugins)) { + ctx.registerPlugin(name, instance as BasePlugin); + } + return ctx; +} + +describe("LifecycleManager", () => { + let exitSpy: MockInstance; + + beforeEach(() => { + mockLoggerError.mockClear(); + exitSpy = vi + .spyOn(process, "exit") + .mockImplementation(((_code?: number) => undefined) as any); + }); + + afterEach(() => { + exitSpy.mockRestore(); + vi.useRealTimers(); + vi.clearAllMocks(); + }); + + describe("shutdown", () => { + test("runs plugin shutdown() hooks concurrently and exits 0", async () => { + // Prove concurrency, not just "both called": hook B only resolves after + // hook A has started. If the manager awaited hooks serially (A fully + // before B), B would never observe A as started and this would hang. + let aStarted: (() => void) | undefined; + const aStartedGate = new Promise((resolve) => { + aStarted = resolve; + }); + const shutdownA = vi.fn(async () => { + aStarted?.(); + }); + const shutdownB = vi.fn(async () => { + await aStartedGate; + }); + const ctx = contextWithPlugins({ + a: { name: "a", shutdown: shutdownA }, + b: { name: "b", shutdown: shutdownB }, + "no-hooks": { name: "no-hooks" }, + }); + + await new LifecycleManager(ctx).shutdown(); + + expect(shutdownA).toHaveBeenCalledTimes(1); + expect(shutdownB).toHaveBeenCalledTimes(1); + expect(exitSpy).toHaveBeenCalledWith(0); + }); + + test("stops the internal-telemetry reporter before aborting", async () => { + const stop = vi.fn(); + const order: string[] = []; + stop.mockImplementation(() => order.push("reporter-stop")); + vi.mocked(TelemetryReporter.getInstance).mockReturnValueOnce({ + stop, + } as any); + const ctx = contextWithPlugins({ + a: { + name: "a", + abortActiveOperations: vi.fn(() => { + order.push("abort"); + }), + }, + }); + + await new LifecycleManager(ctx).shutdown(); + + expect(stop).toHaveBeenCalledTimes(1); + expect(order).toEqual(["reporter-stop", "abort"]); + expect(exitSpy).toHaveBeenCalledWith(0); + }); + + test("aborts active operations with error isolation", async () => { + const okAbort = vi.fn(); + const badAbort = vi.fn(() => { + throw new Error("boom"); + }); + const ctx = contextWithPlugins({ + ok: { name: "ok", abortActiveOperations: okAbort }, + bad: { name: "bad", abortActiveOperations: badAbort }, + }); + + await new LifecycleManager(ctx).shutdown(); + + expect(okAbort).toHaveBeenCalledTimes(1); + expect(badAbort).toHaveBeenCalledTimes(1); + expect( + mockLoggerError.mock.calls.some((c) => + String(c[0]).includes("Error aborting operations for plugin"), + ), + ).toBe(true); + expect(exitSpy).toHaveBeenCalledWith(0); + }); + + test("a failing plugin shutdown() is logged and does not abort shutdown", async () => { + const failing = vi.fn().mockRejectedValue(new Error("drain failed")); + const healthy = vi.fn().mockResolvedValue(undefined); + const ctx = contextWithPlugins({ + failing: { name: "failing", shutdown: failing }, + healthy: { name: "healthy", shutdown: healthy }, + }); + + await new LifecycleManager(ctx).shutdown(); + + expect(failing).toHaveBeenCalledTimes(1); + expect(healthy).toHaveBeenCalledTimes(1); + expect( + mockLoggerError.mock.calls.some((c) => + String(c[0]).includes("Error shutting down plugin"), + ), + ).toBe(true); + expect(exitSpy).toHaveBeenCalledWith(0); + }); + + test("a hanging plugin shutdown() does not block past its 10s timeout", async () => { + vi.useFakeTimers(); + const hanging = vi.fn(() => new Promise(() => {})); + const fast = vi.fn().mockResolvedValue(undefined); + const ctx = contextWithPlugins({ + hanging: { name: "hanging", shutdown: hanging }, + fast: { name: "fast", shutdown: fast }, + }); + + const done = new LifecycleManager(ctx).shutdown(); + await vi.advanceTimersByTimeAsync(10_000); + await done; + + expect(hanging).toHaveBeenCalledTimes(1); + expect(fast).toHaveBeenCalledTimes(1); + expect( + mockLoggerError.mock.calls.some( + (c) => + String(c[0]).includes("Error shutting down plugin") && + c[1] === "hanging" && + String(c[2]).includes("timed out"), + ), + ).toBe(true); + expect(exitSpy).toHaveBeenCalledWith(0); + }); + + test("emits the 'shutdown' lifecycle event", async () => { + const ctx = contextWithPlugins({}); + const hook = vi.fn(); + ctx.onLifecycle("shutdown", hook); + + await new LifecycleManager(ctx).shutdown(); + + expect(hook).toHaveBeenCalledTimes(1); + expect(exitSpy).toHaveBeenCalledWith(0); + }); + + test("is not re-entrant — a second call is a no-op", async () => { + const shutdownHook = vi.fn().mockResolvedValue(undefined); + const ctx = contextWithPlugins({ + a: { name: "a", shutdown: shutdownHook }, + }); + const manager = new LifecycleManager(ctx); + + await Promise.all([manager.shutdown(), manager.shutdown()]); + await manager.shutdown(); + + expect(shutdownHook).toHaveBeenCalledTimes(1); + }); + + test("closes the cache storage and flushes telemetry", async () => { + const close = vi.fn().mockResolvedValue(undefined); + const flush = vi.fn().mockResolvedValue(undefined); + vi.mocked(CacheManager.getInstanceSync).mockReturnValueOnce({ + close, + } as any); + vi.mocked(TelemetryManager.getInstance).mockReturnValueOnce({ + shutdown: flush, + } as any); + + await new LifecycleManager(contextWithPlugins({})).shutdown(); + + expect(close).toHaveBeenCalledTimes(1); + expect(flush).toHaveBeenCalledTimes(1); + expect(exitSpy).toHaveBeenCalledWith(0); + }); + + test("a hanging telemetry flush cannot hang shutdown — the flush timeout still exits 0", async () => { + vi.useFakeTimers(); + const hangingFlush = vi.fn(() => new Promise(() => {})); + vi.mocked(TelemetryManager.getInstance).mockReturnValueOnce({ + shutdown: hangingFlush, + } as any); + + const done = new LifecycleManager(contextWithPlugins({})).shutdown(); + await vi.advanceTimersByTimeAsync(2_000); + await done; + + expect(hangingFlush).toHaveBeenCalledTimes(1); + expect( + mockLoggerError.mock.calls.some( + (c) => + String(c[0]).includes("Error flushing telemetry") && + String(c[1]).includes("timed out"), + ), + ).toBe(true); + expect(exitSpy).toHaveBeenCalledWith(0); + }); + + test("a hanging cache close cannot hang shutdown — the close timeout still exits 0", async () => { + vi.useFakeTimers(); + const hangingClose = vi.fn(() => new Promise(() => {})); + vi.mocked(CacheManager.getInstanceSync).mockReturnValueOnce({ + close: hangingClose, + } as any); + + const done = new LifecycleManager(contextWithPlugins({})).shutdown(); + await vi.advanceTimersByTimeAsync(2_000); + await done; + + expect(hangingClose).toHaveBeenCalledTimes(1); + expect( + mockLoggerError.mock.calls.some( + (c) => + String(c[0]).includes("Error closing cache storage") && + String(c[1]).includes("timed out"), + ), + ).toBe(true); + expect(exitSpy).toHaveBeenCalledWith(0); + }); + + test("a never-initialized cache is skipped without error", async () => { + vi.mocked(CacheManager.getInstanceSync).mockImplementationOnce(() => { + throw new Error("cache not initialized"); + }); + + await new LifecycleManager(contextWithPlugins({})).shutdown(); + + expect( + mockLoggerError.mock.calls.some((c) => + String(c[0]).includes("Error closing cache storage"), + ), + ).toBe(false); + expect(exitSpy).toHaveBeenCalledWith(0); + }); + + test("runs phases in order: abort → plugin hooks → lifecycle emit → cache close + flush (concurrent) → exit", async () => { + const order: string[] = []; + const ctx = contextWithPlugins({ + a: { + name: "a", + abortActiveOperations: vi.fn(() => { + order.push("abort"); + }), + shutdown: vi.fn(async () => { + order.push("plugin-shutdown"); + }), + }, + }); + ctx.onLifecycle("shutdown", () => { + order.push("lifecycle"); + }); + vi.mocked(CacheManager.getInstanceSync).mockReturnValueOnce({ + close: vi.fn(async () => { + order.push("cache-close"); + }), + } as any); + vi.mocked(TelemetryManager.getInstance).mockReturnValueOnce({ + shutdown: vi.fn(async () => { + order.push("flush"); + }), + } as any); + exitSpy.mockImplementationOnce(((_code?: number) => { + order.push("exit"); + }) as any); + + await new LifecycleManager(ctx).shutdown(); + + expect(order.slice(0, 3)).toEqual([ + "abort", + "plugin-shutdown", + "lifecycle", + ]); + // Cache close and telemetry flush run concurrently — assert both land + // after the lifecycle emit and before exit; relative order unspecified. + expect(order.slice(3, 5).sort()).toEqual(["cache-close", "flush"]); + expect(order[5]).toBe("exit"); + expect(order).toHaveLength(6); + }); + + test("a plugin shutdown() that rejects after its timeout already won does not crash", async () => { + vi.useFakeTimers(); + let rejectLate: ((err: Error) => void) | undefined; + const lateRejecting = vi.fn( + () => + new Promise((_, reject) => { + rejectLate = reject; + }), + ); + const ctx = contextWithPlugins({ + late: { name: "late", shutdown: lateRejecting }, + }); + + const done = new LifecycleManager(ctx).shutdown(); + await vi.advanceTimersByTimeAsync(10_000); + await done; + + // The hook loses the race, then rejects afterwards — must be swallowed + // by the pre-attached no-op handler, not crash the process. + rejectLate?.(new Error("late rejection")); + await vi.advanceTimersByTimeAsync(0); + + expect(exitSpy).toHaveBeenCalledWith(0); + }); + }); + + describe("installSignalHandlers", () => { + test("registers SIGTERM/SIGINT once and triggers shutdown", async () => { + const onceSpy = vi.spyOn(process, "once"); + const ctx = contextWithPlugins({}); + const manager = new LifecycleManager(ctx); + + manager.installSignalHandlers(); + + const signals = onceSpy.mock.calls.map((c) => c[0]); + expect(signals).toContain("SIGTERM"); + expect(signals).toContain("SIGINT"); + onceSpy.mockRestore(); + }); + }); +}); diff --git a/packages/appkit/src/plugins/lakebase/lakebase.ts b/packages/appkit/src/plugins/lakebase/lakebase.ts index 9fac029a6..a827051d7 100644 --- a/packages/appkit/src/plugins/lakebase/lakebase.ts +++ b/packages/appkit/src/plugins/lakebase/lakebase.ts @@ -176,15 +176,24 @@ export class LakebasePlugin extends Plugin implements ToolProvider { /** * Gracefully drains and closes all connection pools (SP + OBO). - * Called automatically by AppKit during shutdown. + * + * Runs as the plugin's `shutdown()` hook (phase 3 of the core lifecycle + * manager's graceful shutdown), NOT in `abortActiveOperations()` (phase 2): + * other plugins' `shutdown()` hooks may still need database connections to + * drain state, + * so the pools must outlive the abort phase. `pg.Pool#end()` waits for + * checked-out clients to be released, so hooks running concurrently with + * this one can still finish their in-flight queries. Errors are caught + * and logged; this hook never throws. */ - abortActiveOperations(): void { - super.abortActiveOperations(); + async shutdown(): Promise { if (this.pool) { logger.info("Closing Lakebase SP pool"); - this.pool.end().catch((err) => { + try { + await this.pool.end(); + } catch (err) { logger.error("Error closing Lakebase SP pool: %O", err); - }); + } this.pool = null; } if (this.oboPoolManager) { @@ -192,9 +201,11 @@ export class LakebasePlugin extends Plugin implements ToolProvider { "Closing all Lakebase OBO pools (%d)", this.oboPoolManager.size, ); - this.oboPoolManager.closeAll().catch((err) => { + try { + await this.oboPoolManager.closeAll(); + } catch (err) { logger.error("Error closing Lakebase OBO pools: %O", err); - }); + } this.oboPoolManager = null; } } diff --git a/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts b/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts index 5f96d6efe..b8564ff8e 100644 --- a/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts +++ b/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts @@ -212,6 +212,53 @@ describe("LakebasePlugin — readOnly enforcement", () => { }); }); +describe("LakebasePlugin — shutdown", () => { + test("closes the SP pool and all OBO pools via shutdown()", async () => { + const { createLakebasePool, createLakebasePoolManager } = await import( + "../../../connectors/lakebase" + ); + const plugin = makePlugin({}); + await plugin.setup(); + + const spPool = vi.mocked(createLakebasePool).mock.results.at(-1)?.value as { + end: ReturnType; + }; + const oboManager = vi.mocked(createLakebasePoolManager).mock.results.at(-1) + ?.value as { closeAll: ReturnType }; + + await plugin.shutdown(); + + expect(spPool.end).toHaveBeenCalledTimes(1); + expect(oboManager.closeAll).toHaveBeenCalledTimes(1); + + // Idempotent: a second call has nothing left to close. + await plugin.shutdown(); + expect(spPool.end).toHaveBeenCalledTimes(1); + expect(oboManager.closeAll).toHaveBeenCalledTimes(1); + }); + + test("abortActiveOperations() does NOT tear down the pools (teardown is shutdown()'s job)", async () => { + const { createLakebasePool, createLakebasePoolManager } = await import( + "../../../connectors/lakebase" + ); + const plugin = makePlugin({}); + await plugin.setup(); + + const spPool = vi.mocked(createLakebasePool).mock.results.at(-1)?.value as { + end: ReturnType; + }; + const oboManager = vi.mocked(createLakebasePoolManager).mock.results.at(-1) + ?.value as { closeAll: ReturnType }; + + plugin.abortActiveOperations(); + + // Other plugins' shutdown() hooks may still need database connections + // to drain state — the pools must survive the abort phase. + expect(spPool.end).not.toHaveBeenCalled(); + expect(oboManager.closeAll).not.toHaveBeenCalled(); + }); +}); + describe("LakebasePlugin — destructive mode", () => { test("does NOT wrap in read-only transaction when readOnly: false", async () => { const queryMock = vi.fn((_text: string, _values?: unknown[]) => diff --git a/packages/appkit/src/plugins/server/index.ts b/packages/appkit/src/plugins/server/index.ts index e66abf5ad..926edbbe5 100644 --- a/packages/appkit/src/plugins/server/index.ts +++ b/packages/appkit/src/plugins/server/index.ts @@ -54,6 +54,14 @@ export class ServerPlugin extends Plugin { port: Number(process.env.DATABRICKS_APP_PORT) || 8000, }; + /** + * Budget for awaiting `server.close()` during shutdown. Bounded because + * `closeAllConnections()` runs immediately before the await, so the close + * is expected to resolve promptly; the core lifecycle manager's overall + * force-exit timer is the ultimate backstop if it does not. + */ + private static readonly SERVER_CLOSE_TIMEOUT_MS = 2_000; + /** Plugin manifest declaring metadata and resource requirements */ static manifest = manifest as PluginManifest<"server">; private serverApplication: express.Application; @@ -65,6 +73,12 @@ export class ServerPlugin extends Plugin { protected declare config: ServerConfig; private serverExtensions: ((app: express.Application) => void)[] = []; private rawBodyPaths: Set = new Set(); + /** + * Resolves when `server.close()` completes. Created in + * {@link abortActiveOperations} (which initiates the close) and awaited in + * {@link closeRemainingConnections} after other plugins have drained. + */ + private serverClosed?: Promise; static phase: PluginPhase = "deferred"; constructor(config: ServerConfig) { @@ -160,8 +174,15 @@ export class ServerPlugin extends Plugin { // attach server to remote tunnel controller this.remoteTunnelController.setServer(server); - process.on("SIGTERM", () => this._gracefulShutdown()); - process.on("SIGINT", () => this._gracefulShutdown()); + // Graceful shutdown is orchestrated by the core LifecycleManager, which + // owns the process signal handlers. This plugin only contributes its + // HTTP concerns: it aborts/initiates the socket close in + // abortActiveOperations(), drains dev servers in shutdown(), and force- + // closes remaining sockets in the "shutdown" lifecycle hook below (which + // fires after every plugin's shutdown() has run). + this.context?.onLifecycle("shutdown", () => + this.closeRemainingConnections(), + ); if (process.env.NODE_ENV === "development") { const allRoutes = getRoutes(this.serverApplication._router.stack); @@ -398,52 +419,71 @@ export class ServerPlugin extends Plugin { } } - private async _gracefulShutdown() { - logger.info("Starting graceful shutdown..."); + /** + * Cancel in-flight work and begin closing the HTTP server. + * + * Runs in the first phase of the core LifecycleManager's graceful shutdown, + * before any plugin `shutdown()` hook. In addition to the base behavior + * (aborting SSE streams), it stops the server accepting new connections and + * drops idle keep-alive sockets — without `closeIdleConnections()`, a + * connected browser would pin `server.close()` open until the force-exit + * timer fires. The close is only *initiated* here; the returned promise is + * awaited later in {@link closeRemainingConnections} so peer plugins can + * still drain in-flight requests through the server first. + */ + abortActiveOperations(): void { + super.abortActiveOperations(); + if (this.server) { + const server = this.server; + this.serverClosed = new Promise((resolve) => { + server.close(() => resolve()); + }); + server.closeIdleConnections(); + } + } + + /** + * Drain the dev-only servers as this plugin's `shutdown()` hook. + * + * Runs concurrently with the other plugins' hooks during graceful shutdown. + * These are no-ops in production (neither is created), so this hook is + * effectively dev-only. + */ + async shutdown(): Promise { if (this.viteDevServer) { await this.viteDevServer.close(); } - if (this.remoteTunnelController) { this.remoteTunnelController.cleanup(); } + } - TelemetryReporter.getInstance()?.stop(); - - // 1. abort active operations from plugins - const shutdownPlugins = this.context?.getPlugins(); - if (shutdownPlugins) { - for (const plugin of shutdownPlugins.values()) { - if (plugin.abortActiveOperations) { - try { - plugin.abortActiveOperations(); - } catch (err) { - logger.error( - "Error aborting operations for plugin %s: %O", - plugin.name, - err, - ); - } - } - } - } - - // 2. close the server + /** + * Force-close whatever sockets remain and await the server close. + * + * Registered on the `"shutdown"` lifecycle event in {@link start}, so it + * fires after every plugin's `shutdown()` hook has run — meaning any + * in-flight request a peer plugin was draining has already completed. + * `closeAllConnections()` destroys the leftover sockets (aborted SSE + * responses, keep-alive connections) synchronously so the pending + * `server.close()` can complete; the await is bounded because the close is + * expected to resolve promptly once the sockets are gone. + */ + private async closeRemainingConnections(): Promise { if (this.server) { - this.server.close(() => { - logger.debug("Server closed gracefully"); - process.exit(0); - }); - - // 3. timeout to force shutdown after 15 seconds - setTimeout(() => { - logger.debug("Force shutdown after timeout"); - process.exit(1); - }, 15000); - } else { - process.exit(0); + this.server.closeAllConnections(); } + if (!this.serverClosed) return; + + await Promise.race([ + this.serverClosed, + new Promise((resolve) => { + const timer = setTimeout(resolve, ServerPlugin.SERVER_CLOSE_TIMEOUT_MS); + timer.unref(); + }), + ]); + logger.debug("Server closed gracefully"); } /** diff --git a/packages/appkit/src/plugins/server/tests/server.test.ts b/packages/appkit/src/plugins/server/tests/server.test.ts index bbc961723..df8b6d3f2 100644 --- a/packages/appkit/src/plugins/server/tests/server.test.ts +++ b/packages/appkit/src/plugins/server/tests/server.test.ts @@ -12,6 +12,8 @@ const { } = vi.hoisted(() => { const httpServer = { close: vi.fn((cb: any) => cb?.()), + closeIdleConnections: vi.fn(), + closeAllConnections: vi.fn(), on: vi.fn(), address: vi.fn().mockReturnValue({ port: 8000 }), }; @@ -85,6 +87,9 @@ vi.mock("express", () => { // Mock dependencies before imports vi.mock("../../../telemetry", () => ({ TelemetryManager: { + getInstance: vi.fn().mockReturnValue({ + shutdown: vi.fn().mockResolvedValue(undefined), + }), getProvider: vi.fn().mockReturnValue({ getTracer: vi.fn().mockReturnValue({ startActiveSpan: vi.fn() }), getMeter: vi.fn().mockReturnValue({ @@ -107,6 +112,7 @@ vi.mock("../../../cache", () => ({ get: vi.fn(), set: vi.fn(), delete: vi.fn(), + close: vi.fn().mockResolvedValue(undefined), }), }, })); @@ -186,6 +192,7 @@ vi.mock("../client-config-sanitizer", () => ({ import fs from "node:fs"; import express from "express"; +import { LifecycleManager } from "../../../core/lifecycle-manager"; import { sanitizeClientConfig } from "../client-config-sanitizer"; import { ServerPlugin } from "../index"; import { RemoteTunnelController } from "../remote-tunnel/remote-tunnel-controller"; @@ -694,41 +701,148 @@ describe("ServerPlugin", () => { }); }); - describe("_gracefulShutdown", () => { - test("aborts plugin operations (with error isolation) and closes server", async () => { - vi.useFakeTimers(); + describe("shutdown hooks", () => { + beforeEach(() => { mockLoggerError.mockClear(); - const exitSpy = vi - .spyOn(process, "exit") - .mockImplementation(((_code?: number) => undefined) as any); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + test("abortActiveOperations() stops accepting connections and initiates close", () => { const plugin = new ServerPlugin({ - context: createContextWithPlugins({ - ok: { - name: "ok", - abortActiveOperations: vi.fn(), - }, - bad: { - name: "bad", - abortActiveOperations: vi.fn(() => { - throw new Error("boom"); - }), - }, - }), + context: createContextWithPlugins({}), } as any); + (plugin as any).server = mockHttpServer; + + plugin.abortActiveOperations(); - // pretend started + // Idle keep-alive sockets are dropped and close() is initiated so a + // connected browser cannot pin the server open. + expect(mockHttpServer.closeIdleConnections).toHaveBeenCalledTimes(1); + expect(mockHttpServer.close).toHaveBeenCalledTimes(1); + // The full socket teardown is deferred to the lifecycle hook. + expect(mockHttpServer.closeAllConnections).not.toHaveBeenCalled(); + }); + + test("abortActiveOperations() is a no-op on sockets when no server started", () => { + const plugin = new ServerPlugin({ + context: createContextWithPlugins({}), + } as any); + + expect(() => plugin.abortActiveOperations()).not.toThrow(); + expect(mockHttpServer.close).not.toHaveBeenCalled(); + }); + + test("shutdown() drains the dev servers", async () => { + const plugin = new ServerPlugin({ + context: createContextWithPlugins({}), + } as any); + const viteClose = vi.fn().mockResolvedValue(undefined); + const tunnelCleanup = vi.fn(); + (plugin as any).viteDevServer = { close: viteClose }; + (plugin as any).remoteTunnelController = { cleanup: tunnelCleanup }; + + await plugin.shutdown(); + + expect(viteClose).toHaveBeenCalledTimes(1); + expect(tunnelCleanup).toHaveBeenCalledTimes(1); + }); + + test("closeRemainingConnections() force-closes sockets and awaits the close", async () => { + const plugin = new ServerPlugin({ + context: createContextWithPlugins({}), + } as any); (plugin as any).server = mockHttpServer; - await (plugin as any)._gracefulShutdown(); - vi.runAllTimers(); + // abort initiates the close and captures the serverClosed promise + plugin.abortActiveOperations(); + await (plugin as any).closeRemainingConnections(); + + expect(mockHttpServer.closeAllConnections).toHaveBeenCalledTimes(1); + expect(mockHttpServer.close).toHaveBeenCalledTimes(1); + }); + + test("closeRemainingConnections() does not hang if close() never completes", async () => { + vi.useFakeTimers(); + // A server whose close callback never fires. + const stuckServer = { + ...mockHttpServer, + close: vi.fn(), + closeIdleConnections: vi.fn(), + closeAllConnections: vi.fn(), + }; + const plugin = new ServerPlugin({ + context: createContextWithPlugins({}), + } as any); + (plugin as any).server = stuckServer; + + plugin.abortActiveOperations(); + const done = (plugin as any).closeRemainingConnections(); + // Bounded by SERVER_CLOSE_TIMEOUT_MS (2s) even though close never fires. + await vi.advanceTimersByTimeAsync(2_000); + await done; - expect(mockLoggerError).toHaveBeenCalled(); - expect(mockHttpServer.close).toHaveBeenCalled(); - expect(exitSpy).toHaveBeenCalled(); + expect(stuckServer.closeAllConnections).toHaveBeenCalledTimes(1); + }); + + test("start() registers closeRemainingConnections on the 'shutdown' lifecycle event", async () => { + const ctx = createContextWithPlugins({}); + const onLifecycleSpy = vi.spyOn(ctx, "onLifecycle"); + const plugin = new ServerPlugin({ context: ctx } as any); + await plugin.start(); + + expect(onLifecycleSpy).toHaveBeenCalledWith( + "shutdown", + expect.any(Function), + ); + + // Emitting the event drives the socket teardown end-to-end. + await ctx.emitLifecycle("shutdown"); + expect(mockHttpServer.closeAllConnections).toHaveBeenCalledTimes(1); + }); + + test("real LifecycleManager drives the ServerPlugin's socket teardown in order", async () => { + // End-to-end across the contract seam: a real LifecycleManager driving a + // real ServerPlugin + real PluginContext, with a peer plugin whose + // shutdown() hook must land BETWEEN the server's closeIdle (abort phase) + // and closeAll (lifecycle-emit phase). Mocking only one side would let a + // dropped registration or phase reorder pass — this catches it. + const order: string[] = []; + mockHttpServer.closeIdleConnections.mockImplementationOnce(() => { + order.push("closeIdle"); + }); + mockHttpServer.closeAllConnections.mockImplementationOnce(() => { + order.push("closeAll"); + }); + + const ctx = new PluginContext(); + const server = new ServerPlugin({ context: ctx } as any); + ctx.registerPlugin("server", server as unknown as BasePlugin); + ctx.registerPlugin("peer", { + name: "peer", + shutdown: vi.fn(async () => { + order.push("peer-shutdown"); + }), + } as unknown as BasePlugin); + + const exitSpy = vi.spyOn(process, "exit").mockImplementation((( + _code?: number, + ) => { + order.push("exit"); + return undefined; + }) as any); + + await server.start(); + await new LifecycleManager(ctx).shutdown(); + + // closeIdle fires in the abort phase, the peer drains next, closeAll only + // fires in the later lifecycle-emit phase, then the process exits 0. + expect(order).toEqual(["closeIdle", "peer-shutdown", "closeAll", "exit"]); + expect(exitSpy).toHaveBeenCalledWith(0); exitSpy.mockRestore(); - vi.useRealTimers(); }); }); }); diff --git a/packages/appkit/src/telemetry/telemetry-manager.ts b/packages/appkit/src/telemetry/telemetry-manager.ts index 6660b6b2c..a2642f2ab 100644 --- a/packages/appkit/src/telemetry/telemetry-manager.ts +++ b/packages/appkit/src/telemetry/telemetry-manager.ts @@ -35,6 +35,7 @@ export class TelemetryManager { private static instance?: TelemetryManager; private sdk?: NodeSDK; + private shutdownPromise?: Promise; /** * Create a scoped telemetry provider for a specific plugin. @@ -95,7 +96,6 @@ export class TelemetryManager { }); this.sdk.start(); - this.registerShutdown(); logger.debug("Initialized successfully"); } catch (error) { logger.error("Failed to initialize: %O", error); @@ -158,24 +158,27 @@ export class TelemetryManager { ]; } - private registerShutdown() { - const shutdownFn = async () => { - await TelemetryManager.getInstance().shutdown(); - }; - process.once("SIGTERM", shutdownFn); - process.once("SIGINT", shutdownFn); - } - - private async shutdown(): Promise { - if (!this.sdk) { - return; - } - - try { - await this.sdk.shutdown(); + /** + * Flush and shut down the OpenTelemetry SDK. + * + * Idempotent: the SDK reference is cleared synchronously and concurrent + * or repeated calls await the same in-flight flush. Awaited by the core + * lifecycle manager during graceful shutdown — that manager owns the + * process signal handlers, so telemetry no longer registers its own. + */ + async shutdown(): Promise { + if (this.sdk) { + const sdk = this.sdk; this.sdk = undefined; - } catch (error) { - logger.error("Error shutting down: %O", error); + this.shutdownPromise = (async () => { + try { + await sdk.shutdown(); + } catch (error) { + logger.error("Error shutting down: %O", error); + } + })(); } + + return this.shutdownPromise; } } diff --git a/packages/shared/src/plugin.ts b/packages/shared/src/plugin.ts index da4800548..3c837e5f1 100644 --- a/packages/shared/src/plugin.ts +++ b/packages/shared/src/plugin.ts @@ -16,8 +16,25 @@ export type { ResourceFieldEntry, DiscoveryDescriptor, PluginScaffoldingRules }; export interface BasePlugin { name: string; + /** + * Cancel in-flight work (abort signals, SSE streams). Runs in the first + * phase of graceful shutdown, BEFORE any plugin's `shutdown()` hook — + * so it must not tear down shared resources (e.g. connection pools) + * that other plugins' hooks may still need. Put teardown in `shutdown()`. + */ abortActiveOperations?(): void; + /** + * Optional graceful-shutdown hook. + * + * Invoked by AppKit's core lifecycle manager during graceful shutdown + * (SIGTERM/SIGINT), after `abortActiveOperations()` has run. Use it to + * drain in-flight work or release resources. Each plugin's hook is bounded + * by a per-plugin timeout and runs concurrently with other plugins' hooks; + * errors are logged and never abort the shutdown. + */ + shutdown?(): Promise | void; + setup(): Promise; injectRoutes(router: express.Router): void;