fix: add native lifecycle-aware health endpoint (#7852)

* fix: add native health endpoint

* fix: keep health endpoint dynamic

---------

Co-authored-by: Ravi Tharuma <RaviTharuma@users.noreply.github.com>

src/app/healthz/route.ts

@@ -0,0 +1,31 @@
+import { getServerLifecyclePhase } from "@/lib/serverLifecycle";
+
+export const dynamic = "force-dynamic";
+
+const HEALTH_BODIES = {
+  ready: "ok\n",
+  starting: "starting\n",
+  stopping: "stopping\n",
+} as const;
+
+function createHealthResponse(method: "GET" | "HEAD"): Response {
+  const phase = getServerLifecyclePhase();
+  const body = HEALTH_BODIES[phase];
+
+  return new Response(method === "HEAD" ? null : body, {
+    status: phase === "ready" ? 200 : 503,
+    headers: {
+      "Cache-Control": "no-store",
+      "Content-Length": String(body.length),
+      "Content-Type": "text/plain; charset=utf-8",
+    },
+  });
+}
+
+export function GET(): Response {
+  return createHealthResponse("GET");
+}
+
+export function HEAD(): Response {
+  return createHealthResponse("HEAD");
+}

src/instrumentation-node.ts

@@ -6,6 +6,8 @@
  * and emit spurious "not supported in Edge Runtime" warnings.
  */
 
+import { markServerReady, markServerStarting } from "@/lib/serverLifecycle";
+
 function getRandomBytes(byteLength: number): Uint8Array {
   const bytes = new Uint8Array(byteLength);
   globalThis.crypto.getRandomValues(bytes);
@@ -214,6 +216,8 @@ export async function warmModelCatalogCache(): Promise<void> {
 }
 
 export async function registerNodejs(): Promise<void> {
+  markServerStarting();
+
   // Rename the process title so OmniRoute is identifiable in ps/htop instead
   // of the generic "next-server" standalone server name.
   process.title = renameProcessTitle(process.title);
@@ -565,4 +569,6 @@ export async function registerNodejs(): Promise<void> {
       console.warn("[STARTUP] Live dashboard WebSocket daemon failed to start (non-fatal):", msg);
     }
   }
+
+  markServerReady();
 }

src/lib/gracefulShutdown.ts

@@ -12,6 +12,8 @@
  * @module lib/gracefulShutdown
  */
 
+import { markServerStopping } from "@/lib/serverLifecycle";
+
 /** Grace period before forced exit (default 30s, configurable) */
 const SHUTDOWN_TIMEOUT_MS = parseInt(process.env.SHUTDOWN_TIMEOUT_MS || "30000", 10);
 
@@ -134,6 +136,7 @@ export function initGracefulShutdown(): void {
   const shutdown = async (signal: string) => {
     if (state.shuttingDown) return;
     state.shuttingDown = true;
+    markServerStopping();
 
     console.log(`\n[Shutdown] Received ${signal}. Draining ${state.activeRequests} request(s)...`);
 

src/lib/serverLifecycle.ts

@@ -0,0 +1,23 @@
+export type ServerLifecyclePhase = "starting" | "ready" | "stopping";
+
+declare global {
+  var __omnirouteServerLifecycle: ServerLifecyclePhase | undefined;
+}
+
+export function getServerLifecyclePhase(): ServerLifecyclePhase {
+  return globalThis.__omnirouteServerLifecycle ?? "starting";
+}
+
+export function markServerStarting(): void {
+  globalThis.__omnirouteServerLifecycle = "starting";
+}
+
+export function markServerReady(): void {
+  if (getServerLifecyclePhase() !== "stopping") {
+    globalThis.__omnirouteServerLifecycle = "ready";
+  }
+}
+
+export function markServerStopping(): void {
+  globalThis.__omnirouteServerLifecycle = "stopping";
+}

tests/unit/healthz-route.test.ts

@@ -0,0 +1,92 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import fs from "node:fs";
+import { createRequire } from "node:module";
+
+import {
+  getServerLifecyclePhase,
+  markServerReady,
+  markServerStopping,
+} from "../../src/lib/serverLifecycle.ts";
+
+const routeModule = await import("../../src/app/healthz/route.ts");
+const require = createRequire(import.meta.url);
+const { getMiddlewareMatchers } = require("next/dist/build/analysis/get-page-static-info.js");
+const {
+  getMiddlewareRouteMatcher,
+} = require("next/dist/shared/lib/router/utils/middleware-route-matcher.js");
+
+test("/healthz follows the native server lifecycle without static caching", async () => {
+  assert.equal(routeModule.dynamic, "force-dynamic");
+  assert.equal(getServerLifecyclePhase(), "starting");
+
+  const starting = await routeModule.GET(new Request("http://localhost/healthz"));
+  assert.equal(starting.status, 503);
+  assert.equal(starting.headers.get("Cache-Control"), "no-store");
+  assert.equal(starting.headers.get("Content-Type"), "text/plain; charset=utf-8");
+  assert.equal(starting.headers.get("Content-Length"), "9");
+  assert.equal(await starting.text(), "starting\n");
+
+  markServerReady();
+  const ready = await routeModule.GET(new Request("http://localhost/healthz"));
+  assert.equal(ready.status, 200);
+  assert.equal(ready.headers.get("Content-Length"), "3");
+  assert.equal(await ready.text(), "ok\n");
+
+  const readyHead = await routeModule.HEAD(
+    new Request("http://localhost/healthz", { method: "HEAD" })
+  );
+  assert.equal(readyHead.status, 200);
+  assert.equal(readyHead.headers.get("Content-Length"), "3");
+  assert.equal(await readyHead.text(), "");
+
+  markServerStopping();
+  const stopping = await routeModule.GET(new Request("http://localhost/healthz"));
+  assert.equal(stopping.status, 503);
+  assert.equal(stopping.headers.get("Content-Length"), "9");
+  assert.equal(await stopping.text(), "stopping\n");
+
+  const stoppingHead = await routeModule.HEAD(
+    new Request("http://localhost/healthz", { method: "HEAD" })
+  );
+  assert.equal(stoppingHead.status, 503);
+  assert.equal(stoppingHead.headers.get("Content-Length"), "9");
+  assert.equal(await stoppingHead.text(), "");
+
+  markServerReady();
+  assert.equal(getServerLifecyclePhase(), "stopping");
+});
+
+test("native startup and shutdown hooks drive the health lifecycle", () => {
+  const startupSource = fs.readFileSync("src/instrumentation-node.ts", "utf8");
+  const registerStart = startupSource.indexOf("export async function registerNodejs");
+  const markStarting = startupSource.indexOf("markServerStarting();", registerStart);
+  const firstStartupAwait = startupSource.indexOf("await ", registerStart);
+  const markReady = startupSource.lastIndexOf("markServerReady();");
+
+  assert.ok(markStarting > registerStart && markStarting < firstStartupAwait);
+  assert.ok(markReady > firstStartupAwait);
+  assert.equal(startupSource.slice(markReady).trim(), "markServerReady();\n}");
+
+  const shutdownSource = fs.readFileSync("src/lib/gracefulShutdown.ts", "utf8");
+  const drainingFlag = shutdownSource.indexOf("state.shuttingDown = true;");
+  const markStopping = shutdownSource.indexOf("markServerStopping();", drainingFlag);
+  const drainAwait = shutdownSource.indexOf("await waitForDrain();", drainingFlag);
+
+  assert.ok(markStopping > drainingFlag && markStopping < drainAwait);
+});
+
+test("/healthz bypasses the centralized auth proxy matcher", () => {
+  const proxySource = fs.readFileSync("src/proxy.ts", "utf8");
+  const matcherBlock = proxySource.match(/matcher:\s*\[([\s\S]*?)\]/)?.[1];
+  assert.ok(matcherBlock, "proxy matcher configuration must remain discoverable");
+
+  const configuredMatchers = Array.from(
+    matcherBlock.matchAll(/[\"']([^\"']+)[\"']/g),
+    (match) => match[1]
+  );
+  const matcher = getMiddlewareRouteMatcher(getMiddlewareMatchers(configuredMatchers, "/"));
+
+  assert.equal(matcher("/healthz", { headers: {} }), false);
+  assert.equal(matcher("/api/system/version", { headers: {} }), true);
+});