diff --git a/cli/commands/serve/split-mode.test.ts b/cli/commands/serve/split-mode.test.ts index d22f128a36..4015fa239f 100644 --- a/cli/commands/serve/split-mode.test.ts +++ b/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", () => { diff --git a/cli/commands/serve/split-mode.ts b/cli/commands/serve/split-mode.ts index 24cd29cae2..af34e76afd 100644 --- a/cli/commands/serve/split-mode.ts +++ b/cli/commands/serve/split-mode.ts @@ -42,6 +42,17 @@ const STATIC_ENV: Record = { 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, +): Record { + return { ...env, SHUTDOWN_ADMISSION_GRACE_MS: "0" }; +} + export function buildSplitModeEnvForTests( userEnv: Record, productionServerPort: number, @@ -179,7 +190,7 @@ export async function runSplitMode(options: SplitModeOptions): Promise { // Start proxy const proxyProcess = startProcess( [...veryfront, "serve", "--mode=proxy", `--port=${proxyPort}`], - env, + buildSplitModeProxyEnvForTests(env), `proxy on :${proxyPort}`, ); diff --git a/scripts/build/compile-binary.test.ts b/scripts/build/compile-binary.test.ts index a32a21f26e..2655a65f58 100644 --- a/scripts/build/compile-binary.test.ts +++ b/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( diff --git a/scripts/build/smoke-proxy-binary.sh b/scripts/build/smoke-proxy-binary.sh index 4a7b06af3a..c04f292470 100644 --- a/scripts/build/smoke-proxy-binary.sh +++ b/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=$! diff --git a/scripts/build/smoke-proxy-memory.sh b/scripts/build/smoke-proxy-memory.sh index 31fcf80f4d..893ccadf56 100755 --- a/scripts/build/smoke-proxy-memory.sh +++ b/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" diff --git a/src/agent/react/use-chat/use-chat.csrf.test.tsx b/src/agent/react/use-chat/use-chat.csrf.test.tsx index de9c0cc393..f6f8d0538c 100644 --- a/src/agent/react/use-chat/use-chat.csrf.test.tsx +++ b/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 { diff --git a/src/agent/react/use-chat/use-chat.ts b/src/agent/react/use-chat/use-chat.ts index 9dc78d42d0..ba81e5fcd6 100644 --- a/src/agent/react/use-chat/use-chat.ts +++ b/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, }, diff --git a/src/agent/react/use-completion.ts b/src/agent/react/use-completion.ts index bfd8f80254..9a08ee1b84 100644 --- a/src/agent/react/use-completion.ts +++ b/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, }, diff --git a/src/agent/react/use-streaming.ts b/src/agent/react/use-streaming.ts index a4f9bb949b..8e68385016 100644 --- a/src/agent/react/use-streaming.ts +++ b/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, diff --git a/src/proxy/main.test.ts b/src/proxy/main.test.ts index 220d43e778..14011bbc18 100644 --- a/src/proxy/main.test.ts +++ b/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"'); diff --git a/src/proxy/main.ts b/src/proxy/main.ts index d33f152354..e30e25159f 100644 --- a/src/proxy/main.ts +++ b/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( "ext-auth-jwt", @@ -851,8 +872,14 @@ async function router(req: Request): Promise { 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 = 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 { 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 { }; 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 }); } diff --git a/src/proxy/request-drain.test.ts b/src/proxy/request-drain.test.ts index 359957fbdb..8b98facc2a 100644 --- a/src/proxy/request-drain.test.ts +++ b/src/proxy/request-drain.test.ts @@ -9,8 +9,13 @@ import { describe, it } from "#veryfront/testing/bdd.ts"; import { closeProxyServerWithin, createProxyDrainingResponse, + createProxyReadinessResponse, + isLongLivedProxyRequest, parseProxyDrainTimeoutMs, + PROXY_READY_PATH, + ProxyDrainAdmission, ProxyRequestDrainTracker, + withConnectionClose, } from "./request-drain.ts"; describe("proxy request drain", () => { @@ -378,6 +383,230 @@ describe("proxy request drain", () => { assertEquals(response.headers.get("x-content-type-options"), "nosniff"); }); + it("reports ready until shutdown starts, then fails readiness with the draining response", async () => { + assertEquals(PROXY_READY_PATH, "/_proxy/ready"); + const ready = createProxyReadinessResponse(false); + assertEquals(ready.status, 200); + assertEquals(ready.headers.get("cache-control"), "no-store"); + assertEquals(await ready.json(), { service: "veryfront-proxy", status: "ready" }); + + const draining = createProxyReadinessResponse(true); + assertEquals(draining.status, 503); + assertEquals(draining.headers.get("retry-after"), "1"); + await draining.body?.cancel(); + }); + + it("admits requests for the grace window after SIGTERM, then refuses them", () => { + let now = 1_000; + const admission = new ProxyDrainAdmission(10_000, () => now); + const plain = new Request("http://proxy.test/page"); + assertEquals(admission.admits(plain, "/page"), true); + assertEquals(admission.getGrantedMs(), 0); + + admission.begin("SIGTERM"); + assertEquals(admission.admits(plain, "/page"), true); + assertEquals(admission.getGrantedMs(), 10_000); + + now += 4_000; + admission.begin("SIGTERM"); // a second signal does not extend the window + admission.begin("SIGINT"); // nor does a later SIGINT cut it short + assertEquals(admission.admits(plain, "/page"), true); + assertEquals(admission.getGrantedMs(), 10_000); + + now += 6_000; + assertEquals(admission.admits(plain, "/page"), false); + assertThrows(() => new ProxyDrainAdmission(-1), RangeError); + }); + + it("gives SIGINT no admission grace and no grace wait", async () => { + const admission = new ProxyDrainAdmission(60_000, () => 1_000); + admission.begin("SIGINT"); + assertEquals(admission.getGrantedMs(), 0); + assertEquals(admission.admits(new Request("http://proxy.test/page"), "/page"), false); + + // A 60s window would hold this wait open; SIGINT resolves it at once. + let timeoutId: ReturnType | undefined; + const settled = await Promise.race([ + admission.waitForWindow().then(() => "settled"), + new Promise((resolve) => { + timeoutId = setTimeout(() => resolve("pending"), 200); + }), + ]); + clearTimeout(timeoutId); + assertEquals(settled, "settled"); + }); + + it("refuses new WebSocket and SSE streams from the start of the grace window", () => { + const admission = new ProxyDrainAdmission(10_000, () => 0); + const runStream = new Request("http://proxy.test/api/control-plane/runs/run_1/stream", { + method: "POST", + }); + const eventSource = new Request("http://proxy.test/events", { + headers: { accept: "text/event-stream" }, + }); + const webSocket = new Request("http://proxy.test/ws", { + headers: { upgrade: "WebSocket" }, + }); + const runStatus = new Request("http://proxy.test/api/control-plane/runs/run_1"); + + // Before shutdown every request is served. + assertEquals(admission.admits(runStream, "/api/control-plane/runs/run_1/stream"), true); + + admission.begin("SIGTERM"); + assertEquals(admission.admits(runStream, "/api/control-plane/runs/run_1/stream"), false); + assertEquals(admission.admits(eventSource, "/events"), false); + assertEquals(admission.admits(webSocket, "/ws"), false); + // useCompletion/useStreaming prefer plain text but list SSE. + const textPreferred = new Request("http://proxy.test/api/completion", { + method: "POST", + headers: { accept: "text/plain, text/event-stream;q=0.9, */*;q=0.8" }, + }); + assertEquals(admission.admits(textPreferred, "/api/completion"), false); + assertEquals(admission.admits(runStatus, "/api/control-plane/runs/run_1"), true); + + // useChat POSTs the AG-UI stream without Accept: text/event-stream. + const post = (path: string) => + new Request(`http://proxy.test${path}`, { + method: "POST", + headers: { "content-type": "application/json" }, + }); + assertEquals(admission.admits(post("/api/ag-ui"), "/api/ag-ui"), false); + assertEquals(admission.admits(post("/api/ag-ui/"), "/api/ag-ui/"), false); + // useAgent's JSON invocation, agent metadata reads and other POSTs stay + // admitted (served with a connection-closing response). + assertEquals(admission.admits(post("/api/agents/support"), "/api/agents/support"), true); + assertEquals( + admission.admits(new Request("http://proxy.test/api/agents/support"), "/api/agents/support"), + true, + ); + assertEquals(admission.admits(post("/api/agents"), "/api/agents"), true); + assertEquals(admission.admits(post("/api/forms/contact"), "/api/forms/contact"), true); + + assertEquals(isLongLivedProxyRequest(runStream, "/api/control-plane/runs/run_1/stream"), true); + assertEquals( + isLongLivedProxyRequest( + new Request("http://proxy.test/api/control-plane/runs/run_1/stream"), + "/api/control-plane/runs/run_1/stream", + ), + false, + ); + }); + + it("closes connections only once shutdown has begun", async () => { + const admission = new ProxyDrainAdmission(10_000, () => 0); + const get = new Request("http://proxy.test/page"); + const before = admission.closeIfDraining(get, new Response("ok")); + assertEquals(before.headers.get("connection"), null); + + admission.begin("SIGTERM"); + const during = admission.closeIfDraining(get, new Response("ok")); + assertEquals(during.headers.get("connection"), "close"); + await before.body?.cancel(); + await during.body?.cancel(); + }); + + it("closes error responses produced while draining", async () => { + const admission = new ProxyDrainAdmission(10_000, () => 0); + admission.begin("SIGTERM"); + const error = admission.closeIfDraining( + new Request("http://proxy.test/page"), + new Response("Bad Request", { status: 400 }), + ); + assertEquals(error.status, 400); + assertEquals(error.headers.get("connection"), "close"); + await error.body?.cancel(); + }); + + it("answers a response-negotiated SSE stream admitted during the drain with the draining response", async () => { + const admission = new ProxyDrainAdmission(10_000, () => 0); + const sse = () => { + let cancelled = false; + const body = new ReadableStream({ + cancel() { + cancelled = true; + }, + }); + const response = new Response(body, { + headers: { "content-type": "text/event-stream; charset=utf-8" }, + }); + return { response, wasCancelled: () => cancelled }; + }; + const events = new Request("http://proxy.test/api/workflows/runs/run_1/events"); + + // Before shutdown an SSE response streams normally. + const live = sse(); + assertEquals(admission.closeIfDraining(events, live.response), live.response); + + admission.begin("SIGTERM"); + + // Admitted during the drain: the GET stream is retried elsewhere. + const fresh = sse(); + const refused = admission.closeIfDraining(events, fresh.response, true); + assertEquals(refused.status, 503); + assertEquals(refused.headers.get("connection"), "close"); + await Promise.resolve(); + assertEquals(fresh.wasCancelled(), true); + await refused.body?.cancel(); + + // Started before the drain: keep streaming within the drain budget. + const started = sse(); + const kept = admission.closeIfDraining(events, started.response, false); + assertEquals(kept.headers.get("content-type"), "text/event-stream; charset=utf-8"); + assertEquals(kept.headers.get("connection"), "close"); + await kept.body?.cancel(); + + // A POST already ran upstream, so it is not replayed by refusing it. + const posted = sse(); + const post = new Request("http://proxy.test/api/chat", { method: "POST" }); + const kept2 = admission.closeIfDraining(post, posted.response, true); + assertEquals(kept2.headers.get("content-type"), "text/event-stream; charset=utf-8"); + await kept2.body?.cancel(); + }); + + it("measures the remaining drain budget from the start of shutdown", () => { + let now = 1_000; + const sigterm = new ProxyDrainAdmission(10_000, () => now); + sigterm.begin("SIGTERM"); + now += 10_000; + assertEquals(sigterm.remainingDrainBudgetMs(25_000), 15_000); + // A grace wait that settled late (busy event loop) shrinks the budget + // instead of extending the shutdown past the configured bound. + now += 2_000.5; + assertEquals(sigterm.remainingDrainBudgetMs(25_000), 12_999); + now += 20_000; + assertEquals(sigterm.remainingDrainBudgetMs(25_000), 0); + + const sigint = new ProxyDrainAdmission(10_000, () => 0); + sigint.begin("SIGINT"); + assertEquals(sigint.remainingDrainBudgetMs(25_000), 25_000); + }); + + it("marks admitted drain responses to close their connection", async () => { + const mutable = withConnectionClose(new Response("ok", { headers: { "x-a": "1" } })); + assertEquals(mutable.headers.get("connection"), "close"); + assertEquals(mutable.headers.get("x-a"), "1"); + + const upstream = Response.redirect("https://example.com/next", 307); + const rebuilt = withConnectionClose(upstream); + assertEquals(rebuilt.status, 307); + assertEquals(rebuilt.headers.get("location"), "https://example.com/next"); + assertEquals(rebuilt.headers.get("connection"), "close"); + await rebuilt.body?.cancel(); + await mutable.body?.cancel(); + }); + + it("parses the admission grace with its own variable name", () => { + assertEquals( + parseProxyDrainTimeoutMs(undefined, 10_000, "SHUTDOWN_ADMISSION_GRACE_MS"), + 10_000, + ); + assertThrows( + () => parseProxyDrainTimeoutMs("x", 10_000, "SHUTDOWN_ADMISSION_GRACE_MS"), + TypeError, + "SHUTDOWN_ADMISSION_GRACE_MS", + ); + }); + it("bounds server close when an adapter keeps waiting on open connections", async () => { assertEquals(await closeProxyServerWithin(() => new Promise(() => {}), 5), false); assertEquals(await closeProxyServerWithin(() => Promise.resolve(), 50), true); diff --git a/src/proxy/request-drain.ts b/src/proxy/request-drain.ts index 0ce99708e7..3ea6c03c5c 100644 --- a/src/proxy/request-drain.ts +++ b/src/proxy/request-drain.ts @@ -1,8 +1,26 @@ import { getRequestTransportLifetime } from "#veryfront/platform/adapters/runtime/shared/request-peer.ts"; import { completeOnResponseBodyConsumption } from "#veryfront/platform/compat/http/response-lifecycle.ts"; -import { continueProxyShutdownPromise } from "./shutdown-intrinsics.ts"; +import { continueProxyShutdownPromise, delayProxyShutdown } from "./shutdown-intrinsics.ts"; +import { normalizeProxyRoutePolicyPathname } from "./legacy-managed-agent-route-denial.ts"; +import { isControlPlaneRunStreamPost } from "./retry.ts"; import { MAX_PROXY_TIMER_DELAY_MS } from "./timing.ts"; +// Captured before extensions activate: shutdown timing and drain reporting must +// not consult a replaced Performance.prototype.now, Math, Number or Array. +const apply = Reflect.apply; +const mathCeil = Math.ceil; +const mathFloor = Math.floor; +const mathMax = Math.max; +const mathMin = Math.min; +const isSafeInteger = Number.isSafeInteger; +const arrayFrom = Array.from; +const mapValues = Map.prototype.values; +const monotonicClock = performance; +const monotonicNowMethod = performance.now; +function monotonicNow(): number { + return apply(monotonicNowMethod, monotonicClock, []); +} + export interface TrackedProxyRequest { requestId: string; method: string; @@ -18,7 +36,7 @@ export class ProxyRequestDrainTracker { requestId, method, path, - startTime: performance.now(), + startTime: monotonicNow(), }); } @@ -53,12 +71,15 @@ export class ProxyRequestDrainTracker { } getInFlightRequests(): TrackedProxyRequest[] { - return Array.from(this.inFlight.values(), (request) => ({ ...request })); + return arrayFrom( + apply(mapValues, this.inFlight, []) as MapIterator, + (request) => ({ ...request }), + ); } async waitForDrain(timeoutMs: number, pollIntervalMs = 100): Promise { if ( - !Number.isSafeInteger(timeoutMs) || + !isSafeInteger(timeoutMs) || timeoutMs < 0 || timeoutMs > MAX_PROXY_TIMER_DELAY_MS ) { @@ -67,7 +88,7 @@ export class ProxyRequestDrainTracker { ); } if ( - !Number.isSafeInteger(pollIntervalMs) || + !isSafeInteger(pollIntervalMs) || pollIntervalMs < 1 || pollIntervalMs > MAX_PROXY_TIMER_DELAY_MS ) { @@ -77,15 +98,13 @@ export class ProxyRequestDrainTracker { } if (this.inFlight.size === 0) return true; - const deadline = performance.now() + timeoutMs; + const deadline = monotonicNow() + timeoutMs; while (this.inFlight.size > 0) { - const remainingMs = deadline - performance.now(); + const remainingMs = deadline - monotonicNow(); if (remainingMs <= 0) return false; - await new Promise((resolve) => - setTimeout(resolve, Math.min(pollIntervalMs, remainingMs)) - ); + await delayProxyShutdown(mathMin(pollIntervalMs, remainingMs)); } return true; @@ -95,6 +114,7 @@ export class ProxyRequestDrainTracker { export function parseProxyDrainTimeoutMs( rawValue: string | undefined, defaultValue: number, + name = "SHUTDOWN_DRAIN_TIMEOUT_MS", ): number { if ( !Number.isSafeInteger(defaultValue) || @@ -107,7 +127,7 @@ export function parseProxyDrainTimeoutMs( } if (rawValue === undefined || rawValue === "") return defaultValue; if (!/^\d+$/.test(rawValue)) { - throw new TypeError("SHUTDOWN_DRAIN_TIMEOUT_MS must be a non-negative decimal integer"); + throw new TypeError(`${name} must be a non-negative decimal integer`); } const parsed = Number(rawValue); if ( @@ -115,9 +135,7 @@ export function parseProxyDrainTimeoutMs( parsed < 0 || parsed > 10 * 60 * 1_000 ) { - throw new RangeError( - "SHUTDOWN_DRAIN_TIMEOUT_MS must be between 0 and 600000", - ); + throw new RangeError(`${name} must be between 0 and 600000`); } return parsed; } @@ -134,6 +152,160 @@ export function createProxyDrainingResponse(): Response { }); } +/** + * After SIGTERM the proxy keeps admitting new requests for a short grace + * window. The gateway reaches it over pooled keep-alive connections through the + * Service, and requests already queued on those connections would otherwise + * get the draining 503. Each admitted response closes its connection (see + * `withConnectionClose`), so the gateway reconnects through the Service, which + * no longer routes to this terminating pod. After the window, new requests get + * the draining response. + * + * SIGINT (a local Ctrl-C) has no gateway to wait for, so it starts with no + * grace window. Long-lived requests (WebSocket upgrades and SSE streams) are + * refused from the moment shutdown starts: they could outlive the remaining + * drain budget and belong on a replica that stays. An idempotent request + * admitted during the drain whose response turns out to be an SSE stream + * (negotiated by the response, not the Accept header) is answered with the + * draining response instead, so the client retries on a replica that stays. + */ +export class ProxyDrainAdmission { + private admitUntil: number | null = null; + private startedAt = 0; + private grantedMs = 0; + + constructor( + private readonly windowMs: number, + private readonly now: () => number = () => monotonicNow(), + ) { + if (!Number.isSafeInteger(windowMs) || windowMs < 0) { + throw new RangeError("Proxy drain admission window must be a non-negative integer"); + } + } + + begin(signal: "SIGINT" | "SIGTERM"): void { + if (this.admitUntil !== null) return; + this.grantedMs = signal === "SIGTERM" ? this.windowMs : 0; + this.startedAt = this.now(); + this.admitUntil = this.startedAt + this.grantedMs; + } + + /** Whether a request arriving now is served (always true before `begin`). */ + admits(request: Request, pathname: string): boolean { + if (this.admitUntil === null) return true; + if (isLongLivedProxyRequest(request, pathname)) return false; + return this.now() < this.admitUntil; + } + + /** Whether shutdown has begun. */ + isDraining(): boolean { + return this.admitUntil !== null; + } + + /** + * Once shutdown has begun, close the connection after this response. A GET + * or HEAD admitted during the drain that answers with an SSE stream gets the + * draining response instead: the stream could outlive the drain budget, and + * retrying it on another replica is safe. + */ + closeIfDraining(request: Request, response: Response, admittedDuringDrain = false): Response { + if (this.admitUntil === null) return response; + if (admittedDuringDrain && isIdempotentStreamRetry(request, response)) { + cancelResponseBody(response); + return createProxyDrainingResponse(); + } + return withConnectionClose(response); + } + + /** + * What remains of `drainTimeoutMs` measured from `begin`, so a grace wait + * that settled late (a busy event loop) does not extend the shutdown bound. + */ + remainingDrainBudgetMs(drainTimeoutMs: number): number { + if (this.admitUntil === null) return drainTimeoutMs; + return mathMax(0, mathFloor(drainTimeoutMs - (this.now() - this.startedAt))); + } + + /** The whole grace window granted by `begin`, in integer milliseconds. */ + getGrantedMs(): number { + return this.grantedMs; + } + + /** Wait out the rest of the grace window on a timer captured before extensions load. */ + waitForWindow(): Promise { + const remainingMs = this.admitUntil === null ? 0 : mathMax(0, this.admitUntil - this.now()); + return delayProxyShutdown(mathCeil(remainingMs)); + } +} + +function isIdempotentStreamRetry(request: Request, response: Response): boolean { + const method = request.method.toUpperCase(); + if (method !== "GET" && method !== "HEAD") return false; + return response.headers.get("content-type")?.toLowerCase().startsWith("text/event-stream") ?? + false; +} + +function cancelResponseBody(response: Response): void { + try { + response.body?.cancel().catch(() => { + // The replaced stream has no consumer left to report to. + }); + } catch { + // A locked body is already owned elsewhere; nothing to release here. + } +} + +/** + * The framework AG-UI endpoint answers with an SSE stream although its + * built-in `useChat` client sends no `Accept: text/event-stream`. (`useAgent` + * posts to `/api/agents/:id` and reads JSON, so that route stays admitted.) + */ +function isKnownAgentStreamPost(request: Request, pathname: string): boolean { + return request.method.toUpperCase() === "POST" && + normalizeProxyRoutePolicyPathname(pathname) === "/api/ag-ui"; +} + +/** WebSocket upgrades and SSE streams hold their connection open indefinitely. */ +export function isLongLivedProxyRequest(request: Request, pathname: string): boolean { + if (request.headers.get("upgrade")?.toLowerCase() === "websocket") return true; + if (request.headers.get("accept")?.toLowerCase().includes("text/event-stream")) return true; + return isControlPlaneRunStreamPost(request, pathname) || + isKnownAgentStreamPost(request, pathname); +} + +/** Ask the HTTP/1.1 server to close the connection after this response. */ +export function withConnectionClose(response: Response): Response { + if (response.status === 101) return response; + try { + response.headers.set("Connection", "close"); + return response; + } catch { + // Fetched upstream responses carry immutable headers. + const headers = new Headers(response.headers); + headers.set("Connection", "close"); + return new Response(response.body, { + status: response.status, + statusText: response.statusText, + headers, + }); + } +} + +/** + * Readiness path for the proxy. Unlike the static `/_proxy/health` liveness + * check, it fails as soon as shutdown starts, so Kubernetes removes the + * endpoint instead of routing new requests to a replica that refuses them. + */ +export const PROXY_READY_PATH = "/_proxy/ready"; + +export function createProxyReadinessResponse(shuttingDown: boolean): Response { + if (shuttingDown) return createProxyDrainingResponse(); + return Response.json( + { service: "veryfront-proxy", status: "ready" }, + { headers: { "Cache-Control": "no-store" } }, + ); +} + export async function closeProxyServerWithin( close: () => Promise, timeoutMs: number, diff --git a/src/proxy/retry.ts b/src/proxy/retry.ts index b56a7533d9..7e6109db1b 100644 --- a/src/proxy/retry.ts +++ b/src/proxy/retry.ts @@ -101,7 +101,7 @@ function isIdempotentMethod(method: string): boolean { return normalized === "GET" || normalized === "HEAD" || normalized === "OPTIONS"; } -function isControlPlaneRunStreamPost(request: Request, pathname: string): boolean { +export function isControlPlaneRunStreamPost(request: Request, pathname: string): boolean { return request.method.toUpperCase() === "POST" && CONTROL_PLANE_RUN_STREAM_PATH.test(pathname); } diff --git a/src/proxy/shutdown-intrinsics.ts b/src/proxy/shutdown-intrinsics.ts index f08003dd26..0a69fe26d8 100644 --- a/src/proxy/shutdown-intrinsics.ts +++ b/src/proxy/shutdown-intrinsics.ts @@ -17,6 +17,7 @@ const getOwnPropertyDescriptor = Object.getOwnPropertyDescriptor; const hasOwnProperty = Object.prototype.hasOwnProperty; const nativePromiseThen = Promise.prototype.then; const promiseSpecies = Symbol.species; +const setTimer = setTimeout; function hasOwn(object: PropertyDescriptor, key: PropertyKey): boolean { return apply(hasOwnProperty, object, [key]) as boolean; @@ -121,3 +122,13 @@ export function resolveProxyShutdownValue( resolve(value); }); } + +/** + * Resolve after `delayMs` using the timer captured at module evaluation, so a + * shutdown wait still settles after extension code replaces `setTimeout`. + */ +export function delayProxyShutdown(delayMs: number): Promise { + return createProxyShutdownPromise((resolve) => { + setTimer(resolve, delayMs); + }); +} diff --git a/src/proxy/shutdown-lifecycle.test.ts b/src/proxy/shutdown-lifecycle.test.ts index c4e1a7f818..37e9d2b489 100644 --- a/src/proxy/shutdown-lifecycle.test.ts +++ b/src/proxy/shutdown-lifecycle.test.ts @@ -1,6 +1,7 @@ import "#veryfront/schemas/_test-setup.ts"; import { assertEquals, assertRejects, assertThrows } from "#veryfront/testing/assert.ts"; import { describe, it } from "#veryfront/testing/bdd.ts"; +import { ProxyDrainAdmission, ProxyRequestDrainTracker } from "./request-drain.ts"; import { parseProxyShutdownCleanupTimeoutMs, runProxyShutdownSteps } from "./shutdown-lifecycle.ts"; describe("proxy shutdown lifecycle", () => { @@ -183,6 +184,117 @@ describe("proxy shutdown lifecycle", () => { assertEquals(await settlement, []); }); + it("waits out the admission grace window on the timer captured before extensions load", async () => { + const admission = new ProxyDrainAdmission(5, () => 0); + admission.begin("SIGTERM"); + const timerDescriptor = Object.getOwnPropertyDescriptor(globalThis, "setTimeout")!; + let wait: Promise | undefined; + try { + Object.defineProperty(globalThis, "setTimeout", { + configurable: timerDescriptor.configurable, + enumerable: timerDescriptor.enumerable, + writable: true, + value: () => { + throw new Error("poisoned setTimeout"); + }, + }); + wait = admission.waitForWindow(); + } finally { + Object.defineProperty(globalThis, "setTimeout", timerDescriptor); + } + + assertEquals(await wait, undefined); + }); + + it("starts and waits out the admission grace on the clock captured before extensions load", async () => { + const admission = new ProxyDrainAdmission(1); + // Deno defines `now` on Performance.prototype; Bun on the instance. + const nowOwner: object = Object.hasOwn(performance, "now") + ? performance + : Performance.prototype; + const nowDescriptor = Object.getOwnPropertyDescriptor(nowOwner, "now")!; + let wait: Promise | undefined; + let admitsStream: boolean | undefined; + try { + Object.defineProperty(nowOwner, "now", { + ...nowDescriptor, + value: () => { + throw new Error("poisoned performance.now"); + }, + }); + admission.begin("SIGTERM"); + admitsStream = admission.admits( + new Request("http://proxy.test/events", { headers: { accept: "text/event-stream" } }), + "/events", + ); + wait = admission.waitForWindow(); + } finally { + Object.defineProperty(nowOwner, "now", nowDescriptor); + } + + assertEquals(admitsStream, false); + assertEquals(await wait, undefined); + assertEquals(admission.getGrantedMs(), 1); + }); + + it("computes the grace wait and drain budget with math captured before extensions load", async () => { + const admission = new ProxyDrainAdmission(1, () => 0); + admission.begin("SIGTERM"); + const maxDescriptor = Object.getOwnPropertyDescriptor(Math, "max")!; + const ceilDescriptor = Object.getOwnPropertyDescriptor(Math, "ceil")!; + const floorDescriptor = Object.getOwnPropertyDescriptor(Math, "floor")!; + const poisoned = () => { + throw new Error("poisoned Math"); + }; + let wait: Promise | undefined; + let budget: number | undefined; + try { + Object.defineProperty(Math, "max", { ...maxDescriptor, value: poisoned }); + Object.defineProperty(Math, "ceil", { ...ceilDescriptor, value: poisoned }); + Object.defineProperty(Math, "floor", { ...floorDescriptor, value: poisoned }); + wait = admission.waitForWindow(); + budget = admission.remainingDrainBudgetMs(25); + } finally { + Object.defineProperty(Math, "max", maxDescriptor); + Object.defineProperty(Math, "ceil", ceilDescriptor); + Object.defineProperty(Math, "floor", floorDescriptor); + } + + assertEquals(await wait, undefined); + assertEquals(budget, 25); + }); + + it("polls and reports the request drain with intrinsics captured before extensions load", async () => { + const tracker = new ProxyRequestDrainTracker(); + tracker.start("req_1", "GET", "/slow"); + const poisoned = () => { + throw new Error("poisoned intrinsic"); + }; + const targets = [ + [Math, "min"], + [Number, "isSafeInteger"], + [Array, "from"], + ] as const; + const descriptors = targets.map(([owner, key]) => Object.getOwnPropertyDescriptor(owner, key)!); + let drained: Promise | undefined; + let reported: string[] | undefined; + try { + targets.forEach(([owner, key], index) => + Object.defineProperty(owner, key, { ...descriptors[index], value: poisoned }) + ); + drained = tracker.waitForDrain(5, 1); + reported = tracker.getInFlightRequests().map((request) => request.path); + } finally { + targets.forEach(([owner, key], index) => + Object.defineProperty(owner, key, descriptors[index]!) + ); + } + + // The poll keeps waiting (and then times out) instead of rejecting. + assertEquals(await drained, false); + assertEquals(reported, ["/slow"]); + }); + it("rejects malformed shutdown step configurations", async () => { const ran: string[] = []; const record = (name: string) => ({ diff --git a/src/react/components/chat/hooks-coverage.test.tsx b/src/react/components/chat/hooks-coverage.test.tsx index 25276cff89..9ad1f8cf79 100644 --- a/src/react/components/chat/hooks-coverage.test.tsx +++ b/src/react/components/chat/hooks-coverage.test.tsx @@ -305,9 +305,11 @@ describe("veryfront/chat hook behavior", () => { it("useCompletion streams text and reports the completed value", async () => { let requestBody: unknown; + let accept: string | null = null; let finished = ""; const restoreFetch = replaceFetch(async (_input, init) => { requestBody = JSON.parse(String(observeFetchRequestInit(init).body)); + accept = new Headers(observeFetchRequestInit(init).headers).get("accept"); return chunkedResponse(["Generated ", "answer"]); }); try { @@ -331,6 +333,8 @@ describe("veryfront/chat hook behavior", () => { prompt: "Write an answer", model: "test-model", }); + // Declaring the stream lets a draining proxy refuse it up front. + assertEquals(accept, "text/plain, text/event-stream;q=0.9, */*;q=0.8"); }, ); } finally { @@ -341,7 +345,11 @@ describe("veryfront/chat hook behavior", () => { it("useStreaming accumulates chunks and reports completion", async () => { const chunks: string[] = []; let completions = 0; - const restoreFetch = replaceFetch(async () => chunkedResponse(["streamed ", "value"])); + let accept: string | null = null; + const restoreFetch = replaceFetch(async (_input, init) => { + accept = new Headers(observeFetchRequestInit(init).headers).get("accept"); + return chunkedResponse(["streamed ", "value"]); + }); try { await exerciseAsync( () => @@ -364,6 +372,7 @@ describe("veryfront/chat hook behavior", () => { "each stream chunk is delivered separately to onChunk", ); assertEquals(completions, 1); + assertEquals(accept, "text/plain, text/event-stream;q=0.9, */*;q=0.8"); const stream = current(); flushSync(() => stream.reset()); assertEquals(current().data, "");