Hub shutdown: drain in-flight requests then force-close lingering streams

apps/hub/README.md

@@ -27,6 +27,12 @@ native tenant middleware.
 - `credential-expiry-sweep.ts`, `cron-due.ts`, `routine-launcher.ts`,
   `routine-scheduler.ts`, `tenant-create-guard.ts` are host-level
   background jobs and guards wired at boot, alongside the route mounts.
+- `in-flight-requests.ts` and `shutdown.ts` bound SIGINT/SIGTERM: the hub
+  waits for Hono handlers that have not yet returned a Response (a
+  request still in a Postgres transaction, a git write), then
+  `server.stop(true)` so a live SSE bridge or sidecar websocket cannot
+  hang the drain. The sequence is capped at 10s; a lingering stream is
+  not a shutdown fault.
 
 ## Running
 

apps/hub/src/in-flight-requests.test.ts

@@ -0,0 +1,59 @@
+import { expect, test } from "bun:test";
+import { Hono } from "hono";
+
+import { createInFlightRequestTracker } from "./in-flight-requests";
+
+test("pending starts at zero and whenIdle resolves immediately", async () => {
+  const tracker = createInFlightRequestTracker();
+  expect(tracker.pending).toBe(0);
+  await tracker.whenIdle();
+  expect(tracker.pending).toBe(0);
+});
+
+test("middleware counts a request until the handler returns", async () => {
+  const tracker = createInFlightRequestTracker();
+  const app = new Hono();
+  app.use(tracker.middleware);
+  let pendingDuringHandler: number | undefined;
+  app.get("/work", async () => {
+    pendingDuringHandler = tracker.pending;
+    return new Response("ok");
+  });
+
+  const response = await app.request("/work");
+  expect(response.status).toBe(200);
+  expect(pendingDuringHandler).toBe(1);
+  expect(tracker.pending).toBe(0);
+});
+
+test("whenIdle waits for an in-flight handler, including one that throws", async () => {
+  const tracker = createInFlightRequestTracker();
+  const app = new Hono();
+  app.use(tracker.middleware);
+  app.onError(() => new Response("error", { status: 500 }));
+
+  let release!: () => void;
+  const held = new Promise<void>((resolve) => {
+    release = resolve;
+  });
+  app.get("/held", async () => {
+    await held;
+    throw new Error("handler fault");
+  });
+
+  const request = app.request("/held");
+  const idle = tracker.whenIdle();
+  let idleSettled = false;
+  void idle.then(() => {
+    idleSettled = true;
+  });
+  await new Promise((resolve) => setTimeout(resolve, 20));
+  expect(tracker.pending).toBe(1);
+  expect(idleSettled).toBe(false);
+
+  release();
+  await request;
+  await idle;
+  expect(idleSettled).toBe(true);
+  expect(tracker.pending).toBe(0);
+});

apps/hub/src/in-flight-requests.ts

@@ -0,0 +1,57 @@
+import { Hono, type MiddlewareHandler } from "hono";
+
+export type InFlightRequestTracker = {
+  readonly middleware: MiddlewareHandler;
+  readonly pending: number;
+  whenIdle: () => Promise<void>;
+};
+
+/**
+ * Counts Hono handlers that have not yet returned a Response. Streaming
+ * bodies (SSE) and upgraded websockets stay open after that return, so they
+ * do not keep `pending` above zero — those are lingering connections the
+ * drain force-closes, not in-flight work.
+ */
+export function createInFlightRequestTracker(): InFlightRequestTracker {
+  let pending = 0;
+  const waiters = new Set<() => void>();
+
+  const notifyIfIdle = () => {
+    if (pending !== 0) return;
+    for (const waiter of waiters) waiter();
+    waiters.clear();
+  };
+
+  const middleware: MiddlewareHandler = async (_c, next) => {
+    pending += 1;
+    try {
+      await next();
+    } finally {
+      pending -= 1;
+      notifyIfIdle();
+    }
+  };
+
+  return {
+    middleware,
+    get pending() {
+      return pending;
+    },
+    whenIdle() {
+      if (pending === 0) return Promise.resolve();
+      return new Promise<void>((resolve) => {
+        waiters.add(resolve);
+      });
+    },
+  };
+}
+
+export function withInFlightRequestTracking<E extends object>(
+  app: Hono<E>,
+  tracker: InFlightRequestTracker,
+): Hono<E> {
+  const outer = new Hono<E>();
+  outer.use(tracker.middleware);
+  outer.route("/", app);
+  return outer;
+}

apps/hub/src/index.ts

@@ -367,7 +367,11 @@ import { createToolGrantsForPins } from "./tool-grants";
 import { createMcpCredentialBindingsFor } from "./mcp-credential-bindings";
 import { reconcilePinnedToolPackagesAfterConnect } from "./connection-live-reconcile";
 import { createPinnedPackageCredentialBindingsFor } from "./pinned-package-credential-bindings";
