feat(backend): implement graceful shutdown with drain, flush, and connection teardown (#349)

backend/src/api/app.ts

@@ -97,6 +97,14 @@ export interface AppOptions {
   authService?: AuthService;
   /** Enable background queue worker (default: true) */
   enableQueueWorker?: boolean;
+  /**
+   * How long close() waits for in-flight jobs to finish before closing the
+   * HTTP/WS server anyway. Default: 10000 (10s). A job still running when
+   * this elapses is left in the queue's "active" state — the next worker
+   * start (see JobWorker.start()/recoverIncompleteJobs()) resets it to
+   * "pending" and retries it, rather than losing the work.
+   */
+  jobWorkerStopTimeoutMs?: number;
 }
 
 function tryLoadStellarRelease(): StellarReleasePaymentFn | undefined {
@@ -252,15 +260,21 @@ export function createApp(opts: AppOptions = {}): {
   }));
 
   function close(callback?: () => void): void {
-    jobWorker.stop();
-    heartbeatService.stop();
-    metricsService.setWebSocketProbe(null);
-    detachStream();
-    if (httpServer.listening) {
-      httpServer.close(callback);
-    } else if (callback) {
-      callback();
-    }
+    // Drain first: wait for in-flight jobs to finish (bounded by
+    // jobWorkerStopTimeoutMs) before we stop accepting connections. A job
+    // still active when the drain window elapses is NOT force-failed — it
+    // stays "active" in the store and is picked back up by the next
+    // JobWorker.start() via recoverIncompleteJobs().
+    jobWorker.stop(opts.jobWorkerStopTimeoutMs ?? 10_000).finally(() => {
+      heartbeatService.stop();
+      metricsService.setWebSocketProbe(null);
+      detachStream();
+      if (httpServer.listening) {
+        httpServer.close(callback);
+      } else if (callback) {
+        callback();
+      }
+    });
   }
 
   const routeCount = (app as unknown as { _router?: { stack?: unknown[] } })._router?.stack?.length;

backend/src/api/routes/agents.ts

