fix(v03): drain workbench sse on shutdown (#2001)

internal/cloud/bun-server.ts

@@ -1,12 +1,14 @@
 import { createCloudApiServer } from "./server.ts";
 import { createHwpodNodeWsRegistry } from "./hwpod-node-ws-registry.ts";
+import { drainWorkbenchRealtimeConnections } from "./server-workbench-http.ts";
 import { classifyRuntimeDbError } from "../db/runtime-store.ts";
 
 const POSTGRES_TRANSIENT_UNHANDLED_REJECTION_BOUNDARY = Symbol.for("hwlab.cloud.postgresTransientUnhandledRejectionBoundaryInstalled");
 
 export async function createCloudApiBunServer(options: any = {}) {
   const env = options.env ?? process.env;
-  installPostgresTransientUnhandledRejectionBoundary({ env, logger: options.logger ?? console });
+  const logger = options.logger ?? console;
+  installPostgresTransientUnhandledRejectionBoundary({ env, logger });
   const hwpodNodeWsRegistry = options.hwpodNodeWsRegistry || createHwpodNodeWsRegistry({ ...options, env });
   const innerServer = createCloudApiServer({ ...options, env, hwpodNodeWsRegistry });
   await new Promise((resolve) => innerServer.listen(0, "127.0.0.1", resolve));
@@ -45,19 +47,84 @@ export async function createCloudApiBunServer(options: any = {}) {
     }
   });
 
+  let removeSignalHandlers = () => {};
+  let stopping = false;
+  const stop = async (reason = "manual", signal: string | null = null) => {
+    if (stopping) return { ok: true, alreadyStopping: true };
+    stopping = true;
+    removeSignalHandlers();
+    const sseDrainTimeoutMs = parsePositiveInteger(options.sseDrainTimeoutMs ?? env.HWLAB_WORKBENCH_SSE_DRAIN_TIMEOUT_MS, 2500);
+    const innerCloseTimeoutMs = parsePositiveInteger(options.innerCloseTimeoutMs ?? env.HWLAB_CLOUD_API_INNER_CLOSE_TIMEOUT_MS, 8000);
+    logger?.info?.({ event: "cloud_api_shutdown_started", reason, signal, sseDrainTimeoutMs, innerCloseTimeoutMs, valuesRedacted: true });
+    const sseDrain = await drainWorkbenchRealtimeConnections({ reason, signal, timeoutMs: sseDrainTimeoutMs, env });
+    try {
+      server.stop(false);
+    } catch (error) {
+      logger?.warn?.({ event: "cloud_api_bun_stop_failed", reason, signal, errorName: error?.name ?? "Error", message: error instanceof Error ? error.message : String(error ?? "unknown"), valuesRedacted: true });
+    }
+    const innerClose = await closeNodeHttpServer(innerServer, innerCloseTimeoutMs);
+    logger?.info?.({ event: "cloud_api_shutdown_completed", reason, signal, sseDrain, innerClose, valuesRedacted: true });
+    return { ok: sseDrain.ok === true && innerClose.ok === true, sseDrain, innerClose };
+  };
+  if (options.installSignalHandlers !== false) {
+    removeSignalHandlers = installCloudApiSignalHandlers({ stop, logger });
+  }
+
   return {
     server,
     innerServer,
     hwpodNodeWsRegistry,
     port: server.port,
     url: `http://${host === "0.0.0.0" ? "127.0.0.1" : host}:${server.port}`,
-    stop() {
+    stop
-      server.stop(true);
+  };
-      innerServer.close();
+}
+
+function installCloudApiSignalHandlers({ stop, logger }: { stop: (reason?: string, signal?: string | null) => Promise<any>, logger?: any }) {
+  const handlers: Array<{ signal: string, handler: () => void }> = [];
+  for (const signal of ["SIGINT", "SIGTERM"]) {
+    const handler = () => {
+      void stop("process_signal", signal).catch((error) => {
+        logger?.error?.({ event: "cloud_api_shutdown_failed", signal, errorName: error?.name ?? "Error", message: error instanceof Error ? error.message : String(error ?? "unknown"), valuesRedacted: true });
+      }).finally(() => {
+        process.exit(0);
+      });
+    };
+    process.once(signal, handler);
+    handlers.push({ signal, handler });
+  }
+  return () => {
+    for (const { signal, handler } of handlers.splice(0)) {
+      process.off(signal, handler);
     }
   };
 }
 