-import { shutdownHub } from "./shutdown";
+import { drainHubServer, shutdownHub } from "./shutdown";
+import {
+  createInFlightRequestTracker,
+  withInFlightRequestTracking,
+} from "./in-flight-requests";
 
 // Host policy constants, not configuration.
 const MAX_TARBALL_BYTES = 10 * 1024 * 1024;
@@ -3492,6 +3496,8 @@ export async function createHub(config: HubConfig) {
     guardDeps.operatorTenantId = config.operatorTenantId;
   }
   const guardedApp = guardedHubApp(app, guardDeps);
+  const inFlight = createInFlightRequestTracker();
+  const servingApp = withInFlightRequestTracking(guardedApp, inFlight);
 
   // Env-key auto-plant (CL-6101): runs in-process against the app this
   // function is about to return, so it needs nothing more than that
@@ -3502,11 +3508,12 @@ export async function createHub(config: HubConfig) {
     envProviderKeys: config.envProviderKeys,
     envProviderBaseUrls: config.envProviderBaseUrls,
     admin: config.envCredentialPlantAdmin,
-    fetch: (request) => Promise.resolve(guardedApp.fetch(request)),
+    fetch: (request) => Promise.resolve(servingApp.fetch(request)),
   });
 
   return {
-    app: guardedApp,
+    app: servingApp,
+    whenRequestsIdle: () => inFlight.whenIdle(),
     db,
     close: async () => {
       sidecarAllocationReconciliationStopped = true;
@@ -3552,14 +3559,20 @@ if (import.meta.main) {
   const log = getLogger(["hub"]);
   log.info`Hub serving on port ${port}`;
   const SHUTDOWN_DRAIN_MS = 10_000;
-  // `server.stop()` waits for open connections and websockets by default,
-  // so it sits inside the same bound as the hub's own closes.
+  // In-flight Hono handlers (a request mid-Postgres-transaction, a git
+  // write, anything that has not returned a Response yet) must finish
+  // before connections are torn down. `server.stop()` with no argument
+  // also waits for SSE bridges and idle sidecar websockets, which never
+  // close on their own — so once handlers are idle, force-close what's
+  // left. A live stream must not turn this drain into a timeout fault.
   const shutdown = () =>
     shutdownHub({
-      drain: async () => {
-        await server.stop();
-        await hub.close();
-      },
+      drain: () =>
+        drainHubServer({
+          whenRequestsIdle: hub.whenRequestsIdle,
+          stop: (force) => server.stop(force),
+          close: hub.close,
+        }),
       timeoutMs: SHUTDOWN_DRAIN_MS,
       exit: (code) => process.exit(code),
     });

apps/hub/src/shutdown.test.ts

@@ -1,5 +1,5 @@
 import { expect, test } from "bun:test";
-import { drainWithTimeout, shutdownHub } from "./shutdown";
+import { drainHubServer, drainWithTimeout, shutdownHub } from "./shutdown";
 
 test("drainWithTimeout resolves drained when the drain completes inside the bound", async () => {
   const outcome = await drainWithTimeout(() => Promise.resolve(), 1_000);
@@ -76,3 +76,19 @@ test("shutdownHub exits non-zero and reports the cause when the drain rejects",
   expect(exitCode).toBe(1);
   expect(reported).toEqual([error]);
 });
+
+test("drainHubServer waits for idle handlers then force-stops", async () => {
+  const order: string[] = [];
+  await drainHubServer({
+    whenRequestsIdle: async () => {
+      order.push("idle");
+    },
+    stop: (force) => {
+      order.push(force ? "stop-force" : "stop");
+    },
+    close: async () => {
+      order.push("close");
+    },
+  });
+  expect(order).toEqual(["idle", "stop-force", "close"]);
+});

apps/hub/src/shutdown.ts

@@ -39,6 +39,27 @@ export async function drainWithTimeout(
   return outcome;
 }
 
+export type DrainHubServerArgs = {
+  whenRequestsIdle: () => Promise<void>;
+  stop: (force: boolean) => void | Promise<void>;
+  close: () => void | Promise<void>;
+};
+
+/**
+ * Wait for every in-flight Hono handler to return, then force-stop the
+ * listener so lingering SSE/websocket connections cannot hang `server.stop()`,
+ * then close hub resources.
+ */
+export async function drainHubServer({
+  whenRequestsIdle,
+  stop,
+  close,
+}: DrainHubServerArgs): Promise<void> {
+  await whenRequestsIdle();
+  await stop(true);
+  await close();
+}
+
 export type ShutdownHubDeps = {
   drain: () => Promise<void>;
   timeoutMs: number;