From f7db5e7cc2c75cd797e74497a87abe6a89920f90 Mon Sep 17 00:00:00 2001 From: Lyon <88232613+pikasTech@users.noreply.github.com> Date: Wed, 24 Jun 2026 01:34:45 +0800 Subject: [PATCH] fix(v03): drain workbench sse on shutdown (#2001) --- internal/cloud/bun-server.ts | 75 +++++++++++++++- internal/cloud/server-workbench-http.ts | 114 ++++++++++++++++++++++++ 2 files changed, 185 insertions(+), 4 deletions(-) diff --git a/internal/cloud/bun-server.ts b/internal/cloud/bun-server.ts index 87df3f7d..871e5b18 100644 --- a/internal/cloud/bun-server.ts +++ b/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() { - server.stop(true); - innerServer.close(); + stop + }; +} + +function installCloudApiSignalHandlers({ stop, logger }: { stop: (reason?: string, signal?: string | null) => Promise, 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" }; diff --git a/internal/cloud/server-workbench-http.ts b/internal/cloud/server-workbench-http.ts index e8986c98..e4e50c58 100644 --- a/internal/cloud/server-workbench-http.ts +++ b/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;