mirrors/pikasTech-HWLAB
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Lyon
2026-06-24 01:34:45 +08:00
committed by GitHub
parent 0a8da60bd9
commit f7db5e7cc2
2 changed files with 185 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/cloud/bun-server.ts
+71 -4
View File
@@ -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<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
+114
View File
@@ -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;