+function closeNodeHttpServer(server: any, timeoutMs: number): Promise<{ ok: boolean, timeout: boolean, errorName: string | null, message: string | null }> {
+  return new Promise((resolve) => {
+    let settled = false;
+    const finish = (ok: boolean, timeout = false, error: any = null) => {
+      if (settled) return;
+      settled = true;
+      if (timeoutHandle) clearTimeout(timeoutHandle);
+      resolve({ ok, timeout, errorName: error?.name ?? null, message: error instanceof Error ? error.message : null });
+    };
+    const timeoutHandle = setTimeout(() => {
+      try { server.closeAllConnections?.(); } catch {}
+      finish(false, true, null);
+    }, timeoutMs);
+    timeoutHandle.unref?.();
+    try {
+      server.close((error: any) => {
+        finish(!error, false, error);
+      });
+      server.closeIdleConnections?.();
+    } catch (error) {
+      finish(false, false, error);
+    }
+  });
+}
+
 export function installPostgresTransientUnhandledRejectionBoundary(options: any = {}) {
   const proc: any = options.process ?? process;
   if (!proc || typeof proc.on !== "function") return { ok: false, installed: false, reason: "process-unavailable" };

internal/cloud/server-workbench-http.ts

@@ -22,11 +22,44 @@ const DEFAULT_PAGE_LIMIT = 50;
 const DEFAULT_SESSION_LIST_LIMIT = 20;
 const MAX_PAGE_LIMIT = 100;
 const DEFAULT_WORKBENCH_SSE_HEARTBEAT_MS = 15000;
+const WORKBENCH_REALTIME_DRAIN_TIMEOUT_MS = 2500;
 const WORKBENCH_SESSION_LIST_PAGE_FAMILIES = Object.freeze(["sessions"]);
 const WORKBENCH_SESSION_DETAIL_FAMILIES = Object.freeze(["sessions", "messages", "parts", "turns", "checkpoints"]);
 const WORKBENCH_SESSION_MESSAGE_PAGE_FAMILIES = Object.freeze(["sessions", "messages", "parts", "turns", "checkpoints"]);
 const WORKBENCH_TRACE_EVENT_METADATA_FAMILIES = Object.freeze(["sessions", "messages", "turns", "checkpoints"]);
 const WORKBENCH_TRACE_EVENT_PAGE_FAMILIES = Object.freeze(["traceEvents"]);
+const activeWorkbenchRealtimeConnections = new Set();
+
+export async function drainWorkbenchRealtimeConnections(options = {}) {
+  const reason = textValue(options.reason) || "server_shutdown";
+  const signal = textValue(options.signal) || null;
+  const timeoutMs = parsePositiveInteger(options.timeoutMs, WORKBENCH_REALTIME_DRAIN_TIMEOUT_MS);
+  const connections = [...activeWorkbenchRealtimeConnections].filter((connection) => connection?.isActive?.());
+  for (const connection of connections) {
+    try {
+      connection.close({ reason, signal });
+    } catch {
+      // Best effort: one broken client must not block shutdown drain for the rest.
+    }
+  }
+  const deadline = Date.now() + timeoutMs;
+  while (Date.now() < deadline) {
+    const activeAfterClose = connections.filter((connection) => activeWorkbenchRealtimeConnections.has(connection) && connection?.isActive?.()).length;
+    if (activeAfterClose === 0) break;
+    await sleepMs(Math.min(50, Math.max(1, deadline - Date.now())));
+  }
+  const activeAfter = connections.filter((connection) => activeWorkbenchRealtimeConnections.has(connection) && connection?.isActive?.()).length;
+  return {
+    ok: activeAfter === 0,
+    activeBefore: connections.length,
+    closed: Math.max(0, connections.length - activeAfter),
+    activeAfter,
+    timeout: activeAfter > 0,
+    reason,
+    signal
+  };
+}
+
 export async function handleWorkbenchReadModelHttp(request, response, url, options = {}) {
   try {
     const perf = options.backendPerformance;
@@ -97,6 +130,12 @@ export async function handleWorkbenchRealtimeHttp(request, response, url, option
     const traceStore = options.traceStore ?? defaultCodeAgentTraceStore;
     const readModel = createWorkbenchReadModel(options, auth.actor);
     let closed = false;
+    let realtimeCloseReason = "client_close";
+    let realtimeCloseSignal = null;
+    let realtimeSessionId = requestedSessionId ?? null;
+    let realtimeTraceId = requestedTraceId ?? null;
+    let realtimeThreadId = null;
+    const realtimeStartedAtMs = Date.now();
     const cleanup = [];
 
   response.writeHead(200, {
@@ -118,6 +157,37 @@ export async function handleWorkbenchRealtimeHttp(request, response, url, option
     perf?.recordPhase({ phase: "sse_write", durationMs: nowMs() - startedAt, outcome: "ok" });
   };
 
+  const realtimeConnection = {
+    close(fields = {}) {
+      if (closed || response.destroyed || response.writableEnded) return false;
+      const reason = textValue(fields.reason) || "server_shutdown";
+      realtimeCloseReason = reason;
+      realtimeCloseSignal = textValue(fields.signal) || null;
+      try {
+        writeEvent("workbench.server_draining", {
+          type: "server.draining",
+          status: "closing",
+          reason,
+          signal: realtimeCloseSignal,
+          sessionId: realtimeSessionId,
+          threadId: realtimeThreadId,
+          traceId: realtimeTraceId,
+          observedAt: new Date().toISOString()
+        });
+        response.end();
+        return true;
+      } catch {
+        try { response.destroy(); } catch {}
+        return false;
+      }
+    },
+    isActive() {
+      return !closed && !response.destroyed && !response.writableEnded;
+    }
+  };
+  activeWorkbenchRealtimeConnections.add(realtimeConnection);
+  cleanup.push(() => activeWorkbenchRealtimeConnections.delete(realtimeConnection));
+
   writeEvent("workbench.connected", {
     type: "connected",
     status: "connected",
@@ -135,6 +205,9 @@ export async function handleWorkbenchRealtimeHttp(request, response, url, option
   const activeTraceId = requestedTraceId
     ?? safeTraceId(initialSession?.lastTraceId ?? initialSessionSnapshot.lastTraceId ?? initialSessionSnapshot.currentTraceId)
     ?? null;
+  realtimeSessionId = streamSessionId ?? requestedSessionId ?? null;
+  realtimeTraceId = activeTraceId ?? requestedTraceId ?? null;
+  realtimeThreadId = streamThreadId ?? null;
   attachWorkbenchRealtimeOtelContext(request, {
     sessionId: streamSessionId ?? requestedSessionId,
     traceId: activeTraceId,
@@ -256,7 +329,17 @@ export async function handleWorkbenchRealtimeHttp(request, response, url, option
   cleanup.push(() => clearInterval(heartbeat));
 
   response.on("close", () => {
+    if (closed) return;
     closed = true;
+    emitWorkbenchRealtimeClosedOtelSpan(request, options.env, {
+      reason: realtimeCloseReason,
+      signal: realtimeCloseSignal,
+      startedAtMs: realtimeStartedAtMs,
+      sessionId: realtimeSessionId,
+      threadId: realtimeThreadId,
+      traceId: realtimeTraceId,
+      activeConnectionCount: activeWorkbenchRealtimeConnections.size
+    });
     for (const item of cleanup.splice(0)) item();
   });
   } catch (error) {
@@ -326,6 +409,33 @@ function emitWorkbenchRealtimeAcceptedOtelSpan(request, env = process.env) {
   }).catch(() => undefined);
 }
 
+function emitWorkbenchRealtimeClosedOtelSpan(request, env = process.env, fields = {}) {
+  const context = request?.hwlabHttpRequestContext;
+  if (!context?.traceId) return;
+  const endedAtMs = Date.now();
+  const startedAtMs = Number(fields.startedAtMs ?? context.startedAtMs ?? endedAtMs);
+  const reason = textValue(fields.reason) || "client_close";
+  void emitHttpServerRequestSpan(context, env, {
+    startTimeMs: Number.isFinite(startedAtMs) ? startedAtMs : endedAtMs,
+    endTimeMs: endedAtMs,
+    statusCode: 200,
+    method: "GET",
+    route: "/v1/workbench/events",
+    attributes: {
+      ...(context.otelAttributes ?? {}),
+      "http.closed_by": reason,
+      "workbench.sse.lifecycle": "closed",
+      "workbench.sse.close_reason": reason,
+      "workbench.sse.close_signal": fields.signal ?? null,
+      "workbench.sse.duration_ms": Math.max(0, endedAtMs - (Number.isFinite(startedAtMs) ? startedAtMs : endedAtMs)),
+      "workbench.sse.session_id": fields.sessionId ?? null,
+      "workbench.sse.thread_id": fields.threadId ?? null,
+      "workbench.sse.trace_id": fields.traceId ?? null,
+      "workbench.sse.active_connection_count": Number.isFinite(Number(fields.activeConnectionCount)) ? Number(fields.activeConnectionCount) : null
+    }
+  }).catch(() => undefined);
+}
+
 function handleWorkbenchRealtimeFailure(request, response, url, options = {}, error) {
   const classified = classifyWorkbenchReadModelFailure(error);
   const context = request?.hwlabHttpRequestContext;
@@ -412,6 +522,10 @@ function nowMs() {
   return Date.now();
 }
 
+function sleepMs(ms) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
 function recordWorkbenchTurnReadMetric(options = {}, url, statusCode, body = {}, startedAt = nowMs()) {
   const performanceStore = options.backendPerformanceStore ?? options.backendPerformance;
   if (!performanceStore?.recordWorkbenchTurnRead) return;