@@ -21,6 +21,15 @@ export interface AgentsRouterOptions {
 
 const DEFAULT_HEALTH_TIMEOUT_MS = 3_000;
 
+// Mirrors the RegisterAgentRequest schema documented in api/docs.ts.
+const RegisterAgentSchema = z.object({
+  agentId: z.string().min(1),
+  capabilities: z.array(z.string()).min(1),
+  pricingXLM: z.number().min(0.001),
+  endpoint: z.string().url(),
+  stellarPublicKey: z.string().regex(/^G[A-Z2-7]{55}$/),
+});
+
 export function createAgentsRouter(options: AgentsRouterOptions = {}): Router {
   const router = Router();
   const config = getConfig();

backend/src/api/routes/stream.ts

@@ -15,6 +15,22 @@ const STREAM_PATH = /^\/tasks\/([^/?]+)\/stream(?:\?.*)?$/;
 
 const logger = createLogger({ module: 'ws-stream' });
 
+/**
+ * Every WebSocketServer created by attachTaskStream(), tracked so
+ * getStreamConnectionCount() can report a live total across all of them
+ * (normally just one, per HTTP server) for the health/metrics dashboard.
+ */
+const activeStreamServers = new Set<WebSocketServer>();
+
+/** Total connected WebSocket clients across every attached stream server. */
+export function getStreamConnectionCount(): number {
+  let count = 0;
+  for (const wss of activeStreamServers) {
+    count += wss.clients.size;
+  }
+  return count;
+}
+
 // ---------------------------------------------------------------------------
 // Wire-format normalisation
 // ---------------------------------------------------------------------------

backend/src/index.ts

@@ -9,11 +9,11 @@ import { initializeAgents, globalAgentRegistry } from "./agents";
 import { startAgentSync, stopAgentSync } from "./registry/sync";
 import { loadConfig } from "./config";
 import { AgentCleanupService } from "./services/agentCleanup";
-import { createTaskDb, getTaskDb, closeTaskDb } from "./db/tasks";
 import { createAgentDb, getAgentDb, closeAgentDb } from "./db/agents";
 import { closeDb } from "./db/index";
 import { closeAuthDb } from "./db/auth";
 import { closeJobDb } from "./queue";
+import { eventBus } from "./coordinator/eventBus";
 import { createDefaultReconciliationService } from "./services/reconciliation";
 import { createLogger } from "./utils/logger";
 import { redactedConfigSnapshot } from "./config";
@@ -58,7 +58,9 @@ async function main() {
     errorRegistryMaintenance.start();
 
     // Create and start the server
-    const { httpServer, close } = createApp();
+    const { httpServer, close } = createApp({
+      jobWorkerStopTimeoutMs: config.GRACEFUL_SHUTDOWN_TIMEOUT * 1000,
+    });
 
     const port = config.PORT;
 
@@ -97,10 +99,30 @@ async function main() {
   }
 }
 
+export interface GracefulShutdownExtras {
+  cleanupService?: { stop(): void };
+  reconciliationService?: { stop(): void };
+  globalAgentRegistry?: { shutdown(): void };
+}
+
+/**
+ * SIGTERM/SIGINT handler: stop accepting new work, drain in-flight jobs and
+ * the WebSocket stream, flush the event store, close every database
+ * connection, then exit 0 — or force-exit 1 if any of that takes longer
+ * than `config.GRACEFUL_SHUTDOWN_TIMEOUT` seconds.
+ *
+ * In-flight tasks are drained (via `closeApp`, which awaits the job
+ * worker's stop()) rather than force-failed: anything still running when
+ * the drain window elapses stays "active" in the job store and is resumed
+ * by the next `JobWorker.start()` (`recoverIncompleteJobs()` resets it to
+ * "pending" for retry) — see `docs/e2e-testing.md` and
+ * `tests/shutdown.test.ts` for the restart-mid-stream scenario.
+ */
 export function setupGracefulShutdown(
   httpServer: any,
   closeApp: (callback?: () => void) => void,
-  config: { GRACEFUL_SHUTDOWN_TIMEOUT?: number }
+  config: { GRACEFUL_SHUTDOWN_TIMEOUT?: number },
+  extras: GracefulShutdownExtras = {},
 ) {
   const logger = createLogger({ module: "shutdown" });
   let isShuttingDown = false;
@@ -128,6 +150,9 @@ export function setupGracefulShutdown(
 
       logger.info("stopping agent sync service");
       stopAgentSync();
+      extras.cleanupService?.stop();
+      extras.reconciliationService?.stop();
+      extras.globalAgentRegistry?.shutdown();
 
       logger.info("failing running tasks");
       try {

backend/src/queue/worker.test.ts

@@ -368,6 +368,84 @@ describe("Background Job Queue & Worker", () => {
     });
   });
 
+  describe("Restart mid-stream (#349 — in-flight jobs resume, not fail)", () => {
+    it("a job still running when the worker stops is resumed and completed by a fresh worker instance", async () => {
+      // Simulates a graceful shutdown that catches a job mid-execution: the
+      // handler never resolves before stop() gives up waiting, so the job
+      // stays "active" in the store rather than being marked failed — exactly
+      // what api/app.ts's close() -> jobWorker.stop(timeoutMs) produces when
+      // the drain window elapses with work still outstanding.
+      let releaseHandler!: () => void;
+      let resolveHandlerStarted!: () => void;
+      const handlerStarted = new Promise<void>((resolve) => {
+        resolveHandlerStarted = resolve;
+      });
+
+      const firstWorker = new JobWorker({
+        jobStore: store,
+        handler: async () => {
+          resolveHandlerStarted();
+          // Hang until explicitly released — outlives the worker's stop()
+          // timeout, so stop() returns while this job is still "active".
+          await new Promise<void>((releaseResolve) => {
+            releaseHandler = releaseResolve;
+          });
+          return { success: true };
+        },
+        pollIntervalMs: 20,
+        autoStart: false,
+      });
+
+      const queue = new JobQueue(store, firstWorker);
+      const job = queue.enqueue({ taskId: "task_mid_stream" });
+
+      firstWorker.start();
+      await handlerStarted;
+
+      // Job is now actively executing. "Shut down" with a short drain
+      // window — the handler is hung, so stop() times out waiting and
+      // returns with the job still active, exactly like a real deploy that
+      // catches a slow task.
+      await firstWorker.stop(50);
+
+      const midShutdownState = store.findById(job.id);
+      expect(midShutdownState?.status).toBe("active");
+
+      // "Restart": a brand-new JobWorker over the SAME store (in a real
+      // process this is the same jobs.db file reopened) — its start() calls
+      // recoverIncompleteJobs(), which is what actually makes the job
+      // resumable rather than lost.
+      const secondWorker = new JobWorker({
+        jobStore: store,
+        handler: async (_job, updateProgress) => {
+          updateProgress(100);
+          return { success: true, resumed: true };
+        },
+        pollIntervalMs: 20,
+        autoStart: false,
+      });
+
+      const completed = new Promise<void>((resolve) => {
+        secondWorker.onJobCompleted = (completedJob) => {
+          if (completedJob.id === job.id) resolve();
+        };
+      });
+
+      secondWorker.start();
+      await completed;
+      await secondWorker.stop();
+
+      const finalState = store.findById(job.id);
+      expect(finalState?.status).toBe("completed");
+      expect(finalState?.progress).toBe(100);
+      expect(finalState?.completedAt).toBeDefined();
+
+      // Release the first handler's promise so it doesn't leak a dangling
+      // timer/microtask into later tests.
+      releaseHandler();
+    });
+  });
+
   describe("Queue Stats & Admin Operations", () => {
     it("reports accurate stats across pending, active, completed, failed and dead-letter", () => {
       const now = new Date().toISOString();