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

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;