fix(proxy): keep serving briefly after SIGTERM so the gateway moves off a draining pod

cli/commands/serve/split-mode.test.ts

@@ -5,7 +5,12 @@ import "#veryfront/schemas/_test-setup.ts";
 
 import { assertEquals } from "#veryfront/testing/assert.ts";
 import { describe, it } from "#veryfront/testing/bdd.ts";
-import { buildSplitModeEnvForTests, runSplitMode, waitForPort } from "./split-mode.ts";
+import {
+  buildSplitModeEnvForTests,
+  buildSplitModeProxyEnvForTests,
+  runSplitMode,
+  waitForPort,
+} from "./split-mode.ts";
 
 describe("serve-split command", () => {
   describe("runSplitMode", () => {
@@ -38,6 +43,16 @@ describe("serve-split command", () => {
       assertEquals(env.VERYFRONT_SERVER_URL, "http://localhost:3000");
       assertEquals(env.CACHE_TYPE, undefined);
     });
+
+    it("starts the proxy child without the rollout admission grace", () => {
+      const env = buildSplitModeEnvForTests({ SHUTDOWN_ADMISSION_GRACE_MS: "10000" }, 3000);
+      const proxyEnv = buildSplitModeProxyEnvForTests(env);
+
+      assertEquals(proxyEnv.SHUTDOWN_ADMISSION_GRACE_MS, "0");
+      assertEquals(proxyEnv.VERYFRONT_SERVER_URL, "http://localhost:3000");
+      // The production server child keeps the shared environment unchanged.
+      assertEquals(env.SHUTDOWN_ADMISSION_GRACE_MS, "10000");
+    });
   });
 
   describe("waitForPort", () => {

cli/commands/serve/split-mode.ts

@@ -42,6 +42,17 @@ const STATIC_ENV: Record<string, string> = {
   PROJECT_CIRCUIT_RESET_MS: "15000",
 };
 
+/**
+ * Local split mode has no Kubernetes gateway to move off the proxy, and its
+ * signal handler forwards SIGTERM, so the proxy skips the rollout admission
+ * grace and exits as soon as in-flight requests finish.
+ */
+export function buildSplitModeProxyEnvForTests(
+  env: Record<string, string>,
+): Record<string, string> {
+  return { ...env, SHUTDOWN_ADMISSION_GRACE_MS: "0" };
+}
+
 export function buildSplitModeEnvForTests(
   userEnv: Record<string, string>,
   productionServerPort: number,
@@ -179,7 +190,7 @@ export async function runSplitMode(options: SplitModeOptions): Promise<void> {
   // Start proxy
   const proxyProcess = startProcess(
     [...veryfront, "serve", "--mode=proxy", `--port=${proxyPort}`],
-    env,
+    buildSplitModeProxyEnvForTests(env),
     `proxy on :${proxyPort}`,
   );
 

scripts/build/compile-binary.test.ts

@@ -476,6 +476,8 @@ it("compiled proxy smoke covers cache and observability providers", async () =>
       "run_smoke sentry",
       "VERYFRONT_ERROR_REPORTER=sentry",
       "SENTRY_DSN=https://public@example.com/1",
+      // Five SIGTERMed probes must not each wait out the 10s admission grace.
+      "SHUTDOWN_ADMISSION_GRACE_MS=0",
     ]
   ) {
     assertEquals(
@@ -555,6 +557,8 @@ it("proxy release enforces the cold-start cgroup budget", async () => {
       '--memory "$memory_limit"',
       "{{.State.OOMKilled}}",
       '"/_proxy/health"',
+      // The SIGTERM admission grace must fit inside `docker stop --time 5`.
+      "--env SHUTDOWN_ADMISSION_GRACE_MS=1000",
     ]
   ) {
     assertEquals(

scripts/build/smoke-proxy-binary.sh

@@ -31,7 +31,8 @@ run_smoke() {
   shift 3
   local log_file="${tmp_dir}/${name}.log"
 
-  env PORT="$port" HOST=127.0.0.1 NODE_ENV=development "$@" \
+  # Each probe SIGTERMs an idle proxy; skip the rollout admission grace.
+  env PORT="$port" HOST=127.0.0.1 NODE_ENV=development SHUTDOWN_ADMISSION_GRACE_MS=0 "$@" \
     "$binary" >"$log_file" 2>&1 &
   proxy_pid=$!
 

scripts/build/smoke-proxy-memory.sh

@@ -59,6 +59,7 @@ for ((attempt = 1; attempt <= attempts; attempt++)); do
     --env HOST=0.0.0.0 \
     --env NODE_ENV=development \
     --env PORT="$port" \
+    --env SHUTDOWN_ADMISSION_GRACE_MS=1000 \
     --entrypoint /usr/local/bin/veryfront-proxy \
     "$container_image" >/dev/null
 
@@ -93,6 +94,8 @@ for ((attempt = 1; attempt <= attempts; attempt++)); do
     exit 1
   fi
 
+  # The SIGTERM admission grace above (1s) must fit inside the 5s stop window;
+  # the production default (10s) would make docker SIGKILL the idle proxy.
   docker stop --time 5 "$container" >/dev/null
   read -r oom_killed exit_code < <(
     docker inspect --format '{{.State.OOMKilled}} {{.State.ExitCode}}' "$container"

src/agent/react/use-chat/use-chat.csrf.test.tsx

@@ -101,6 +101,17 @@ describe("react/agent/useChat CSRF double-submit", () => {
     }
   });
 
+  it("declares the SSE reply on every chat endpoint, custom paths included", async () => {
+    const restoreDom = installDom();
+    try {
+      const headers = await captureSendHeaders({ api: "/api/support-chat" });
+      assertEquals(headers.get("accept"), "text/event-stream");
+      assertEquals(headers.get("content-type"), "application/json");
+    } finally {
+      restoreDom();
+    }
+  });
+
   it("leaves the header alone when the page has no CSRF cookie", async () => {
     const restoreDom = installDom();
     try {

src/agent/react/use-chat/use-chat.ts

@@ -256,6 +256,10 @@ function useChatState(options: UseChatOptions): ResettableUseChatResult {
           // discovers configured and origin-specific cookie names.
           headers: csrfMutationHeaders(api, {
             headers: {
+              // The reply is an SSE stream. Saying so lets a draining proxy
+              // send the chat to a replica that stays, whatever `api` path the
+              // app configured.
+              Accept: "text/event-stream",
               "Content-Type": "application/json",
               ...options.headers,
             },

src/agent/react/use-completion.ts

@@ -77,6 +77,10 @@ export function useCompletion(options: UseCompletionOptions): UseCompletionResul
           // header itself through `options.headers` keeps its own value.
           headers: csrfMutationHeaders(options.api, {
             headers: {
+              // The hook decodes raw text, so plain text is preferred. Listing
+              // SSE marks the request as a stream, which lets a draining proxy
+              // send it to a replica that stays.
+              Accept: "text/plain, text/event-stream;q=0.9, */*;q=0.8",
               "Content-Type": "application/json",
               ...options.headers,
             },

src/agent/react/use-streaming.ts

@@ -61,7 +61,13 @@ export function useStreaming(options: UseStreamingOptions): UseStreamingResult {
           // so this POST has to echo the configured CSRF cookie the document
           // response issued or the server answers 403.
           headers: csrfMutationHeaders(options.url, {
-            headers: { "Content-Type": "application/json" },
+            headers: {
+              // The hook decodes raw text, so plain text is preferred. Listing
+              // SSE marks the request as a stream, which lets a draining proxy
+              // send it to a replica that stays.
+              Accept: "text/plain, text/event-stream;q=0.9, */*;q=0.8",
+              "Content-Type": "application/json",
+            },
           }),
           body: body ? JSON.stringify(body) : undefined,
           signal: abortController.signal,

src/proxy/main.test.ts

@@ -56,7 +56,15 @@ describe("proxy main request URL parsing", () => {
     // only move the crash (veryfront-issue-inbox#828).
     const catchIndex = source.indexOf("proxyRequestDrainTracker.complete(requestId);");
     assertEquals(catchIndex >= 0, true);
-    const routerCatch = source.slice(catchIndex, catchIndex + 700);
+    // While draining, error responses close their connection too.
+    assertStringIncludes(
+      source.slice(catchIndex, catchIndex + 300),
+      "return drainAdmission.closeIfDraining(req, handleRouterError(req, url, error));",
+    );
+    const handlerIndex = source.indexOf("function handleRouterError(");
+    assertEquals(handlerIndex > catchIndex, true);
+    const routerCatch = source.slice(handlerIndex, handlerIndex + 1_200);
+    assertStringIncludes(routerCatch, "if (!drainAdmission.isDraining()) throw error;");
     assertStringIncludes(routerCatch, "error instanceof ProxyRequestHostError");
     assertStringIncludes(
       routerCatch,
@@ -141,7 +149,32 @@ describe("proxy main request URL parsing", () => {
   it("drains tracked responses before closing the proxy server", async () => {
     const source = await Deno.readTextFile(new URL("./main.ts", import.meta.url));
 
-    assertStringIncludes(source, "if (shuttingDown) return createProxyDrainingResponse()");
+    // Readiness is answered before the draining short-circuit so it can fail
+    // with 503 while liveness (/_proxy/health) stays a static 200. After the
+    // admission grace window (and for new WebSockets and SSE streams, which the
+    // admission decides per request) requests get the 503.
+    const readyIndex = source.indexOf(
+      "if (url.pathname === PROXY_READY_PATH) return createProxyReadinessResponse(shuttingDown);",
+    );
+    const refuseIndex = source.indexOf(
+      "if (shuttingDown && !drainAdmission.admits(req, url.pathname)) {\n    return createProxyDrainingResponse();",
+    );
+    const trackIndex = source.indexOf("proxyRequestDrainTracker.start(requestId");
+    assertEquals(readyIndex >= 0 && refuseIndex > readyIndex && trackIndex > refuseIndex, true);
+    // Admitted drain responses close their connection so the gateway reconnects
+    // through the Service, which no longer routes to this pod.
+    assertStringIncludes(
+      source,
+      "response = drainAdmission.closeIfDraining(req, response, admittedDuringDrain);\n    return proxyRequestDrainTracker.completeOnResponseEnd(requestId, req, response);",
+    );
+    // Admission time, not response time, decides whether an SSE response is
+    // swapped for the draining response.
+    const admittedIndex = source.indexOf(
+      "const admittedDuringDrain = drainAdmission.isDraining();",
+    );
+    assertEquals(admittedIndex > refuseIndex && trackIndex > admittedIndex, true);
+    // The signal decides the grace: SIGINT (Ctrl-C) gets none.
+    assertStringIncludes(source, "shuttingDown = true;\n  drainAdmission.begin(signal);");
     assertStringIncludes(
       source,
       "proxyRequestDrainTracker.start(requestId, req.method, url.pathname)",
@@ -155,7 +188,7 @@ describe("proxy main request URL parsing", () => {
     // close the server underneath in-flight responses.
     assertStringIncludes(
       source,
-      "await proxyRequestDrainTracker.waitForDrain(SHUTDOWN_DRAIN_TIMEOUT_MS)",
+      "await proxyRequestDrainTracker.waitForDrain(\n        drainAdmission.remainingDrainBudgetMs(SHUTDOWN_DRAIN_TIMEOUT_MS),",
       "shutdown must wait for the configured drain budget, not a zero timeout",
     );
     assertStringIncludes(
@@ -164,7 +197,14 @@ describe("proxy main request URL parsing", () => {
       "the default drain budget must stay a real waiting window",
     );
 
+    // The grace wait uses the timer captured before extensions load, never the
+    // live global setTimeout.
+    const graceIndex = source.indexOf("await drainAdmission.waitForWindow();");
+    assertEquals(source.includes("setTimeout(resolve, graceMs)"), false);
     const drainIndex = source.indexOf("await proxyRequestDrainTracker.waitForDrain");
+    assertEquals(graceIndex >= 0 && drainIndex > graceIndex, true);
+    // Operators can discover the setting next to the other shutdown variables.
+    assertStringIncludes(source, " * - SHUTDOWN_ADMISSION_GRACE_MS: ");
     const closeIndex = source.indexOf("await closeProxyServerWithin");
     const handlerCloseIndex = source.indexOf('name: "proxy_handler"');
     const shutdownHooksIndex = source.indexOf('name: "extension_owners"');

src/proxy/main.ts

@@ -21,6 +21,12 @@
  * - VERYFRONT_API_INTERNAL_USER: Basic auth user for internal API
  * - VERYFRONT_API_INTERNAL_PASS: Basic auth pass for internal API
  * - SHUTDOWN_DRAIN_TIMEOUT_MS: Time to wait for active SSE responses during shutdown
+ *   (default 25000, max 600000)
+ * - SHUTDOWN_ADMISSION_GRACE_MS: After SIGTERM, keep serving new ordinary requests
+ *   (each closing its connection) for this long while the gateway moves off the
+ *   draining pod (default 10000, max 600000, 0 disables). It is spent inside
+ *   SHUTDOWN_DRAIN_TIMEOUT_MS and capped at it, so in-flight requests get the rest
+ *   of the drain budget. SIGINT, WebSocket upgrades and SSE streams get no grace.
  * - SHUTDOWN_CLEANUP_TIMEOUT_MS: Total time allowed for post-drain cleanup
  */
 
@@ -106,7 +112,10 @@ import { removeStickyCookieFromPublicCacheableResponse } from "./response-header
 import {
   closeProxyServerWithin,
   createProxyDrainingResponse,
+  createProxyReadinessResponse,
   parseProxyDrainTimeoutMs,
+  PROXY_READY_PATH,
+  ProxyDrainAdmission,
   ProxyRequestDrainTracker,
 } from "./request-drain.ts";
 import {
@@ -213,6 +222,17 @@ const SHUTDOWN_DRAIN_TIMEOUT_MS = parseProxyDrainTimeoutMs(
   getEnv("SHUTDOWN_DRAIN_TIMEOUT_MS"),
   DEFAULT_SHUTDOWN_DRAIN_TIMEOUT_MS,
 );
+// New requests are still admitted for this long after SIGTERM (within the
+// drain budget) while the gateway moves its pooled connections off this pod.
+const DEFAULT_SHUTDOWN_ADMISSION_GRACE_MS = 10_000;
+const SHUTDOWN_ADMISSION_GRACE_MS = Math.min(
+  parseProxyDrainTimeoutMs(
+    getEnv("SHUTDOWN_ADMISSION_GRACE_MS"),
+    DEFAULT_SHUTDOWN_ADMISSION_GRACE_MS,
+    "SHUTDOWN_ADMISSION_GRACE_MS",
+  ),
+  SHUTDOWN_DRAIN_TIMEOUT_MS,
+);
 const SHUTDOWN_CLEANUP_TIMEOUT_MS = parseProxyShutdownCleanupTimeoutMs(
   getEnv("SHUTDOWN_CLEANUP_TIMEOUT_MS"),
   DEFAULT_PROXY_SHUTDOWN_CLEANUP_TIMEOUT_MS,
@@ -233,6 +253,7 @@ if (isProduction() && routingInvalidationSecretBytes < 32) {
 }
 const proxyRequestDrainTracker = new ProxyRequestDrainTracker();
 let shuttingDown = false;
+const drainAdmission = new ProxyDrainAdmission(SHUTDOWN_ADMISSION_GRACE_MS);
 
 const { createAuthProvider } = await importFirstPartyExtensionModule<AuthJwtExtensionModule>(
   "ext-auth-jwt",
@@ -851,8 +872,14 @@ async function router(req: Request): Promise<Response> {
   if (url.pathname === "/_proxy/health") {
     return Response.json({ service: "veryfront-proxy", status: "ok" });
   }
-  if (shuttingDown) return createProxyDrainingResponse();
+  if (url.pathname === PROXY_READY_PATH) return createProxyReadinessResponse(shuttingDown);
+  // Inside the admission grace window a draining replica still serves ordinary
+  // requests; a new WebSocket or SSE stream belongs on a replica that stays.
+  if (shuttingDown && !drainAdmission.admits(req, url.pathname)) {
+    return createProxyDrainingResponse();
+  }
 
+  const admittedDuringDrain = drainAdmission.isDraining();
   const requestId = crypto.randomUUID();
   proxyRequestDrainTracker.start(requestId, req.method, url.pathname);
 
@@ -879,21 +906,33 @@ async function router(req: Request): Promise<Response> {
       response = await forwardToServer(req, url);
     }
 
+    // Read after the await: shutdown may have started while this request ran.
+    response = drainAdmission.closeIfDraining(req, response, admittedDuringDrain);
     return proxyRequestDrainTracker.completeOnResponseEnd(requestId, req, response);
   } catch (error) {
     proxyRequestDrainTracker.complete(requestId);
-    if (error instanceof ProxyRequestHostError) {
-      // An absolute-form request target with an empty Host header parses as a
-      // URL, but Host validation rejects it downstream ("" is not nullish, so
-      // the url.host fallback never fires). Map it to 400 here instead of
-      // letting it escape as Deno's generic 500.
-      proxyLogger.warn(`400 ${req.method} ${url.pathname}`, {
-        host: describeRejectedHostHeader(req.headers.get("host")),
-      });
-      return jsonErrorResponse(400, { error: "Bad Request" });
-    }
-    throw error;
+    // Error responses close their connection while draining too, so the
+    // gateway does not stay pinned to this pod.
+    return drainAdmission.closeIfDraining(req, handleRouterError(req, url, error));
+  }
+}
+
+function handleRouterError(req: Request, url: URL, error: unknown): Response {
+  if (error instanceof ProxyRequestHostError) {
+    // An absolute-form request target with an empty Host header parses as a
+    // URL, but Host validation rejects it downstream ("" is not nullish, so
+    // the url.host fallback never fires). Map it to 400 here instead of
+    // letting it escape as Deno's generic 500.
+    proxyLogger.warn(`400 ${req.method} ${url.pathname}`, {
+      host: describeRejectedHostHeader(req.headers.get("host")),
+    });
+    return jsonErrorResponse(400, { error: "Bad Request" });
   }
+  // Outside a drain the server's own error handling owns unexpected failures.
+  if (!drainAdmission.isDraining()) throw error;
+  captureApplicationError(error, { boundary: "proxy.request", method: req.method });
+  proxyLogger.error(`500 ${req.method} ${url.pathname}`, {}, error);
+  return jsonErrorResponse(500, { error: "Internal Server Error" });
 }
 
 // Create server before signal registration so early SIGTERM/SIGINT can close it safely.
@@ -903,10 +942,12 @@ const server = createHttpServer();
 async function shutdown(signal: "SIGINT" | "SIGTERM"): Promise<void> {
   if (shuttingDown) return;
   shuttingDown = true;
+  drainAdmission.begin(signal);
 
   proxyLogger.info(`Received ${signal}, initiating graceful shutdown`, {
     inFlightRequests: proxyRequestDrainTracker.getInFlightCount(),
     drainTimeoutMs: SHUTDOWN_DRAIN_TIMEOUT_MS,
+    admissionGraceMs: drainAdmission.getGrantedMs(),
     cleanupTimeoutMs: SHUTDOWN_CLEANUP_TIMEOUT_MS,
   });
 
@@ -936,11 +977,17 @@ async function shutdown(signal: "SIGINT" | "SIGTERM"): Promise<void> {
   };
 
   try {
-    // New requests receive the draining response after shuttingDown is set.
-    // Keep this replica subscribed while already-started responses finish.
+    // New requests are admitted (closing their connections) until the grace
+    // window ends, then receive the draining response. Keep this replica
+    // subscribed while admitted and already-started responses finish.
+    // SIGINT gets no grace window. The wait uses a timer captured before
+    // extensions load, so a replaced setTimeout cannot skip cleanup.
+    await drainAdmission.waitForWindow();
     let drained = false;
     try {
-      drained = await proxyRequestDrainTracker.waitForDrain(SHUTDOWN_DRAIN_TIMEOUT_MS);
+      drained = await proxyRequestDrainTracker.waitForDrain(
+        drainAdmission.remainingDrainBudgetMs(SHUTDOWN_DRAIN_TIMEOUT_MS),
+      );
     } catch (error) {
       reportShutdownFailure({ step: "request_drain", error, timedOut: false });
     }