diff --git a/.env.example b/.env.example index 102f5c5fa..19b417d51 100644 --- a/.env.example +++ b/.env.example @@ -20,11 +20,15 @@ ENABLE_API_KEY_VACUUM_FILTER="true" # PostgreSQL 连接池配置(postgres.js) # 说明: -# - 这些值是“每个应用进程”的连接池上限;k8s 多副本时需要按副本数分摊 -# - 默认值:生产环境 20,开发环境 10(可按需覆盖) +# - DB_POOL_MAX 是每个应用进程内 data/control/writer 三类 pool 的连接总预算 +# - 默认拆分:生产环境 20 = 15/4/1,开发与测试环境 10 = 7/2/1 +# - 每条 pool 另有有界 outstanding admission,满载时以 DB_POOL_ADMISSION_EXCEEDED 快速失败 +# - k8s 多副本时仍需按副本数分摊这个总预算 DB_POOL_MAX=20 DB_POOL_IDLE_TIMEOUT=20 # 空闲连接回收(秒) DB_POOL_CONNECT_TIMEOUT=10 # 建立连接超时(秒) +DB_STATEMENT_TIMEOUT_MS=90000 # 活动 SQL 最长执行时间(毫秒,需低于流式结算 120 秒上限) +DB_LOCK_TIMEOUT_MS=5000 # SQL 等待数据库锁的最长时间(毫秒) # message_request 写入模式 # - async:异步批量写入(默认,降低 DB 写放大与连接占用) @@ -41,6 +45,11 @@ DB_USER=postgres DB_PASSWORD=your-secure-password_change-me DB_NAME=claude_code_hub +# Redis 活动命令超时(毫秒) +# 命令 Promise 超时后 5 秒若仍无任何响应进展,客户端会销毁并重建 socket。 +# socket timer 仅在存在等待回复的命令时启用,不影响空闲共享连接。 +REDIS_COMMAND_TIMEOUT_MS=10000 + # 应用配置 APP_PORT=23000 APP_URL= # 应用访问地址(留空自动检测,生产环境建议显式配置) diff --git a/Dockerfile b/Dockerfile index 5518ae51e..e31d5a301 100644 --- a/Dockerfile +++ b/Dockerfile @@ -31,5 +31,7 @@ RUN mkdir -p /app/reports # --report-on-fatalerror / --report-uncaught-exception:在 native 段错误或 # 未捕获异常时写出 JSON 诊断报告(包含原生堆栈、libuv 句柄、JS 堆等) +# --report-exclude-env:诊断报告不得持久化 ADMIN_TOKEN、DSN、Redis、Langfuse +# 与 Provider credentials 等运行时环境变量 # --report-directory:指向 /app/reports 以便挂卷持久化 -CMD ["node", "--report-on-fatalerror", "--report-uncaught-exception", "--report-directory=/app/reports", "server.js"] +CMD ["node", "--report-on-fatalerror", "--report-uncaught-exception", "--report-exclude-env", "--report-directory=/app/reports", "server.js"] diff --git a/deploy/Dockerfile b/deploy/Dockerfile index 27d5c62bb..b426332f4 100644 --- a/deploy/Dockerfile +++ b/deploy/Dockerfile @@ -58,7 +58,9 @@ COPY --from=build --chown=node:node /app/.next/standalone ./ COPY --from=build --chown=node:node /app/.next/server ./.next/server COPY --from=build --chown=node:node /app/.next/static ./.next/static +RUN mkdir -p /app/reports && chown node:node /app/reports + USER node EXPOSE 3000 -CMD ["node", "server.js"] +CMD ["node", "--report-on-fatalerror", "--report-uncaught-exception", "--report-exclude-env", "--report-directory=/app/reports", "server.js"] diff --git a/deploy/k8s/app/deployment.yaml b/deploy/k8s/app/deployment.yaml index 359d99aa4..430bb5029 100644 --- a/deploy/k8s/app/deployment.yaml +++ b/deploy/k8s/app/deployment.yaml @@ -52,7 +52,7 @@ spec: - name: AUTO_MIGRATE value: "true" - name: DB_POOL_MAX - value: "30" + value: "24" - name: DB_POOL_IDLE_TIMEOUT value: "20" - name: DB_POOL_CONNECT_TIMEOUT diff --git a/server.js b/server.js index f7ede45ed..8c4469a34 100644 --- a/server.js +++ b/server.js @@ -60,6 +60,9 @@ const RESERVED_INTERNAL_HEADER_PREFIX = "x-cch-"; // bytes to make a misbehaving / malicious client a bounded-memory event. const MAX_PENDING_FRAMES = 64; const MAX_PENDING_BYTES = 64 * 1024 * 1024; // 64 MiB across all queued frames +const MAX_PENDING_OUTBOUND_BYTES = 1024 * 1024; // 1 MiB per client WebSocket +const REQUEST_BODY_DRAIN_TIMEOUT_MS = 30_000; +const OUTBOUND_SEND_TIMEOUT_MS = 30_000; // Maximum payload size for any single inbound WS frame. The default `ws` // limit is 100 MiB. We pick 32 MiB to accommodate Codex requests that ship @@ -88,16 +91,137 @@ function log(level, msg, extra) { } } -function safeSend(ws, data) { - try { - if (ws.readyState === 1 /* OPEN */) { - ws.send(typeof data === "string" ? data : JSON.stringify(data)); - return true; +const outboundSendStates = new WeakMap(); + +function invalidateOutboundSends(ws, options = {}) { + const state = outboundSendStates.get(ws); + if (!state) return; + state.active = false; + state.generation += 1; + if (state.callbackDeadlineId) clearTimeout(state.callbackDeadlineId); + state.callbackDeadlineId = null; + state.pending.length = 0; + state.pendingBytes = 0; + state.inFlight = false; + if (options.destroyResponses) { + for (const response of state.pressuredResponses) { + if (!response.destroyed) response.destroy(); } + } + state.pressuredResponses.clear(); + outboundSendStates.delete(ws); +} + +function failOutboundSends(ws, state, failure) { + if (!state.active) return; + log("warn", "ws_send_failed", { + reason: failure.reason, + error: failure.error ? String(failure.error) : undefined, + }); + invalidateOutboundSends(ws, { destroyResponses: true }); + if (failure.onFailure) { + failure.onFailure(failure.reason); + return; + } + try { + ws.close(1011, failure.reason); + } catch (err) { + log("warn", "ws_client_close_failed", { error: String(err) }); + } +} + +function flushOutboundSends(ws, state) { + if (!state.active || state.inFlight) return; + const next = state.pending.shift(); + if (!next) { + for (const response of state.pressuredResponses) response.resume(); + state.pressuredResponses.clear(); + return; + } + if (ws.readyState !== 1 /* OPEN */) { + failOutboundSends(ws, state, { + reason: "outbound_socket_closed", + onFailure: next.onFailure, + }); + return; + } + + state.inFlight = true; + const generation = state.generation; + state.callbackDeadlineId = setTimeout(() => { + failOutboundSends(ws, state, { + reason: "outbound_send_timeout", + onFailure: next.onFailure, + }); + }, OUTBOUND_SEND_TIMEOUT_MS); + try { + ws.send(next.payload, (err) => { + if (!state.active || state.generation !== generation) return; + if (state.callbackDeadlineId) clearTimeout(state.callbackDeadlineId); + state.callbackDeadlineId = null; + state.inFlight = false; + state.pendingBytes -= next.bytes; + if (err) { + failOutboundSends(ws, state, { + reason: "outbound_send_error", + error: err, + onFailure: next.onFailure, + }); + return; + } + next.onSuccess?.(); + flushOutboundSends(ws, state); + }); } catch (err) { - log("warn", "ws_send_failed", { error: String(err) }); + failOutboundSends(ws, state, { + reason: "outbound_send_error", + error: err, + onFailure: next.onFailure, + }); } - return false; +} + +function safeSend(ws, data, options = {}) { + if (ws.readyState !== 1 /* OPEN */) { + options.onFailure?.("outbound_socket_closed"); + return false; + } + const payload = typeof data === "string" ? data : JSON.stringify(data); + const bytes = Buffer.byteLength(payload, "utf8"); + let state = outboundSendStates.get(ws); + if (!state) { + state = { + active: true, + generation: 0, + inFlight: false, + callbackDeadlineId: null, + pending: [], + pendingBytes: 0, + pressuredResponses: new Set(), + }; + outboundSendStates.set(ws, state); + } + if (options.response) { + options.response.pause(); + state.pressuredResponses.add(options.response); + } + if (!state.active || state.pendingBytes + bytes > MAX_PENDING_OUTBOUND_BYTES) { + failOutboundSends(ws, state, { + reason: "outbound_backpressure", + onFailure: options.onFailure, + }); + return false; + } + + state.pending.push({ + payload, + bytes, + onSuccess: options.onSuccess, + onFailure: options.onFailure, + }); + state.pendingBytes += bytes; + flushOutboundSends(ws, state); + return true; } function emitErrorEvent(ws, code, message) { @@ -132,22 +256,32 @@ async function handleWebSocketConnection(ws, req) { const pending = []; let pendingBytes = 0; let closed = false; + let closing = false; + let turnSequence = 0; // Track the in-flight internal HTTP ClientRequest so we can abort it when // the client WebSocket disconnects mid-stream — otherwise the SSE consumer // (and provider concurrency / breaker counters) keep running for minutes. let currentInternalReq = null; + let currentInternalRes = null; + let currentTurnSettle = null; const abortCurrentInternalReq = () => { - if (!currentInternalReq) return; const reqToDestroy = currentInternalReq; currentInternalReq = null; try { - if (!reqToDestroy.destroyed) { + if (reqToDestroy && !reqToDestroy.destroyed) { reqToDestroy.destroy(); } } catch { // ignore } + const resToDestroy = currentInternalRes; + currentInternalRes = null; + try { + if (resToDestroy && !resToDestroy.destroyed) resToDestroy.destroy(); + } catch { + currentInternalRes = null; + } }; const cleanupUpstreamWsSession = () => { @@ -176,8 +310,11 @@ async function handleWebSocketConnection(ws, req) { const finalize = () => { if (closed) return; closed = true; + currentTurnSettle?.(); + currentTurnSettle = null; abortCurrentInternalReq(); dropPendingFrames(); + invalidateOutboundSends(ws); cleanupUpstreamWsSession(); }; @@ -199,8 +336,11 @@ async function handleWebSocketConnection(ws, req) { return; } closed = true; + currentTurnSettle?.(); + currentTurnSettle = null; abortCurrentInternalReq(); dropPendingFrames(); + invalidateOutboundSends(ws); cleanupUpstreamWsSession(); log("info", "ws_client_close_initiated", { code, reason }); try { @@ -209,6 +349,14 @@ async function handleWebSocketConnection(ws, req) { log("warn", "ws_client_close_failed", { error: String(err) }); } }; + const sendErrorAndClose = (error, close) => { + if (closed || closing) return; + closing = true; + abortCurrentInternalReq(); + dropPendingFrames(); + const finish = () => requestClose(close.code, close.reason); + safeSend(ws, { type: "error", error }, { onSuccess: finish, onFailure: finish }); + }; ws.on("close", finalize); ws.on("error", (err) => { @@ -219,11 +367,41 @@ async function handleWebSocketConnection(ws, req) { }); const processFrame = async (raw) => { - if (closed) return; + if (closed || closing) return; + const turnId = ++turnSequence; + let turnActive = true; + let turnReq = null; + let turnRes = null; + let turnSettle = null; + const destroyLateResource = (resource) => { + try { + if (resource && !resource.destroyed) resource.destroy(); + } catch { + // ignore late transport cleanup errors + } + }; + const registerTurnResource = (clientReq, clientRes, settleTurn) => { + if (closed || closing || !turnActive) { + destroyLateResource(clientReq); + destroyLateResource(clientRes); + return false; + } + turnReq = clientReq; + turnRes = clientRes || turnRes; + currentInternalReq = clientReq; + if (clientRes) currentInternalRes = clientRes; + if (typeof settleTurn === "function") { + turnSettle = settleTurn; + currentTurnSettle = settleTurn; + } + return true; + }; if (typeof raw !== "string") { - emitErrorEvent(ws, "invalid_frame_type", "Only text WebSocket frames are supported"); - requestClose(1003, "binary_not_supported"); + sendErrorAndClose( + { code: "invalid_frame_type", message: "Only text WebSocket frames are supported" }, + { code: 1003, reason: "binary_not_supported" } + ); return; } @@ -267,18 +445,21 @@ async function handleWebSocketConnection(ws, req) { hasPreviousResponseId: typeof body.previous_response_id === "string", }); - await forwardToInternalHttp( - ws, - req, - body, - responsesWsSessionId, - (clientReq) => { - currentInternalReq = clientReq; - }, - requestClose - ); - if (!closed) { - currentInternalReq = null; + try { + await forwardToInternalHttp( + ws, + req, + body, + responsesWsSessionId, + registerTurnResource, + requestClose + ); + } finally { + turnActive = false; + if (currentInternalReq === turnReq) currentInternalReq = null; + if (currentInternalRes === turnRes) currentInternalRes = null; + if (currentTurnSettle === turnSettle) currentTurnSettle = null; + log("debug", "ws_turn_resources_released", { turnId }); } }; @@ -293,23 +474,27 @@ async function handleWebSocketConnection(ws, req) { await processFrame(next); } finally { inFlight = false; - if (pending.length > 0 && !closed) { + if (pending.length > 0 && !closed && !closing) { void drain().catch((err) => { log("error", "ws_drain_failed", { error: String(err && err.message ? err.message : err), }); - emitErrorEvent(ws, "internal_error", "Failed to process queued request"); - requestClose(1011, "internal_error"); + sendErrorAndClose( + { code: "internal_error", message: "Failed to process queued request" }, + { code: 1011, reason: "internal_error" } + ); }); } } }; ws.on("message", (data, isBinary) => { - if (closed) return; + if (closed || closing) return; if (isBinary) { - emitErrorEvent(ws, "invalid_frame_type", "Only text WebSocket frames are supported"); - requestClose(1003, "binary_not_supported"); + sendErrorAndClose( + { code: "invalid_frame_type", message: "Only text WebSocket frames are supported" }, + { code: 1003, reason: "binary_not_supported" } + ); return; } const text = data.toString("utf8"); @@ -320,8 +505,10 @@ async function handleWebSocketConnection(ws, req) { pendingBytes, attemptedFrameSize: size, }); - emitErrorEvent(ws, "too_many_requests", "Pending frame limit exceeded"); - requestClose(1008, "too_many_requests"); + sendErrorAndClose( + { code: "too_many_requests", message: "Pending frame limit exceeded" }, + { code: 1008, reason: "too_many_requests" } + ); return; } pending.push(text); @@ -330,8 +517,10 @@ async function handleWebSocketConnection(ws, req) { log("error", "ws_drain_failed", { error: String(err && err.message ? err.message : err), }); - emitErrorEvent(ws, "internal_error", "Failed to process request"); - requestClose(1011, "internal_error"); + sendErrorAndClose( + { code: "internal_error", message: "Failed to process request" }, + { code: 1011, reason: "internal_error" } + ); }); }); } @@ -415,6 +604,15 @@ async function forwardToInternalHttp( internalHeaders["content-length"] = String(payload.length); await new Promise((resolve) => { + let cleanupRequestBody = () => {}; + let turnFinished = false; + const forceSettleTurn = () => { + if (turnFinished) return false; + turnFinished = true; + cleanupRequestBody(); + resolve(); + return true; + }; const req = http.request( { method: "POST", @@ -427,12 +625,49 @@ async function forwardToInternalHttp( const contentType = (res.headers["content-type"] || "").toLowerCase(); const isSse = contentType.includes("text/event-stream"); let responseSettled = false; + let responseBodyEnded = false; + let terminalSendAcknowledged = false; const settleResponse = () => { if (responseSettled) return false; + if (!responseBodyEnded || !terminalSendAcknowledged) return false; responseSettled = true; + turnFinished = true; + cleanupRequestBody(); resolve(); return true; }; + const acknowledgeTerminalSend = () => { + terminalSendAcknowledged = true; + settleResponse(); + }; + const forceSettleResponse = () => { + if (responseSettled) return false; + responseSettled = true; + return forceSettleTurn(); + }; + if (typeof registerInternalReq === "function") { + const accepted = registerInternalReq(req, res, forceSettleResponse); + if (accepted === false) { + forceSettleResponse(); + return; + } + } + const settleAndClose = (reason) => { + initiateClose(1011, reason); + forceSettleResponse(); + }; + const sendFatalError = (code, message, closeReason) => { + const sent = safeSend( + ws, + { type: "error", error: { code, message } }, + { + response: res, + onSuccess: () => settleAndClose(closeReason), + onFailure: settleAndClose, + } + ); + if (!sent) settleAndClose(closeReason); + }; if (!isSse) { // Upstream returned non-stream JSON (e.g. error response). Collect @@ -441,6 +676,7 @@ async function forwardToInternalHttp( res.on("data", (c) => chunks.push(c)); res.on("end", () => { if (responseSettled) return; + responseBodyEnded = true; const text = Buffer.concat(chunks).toString("utf8"); let parsed; try { @@ -450,47 +686,52 @@ async function forwardToInternalHttp( } const isHttpError = !!(res.statusCode && res.statusCode >= 400); if (isHttpError) { - safeSend(ws, { - type: "error", - status: res.statusCode, - error: - typeof parsed === "object" && parsed && parsed.error - ? parsed.error - : { code: `http_${res.statusCode}`, message: text.slice(0, 512) }, - }); + safeSend( + ws, + { + type: "error", + status: res.statusCode, + error: + typeof parsed === "object" && parsed && parsed.error + ? parsed.error + : { code: `http_${res.statusCode}`, message: text.slice(0, 512) }, + }, + { response: res, onSuccess: acknowledgeTerminalSend, onFailure: settleAndClose } + ); log("info", "ws_terminal_event_sent", { type: "error", source: "json", status: res.statusCode, }); } else { - safeSend(ws, { - type: "response.completed", - response: parsed, - }); + safeSend( + ws, + { type: "response.completed", response: parsed }, + { response: res, onSuccess: acknowledgeTerminalSend, onFailure: settleAndClose } + ); log("info", "ws_terminal_event_sent", { type: "response.completed", source: "json" }); } - settleResponse(); }); res.on("error", (err) => { if (responseSettled) return; - emitErrorEvent( - ws, + sendFatalError( "internal_response_error", - String(err && err.message ? err.message : err) + String(err && err.message ? err.message : err), + "internal_response_error" ); - initiateClose(1011, "internal_response_error"); - settleResponse(); }); res.on("close", () => { if (responseSettled) return; - emitErrorEvent( - ws, + if (responseBodyEnded || res.complete) { + responseBodyEnded = true; + settleResponse(); + return; + } + sendFatalError( "internal_response_closed", - "Internal response closed before a complete JSON body was received" + "Internal response closed before a complete JSON body was received", + "internal_response_closed" ); - initiateClose(1011, "internal_response_closed"); - settleResponse(); }); return; } @@ -504,11 +745,16 @@ async function forwardToInternalHttp( const EVENT_DELIMITER = /\r?\n\r?\n/; const failIfUnsettled = (code, message, closeReason) => { if (responseSettled) return; - if (!sawTerminal) { - emitErrorEvent(ws, code, message); - initiateClose(1011, closeReason); + if (sawTerminal) { + // A terminal protocol event is authoritative. Once the internal + // transport closes, wait only for its WS send acknowledgement; + // otherwise this persistent connection would retain inFlight + // ownership forever when close/error arrives before `end`. + responseBodyEnded = true; + settleResponse(); + return; } - settleResponse(); + sendFatalError(code, message, closeReason); }; const flushEvents = () => { @@ -532,7 +778,11 @@ async function forwardToInternalHttp( // Some upstreams close SSE with [DONE] without a preceding // response.completed. Synthesize one so the client sees a // clean terminal event. - safeSend(ws, { type: "response.completed", response: null }); + safeSend(ws, { type: "response.completed", response: null }, { + response: res, + onSuccess: acknowledgeTerminalSend, + onFailure: settleAndClose, + }); sawTerminal = true; } continue; @@ -542,11 +792,20 @@ async function forwardToInternalHttp( event = JSON.parse(dataText); } catch { // Not JSON; forward as raw string event. - safeSend(ws, { type: "response.output_text.delta", delta: dataText }); + safeSend(ws, { type: "response.output_text.delta", delta: dataText }, { + response: res, + onFailure: settleAndClose, + }); continue; } - safeSend(ws, event); - if (event && typeof event.type === "string" && TERMINAL_EVENT_TYPES.has(event.type)) { + const isTerminalEvent = + event && typeof event.type === "string" && TERMINAL_EVENT_TYPES.has(event.type); + safeSend(ws, event, { + response: res, + onSuccess: isTerminalEvent ? acknowledgeTerminalSend : undefined, + onFailure: settleAndClose, + }); + if (isTerminalEvent) { sawTerminal = true; terminalEventType = event.type; log("info", "ws_terminal_event_sent", { type: event.type, source: "sse" }); @@ -562,26 +821,27 @@ async function forwardToInternalHttp( }); res.on("end", () => { if (responseSettled) return; + responseBodyEnded = true; // Flush any remaining buffered event if (buffer.trim().length > 0) { buffer += "\n\n"; flushEvents(); } if (!sawTerminal) { - emitErrorEvent( - ws, + sendFatalError( "stream_ended_without_terminal", - "Upstream stream ended before emitting a terminal response event" + "Upstream stream ended before emitting a terminal response event", + "stream_ended_without_terminal" ); - initiateClose(1011, "stream_ended_without_terminal"); } else { // OpenAI Responses WebSocket mode is persistent: after a terminal // event, the same client connection can send the next // response.create. Do not close here; only fatal transport/protocol // errors initiate a close handshake. log("info", "ws_turn_completed", { terminalEventType }); + settleResponse(); } - settleResponse(); + if (!sawTerminal) return; }); res.on("error", (err) => { failIfUnsettled( @@ -591,6 +851,11 @@ async function forwardToInternalHttp( ); }); res.on("close", () => { + if (responseBodyEnded || res.complete) { + responseBodyEnded = true; + settleResponse(); + return; + } failIfUnsettled( "internal_response_closed", "Internal response closed before emitting a terminal response event", @@ -600,27 +865,130 @@ async function forwardToInternalHttp( } ); - req.on("error", (err) => { - // ECONNRESET when we destroy() the request on client disconnect is - // expected; downgrade to debug to avoid noisy logs in normal traffic. - const errCode = err && (err.code || err.name); - const isAbort = errCode === "ECONNRESET" || errCode === "ERR_STREAM_PREMATURE_CLOSE"; - if (!isAbort) { - emitErrorEvent( + if (typeof registerInternalReq === "function") { + const accepted = registerInternalReq(req, null, forceSettleTurn); + if (accepted === false) { + if (!req.destroyed) req.destroy(); + forceSettleTurn(); + return; + } + } + + const handleRequestError = (err) => { + if (turnFinished) { + return; + } + turnFinished = true; + cleanupRequestBody(); + let finished = false; + const finish = () => { + if (finished) return; + finished = true; + initiateClose(1011, "internal_request_error"); + resolve(); + }; + const sent = safeSend( + ws, + { + type: "error", + error: { + code: "internal_request_error", + message: String(err && err.message ? err.message : err), + }, + }, + { onSuccess: finish, onFailure: finish } + ); + if (!sent) finish(); + }; + req.on("error", handleRequestError); + let requestEnded = false; + let requestBodyFinished = false; + let requestBodyDeadlineId = null; + const clearRequestBodyListeners = () => { + req.removeListener("drain", finishRequestBody); + req.removeListener("close", abandonRequestBody); + req.removeListener("error", abandonRequestBody); + req.removeListener("abort", abandonRequestBody); + if (requestBodyDeadlineId) { + clearTimeout(requestBodyDeadlineId); + requestBodyDeadlineId = null; + } + }; + const endRequestOnce = () => { + if (requestEnded) return; + requestEnded = true; + req.end(); + }; + const finishRequestBody = () => { + if (requestBodyFinished) return; + requestBodyFinished = true; + clearRequestBodyListeners(); + endRequestOnce(); + }; + const abandonRequestBody = () => { + if (requestBodyFinished) return; + requestBodyFinished = true; + clearRequestBodyListeners(); + if (!turnFinished) { + turnFinished = true; + const finish = () => { + initiateClose(1011, "internal_request_body_closed"); + resolve(); + }; + const sent = safeSend( ws, - "internal_request_error", - String(err && err.message ? err.message : err) + { + type: "error", + error: { + code: "internal_request_body_closed", + message: "Internal request body closed before it was fully written", + }, + }, + { onSuccess: finish, onFailure: finish } ); - initiateClose(1011, "internal_request_error"); + if (!sent) finish(); + return; } resolve(); - }); - - if (typeof registerInternalReq === "function") { - registerInternalReq(req); + }; + const expireRequestBody = () => { + if (requestBodyFinished) return; + requestBodyFinished = true; + clearRequestBodyListeners(); + turnFinished = true; + if (!req.destroyed) req.destroy(); + const finish = () => { + initiateClose(1011, "internal_request_drain_timeout"); + resolve(); + }; + const sent = safeSend( + ws, + { + type: "error", + error: { + code: "internal_request_drain_timeout", + message: "Internal request body remained backpressured past its deadline", + }, + }, + { onSuccess: finish, onFailure: finish } + ); + if (!sent) finish(); + }; + cleanupRequestBody = () => { + if (requestBodyFinished) return; + requestBodyFinished = true; + clearRequestBodyListeners(); + if (!requestEnded && !req.destroyed) req.destroy(); + }; + if (req.write(payload)) { + finishRequestBody(); + } else { + req.once("drain", finishRequestBody); + req.once("close", abandonRequestBody); + req.once("error", abandonRequestBody); + req.once("abort", abandonRequestBody); + requestBodyDeadlineId = setTimeout(expireRequestBody, REQUEST_BODY_DRAIN_TIMEOUT_MS); } - req.write(payload); - req.end(); }); } @@ -742,10 +1110,10 @@ async function main() { // 2. server.close() -> stop accepting; in-flight HTTP finishes // 3. wss.close() -> reject new WS upgrades // 4. Wait for drain -> bounded by SHUTDOWN_DRAIN_MS -// 5. runApplicationCleanup -> Redis / Langfuse / msg buffer / schedulers; bounded -// by SHUTDOWN_CLEANUP_MS. Inside cleanup, asyncTaskManager.cleanupAll() runs -// LAST so streaming responses had a chance to finish during step 4. -// 6. process.exit(0) +// 5. runApplicationCleanup -> abort + join tasks, flush writer, close DB pools, +// then release non-critical resources. SHUTDOWN_CLEANUP_MS is a soft warning; +// the referenced hard watchdog is the final bound for critical barriers. +// 6. Success logs shutdown_complete and exits 0; cleanup failure exits 1. function registerOrchestratedShutdown(server, wss) { let shuttingDown = false; @@ -770,13 +1138,12 @@ function registerOrchestratedShutdown(server, wss) { shuttingDown = true; log("info", "shutdown_received", { signal, drainMs, cleanupMs, hardExitMs }); - // Final safety: even if every step below hangs, this terminates the process. - // .unref() so the timer itself doesn't keep the event loop alive. + // Final safety: even if every step below hangs, this referenced timer keeps + // the process alive until it can terminate with a truthful non-zero status. const hardExit = setTimeout(() => { log("error", "shutdown_hard_exit_watchdog", { hardExitMs }); process.exit(1); }, hardExitMs); - if (typeof hardExit.unref === "function") hardExit.unref(); // 1. Flip readiness BEFORE closing the listener so probes already in flight // see 503 and the Service starts removing this pod from endpoints. @@ -805,29 +1172,42 @@ function registerOrchestratedShutdown(server, wss) { resolve(); } }); - if (wss && typeof wss.close === "function") { + + const closeWss = new Promise((resolve) => { + if (!wss || typeof wss.close !== "function") { + resolve(); + return; + } + try { - wss.close(); + if (wss.close.length === 0) { + wss.close(); + resolve(); + return; + } + wss.close(() => resolve()); } catch (err) { log("warn", "shutdown_wss_close_error", { error: String(err && err.message ? err.message : err), }); + resolve(); } - } + }); + const closeTransports = Promise.all([closeServer, closeWss]); - // 4. Bounded drain — server.close() resolves only after every in-flight - // connection completes; we cap it so a stuck client can't hold us forever. + // 4. Bounded drain — HTTP and WebSocket close only settle after every in-flight + // connection completes; we cap them so a stuck client can't hold us forever. // Clearing the timer on natural close avoids a misleading // "shutdown_drain_timeout" warning during the subsequent cleanup phase. await Promise.race([ - closeServer, + closeTransports, new Promise((resolve) => { const t = setTimeout(() => { log("warn", "shutdown_drain_timeout", { drainMs }); resolve(); }, drainMs); if (typeof t.unref === "function") t.unref(); - closeServer.finally(() => clearTimeout(t)); + closeTransports.finally(() => clearTimeout(t)); }), ]); @@ -839,11 +1219,15 @@ function registerOrchestratedShutdown(server, wss) { log("warn", "shutdown_cleanup_unavailable", { reason: "lifecycle_globals_not_bound", }); + process.exit(1); + return; } } catch (err) { - log("warn", "shutdown_cleanup_error", { + log("error", "shutdown_cleanup_error", { error: String(err && err.message ? err.message : err), }); + process.exit(1); + return; } log("info", "shutdown_complete", { signal }); @@ -864,6 +1248,7 @@ module.exports = { registerOrchestratedShutdown, WS_MAX_PAYLOAD_BYTES, MAX_PENDING_BYTES, + MAX_PENDING_OUTBOUND_BYTES, }; if (require.main === module) { diff --git a/src/app/v1/[...route]/route.ts b/src/app/v1/[...route]/route.ts index ec4c22c9f..3a3c82922 100644 --- a/src/app/v1/[...route]/route.ts +++ b/src/app/v1/[...route]/route.ts @@ -8,6 +8,7 @@ import { handleOpenAICompatibleModels, } from "@/app/v1/_lib/models/available-models"; import { handleProxyRequest } from "@/app/v1/_lib/proxy-handler"; +import { withDataDbScope } from "@/drizzle/db"; import { logger } from "@/lib/logger"; import { sensitiveWordDetector } from "@/lib/sensitive-word-detector"; import { SessionTracker } from "@/lib/session-tracker"; @@ -58,10 +59,14 @@ app.all("*", handleProxyRequest); export { app as v1App }; -export const GET = handle(app); -export const POST = handle(app); -export const PUT = handle(app); -export const DELETE = handle(app); -export const PATCH = handle(app); -export const OPTIONS = handle(app); -export const HEAD = handle(app); +const routeHandler = withDataDbScope(handle(app)); + +export { + routeHandler as GET, + routeHandler as POST, + routeHandler as PUT, + routeHandler as DELETE, + routeHandler as PATCH, + routeHandler as OPTIONS, + routeHandler as HEAD, +}; diff --git a/src/app/v1/_lib/proxy-handler.ts b/src/app/v1/_lib/proxy-handler.ts index 1fbde5c27..ad5fd55c2 100644 --- a/src/app/v1/_lib/proxy-handler.ts +++ b/src/app/v1/_lib/proxy-handler.ts @@ -1,4 +1,5 @@ import type { Context } from "hono"; +import { findSafeDatabaseError } from "@/drizzle/admitted-client"; import { getCachedSystemSettings } from "@/lib/config"; import { logger } from "@/lib/logger"; import { ProxyStatusTracker } from "@/lib/proxy-status-tracker"; @@ -18,6 +19,7 @@ import { ProxySession } from "./proxy/session"; export async function handleProxyRequest(c: Context): Promise { let session: ProxySession | null = null; let cachedSystemSettings: Awaited> | null = null; + let acquiredConcurrencySessionId: string | null = null; try { session = await ProxySession.fromContext(c); try { @@ -29,10 +31,14 @@ export async function handleProxyRequest(c: Context): Promise { cachedSystemSettings.allowNonConversationEndpointProviderFallback ?? true ); } catch (settingsError) { + const databaseError = findSafeDatabaseError(settingsError); logger.warn( "[ProxyHandler] Failed to load proxy system settings, fallback highConcurrency=false and rawCrossProviderFallback=false", { - error: settingsError, + error: + databaseError?.message ?? + (settingsError instanceof Error ? settingsError.message : String(settingsError)), + databaseCode: databaseError?.code, } ); session.setHighConcurrencyModeEnabled(false); @@ -88,6 +94,7 @@ export async function handleProxyRequest(c: Context): Promise { // 9. 增加并发计数(在所有检查通过后,请求开始前)- 跳过 count_tokens if (session.sessionId && session.getEndpointPolicy().trackConcurrentRequests) { await SessionTracker.incrementConcurrentCount(session.sessionId); + acquiredConcurrencySessionId = session.sessionId; } // 10. 记录请求开始 @@ -134,7 +141,12 @@ export async function handleProxyRequest(c: Context): Promise { return finalResponse; } catch (error) { - logger.error("Proxy handler error:", error); + const databaseError = findSafeDatabaseError(error); + logger.error("Proxy handler error:", { + error: databaseError?.message ?? (error instanceof Error ? error.message : String(error)), + databaseCode: databaseError?.code, + databasePool: databaseError?.pool, + }); if (session) { return await ProxyErrorHandler.handle(session, error); } @@ -146,8 +158,8 @@ export async function handleProxyRequest(c: Context): Promise { return ProxyResponses.buildError(500, "代理请求发生未知错误"); } finally { // 11. 减少并发计数(确保无论成功失败都执行)- 跳过 count_tokens - if (session?.sessionId && session.getEndpointPolicy().trackConcurrentRequests) { - await SessionTracker.decrementConcurrentCount(session.sessionId); + if (acquiredConcurrencySessionId) { + await SessionTracker.decrementConcurrentCount(acquiredConcurrencySessionId); } } } diff --git a/src/app/v1/_lib/proxy/demand-driven-response-pump.test.ts b/src/app/v1/_lib/proxy/demand-driven-response-pump.test.ts new file mode 100644 index 000000000..965b35d62 --- /dev/null +++ b/src/app/v1/_lib/proxy/demand-driven-response-pump.test.ts @@ -0,0 +1,406 @@ +import { describe, expect, it, vi } from "vitest"; +import { + createDemandDrivenResponsePump, + type DemandDrivenResponsePumpCompletion, +} from "./demand-driven-response-pump"; + +const encoder = new TextEncoder(); +const decoder = new TextDecoder(); + +function nextTurn(): Promise { + return new Promise((resolve) => setImmediate(resolve)); +} + +function trackReaderRelease(source: ReadableStream) { + const reader = source.getReader(); + const releaseLock = vi.spyOn(reader, "releaseLock"); + vi.spyOn(source, "getReader").mockReturnValue(reader); + return releaseLock; +} + +async function readText(reader: ReadableStreamDefaultReader): Promise { + let text = ""; + while (true) { + const result = await reader.read(); + if (result.done) return text; + text += decoder.decode(result.value, { stream: true }); + } +} + +describe("createDemandDrivenResponsePump", () => { + it("primes one chunk without reading ahead past an unconsumed pending chunk", async () => { + const chunks = ["one", "two", "three"].map((chunk) => encoder.encode(chunk)); + let pullCount = 0; + const onReadStart = vi.fn(); + const observed: string[] = []; + const source = new ReadableStream( + { + pull(controller) { + pullCount++; + const chunk = chunks.shift(); + if (chunk) { + controller.enqueue(chunk); + } else { + controller.close(); + } + }, + }, + { highWaterMark: 0 } + ); + + const pump = createDemandDrivenResponsePump({ + source, + onReadStart, + onChunk: (chunk) => observed.push(decoder.decode(chunk)), + }); + await nextTurn(); + await nextTurn(); + + expect(pullCount).toBe(1); + expect(onReadStart).toHaveBeenCalledTimes(1); + expect(observed).toEqual(["one"]); + + const reader = pump.stream.getReader(); + await expect(reader.read()).resolves.toMatchObject({ done: false }); + await nextTurn(); + await nextTurn(); + + expect(pullCount).toBe(2); + expect(onReadStart).toHaveBeenCalledTimes(2); + expect(observed).toEqual(["one", "two"]); + + await reader.cancel("test complete"); + await expect(pump.completion).resolves.toMatchObject({ + streamEndedNormally: true, + clientAborted: true, + error: null, + }); + }); + + it("delivers chunks in order and settles after the lookahead discovers EOF", async () => { + const chunks = ["one", "two", "three"].map((chunk) => encoder.encode(chunk)); + const observed: string[] = []; + const source = new ReadableStream( + { + pull(controller) { + const chunk = chunks.shift(); + if (chunk) { + controller.enqueue(chunk); + } else { + controller.close(); + } + }, + }, + { highWaterMark: 0 } + ); + const releaseLock = trackReaderRelease(source); + + const pump = createDemandDrivenResponsePump({ + source, + onChunk: (chunk) => observed.push(decoder.decode(chunk)), + }); + const text = await readText(pump.stream.getReader()); + const completion = await pump.completion; + + expect(text).toBe("onetwothree"); + expect(observed).toEqual(["one", "two", "three"]); + expect(completion).toEqual({ + streamEndedNormally: true, + clientAborted: false, + error: null, + }); + expect(pump.getState()).toBe("closed"); + expect(releaseLock).toHaveBeenCalledTimes(1); + }); + + it("returns from client cancellation before a blocked source read and drains in background", async () => { + const sourceState: { + controller: ReadableStreamDefaultController | null; + } = { controller: null }; + let pullCount = 0; + const observed: string[] = []; + const source = new ReadableStream( + { + start(controller) { + sourceState.controller = controller; + }, + pull(controller) { + pullCount++; + if (pullCount === 2) { + controller.enqueue(encoder.encode("tail")); + } else if (pullCount === 3) { + controller.close(); + } + }, + }, + { highWaterMark: 0 } + ); + const releaseLock = trackReaderRelease(source); + const onClientCancel = vi.fn(); + const pump = createDemandDrivenResponsePump({ + source, + onChunk: (chunk) => observed.push(decoder.decode(chunk)), + onClientCancel, + }); + const reader = pump.stream.getReader(); + + const cancelResult = await Promise.race([ + reader.cancel("client disconnected").then(() => "cancelled" as const), + nextTurn().then(() => "blocked" as const), + ]); + + expect(cancelResult).toBe("cancelled"); + expect(pump.getState()).toBe("draining"); + expect(pump.wasClientAborted()).toBe(true); + expect(onClientCancel).toHaveBeenCalledWith("client disconnected"); + expect(pullCount).toBe(1); + + sourceState.controller?.enqueue(encoder.encode("head")); + const completion = await pump.completion; + + expect(completion).toEqual({ + streamEndedNormally: true, + clientAborted: true, + error: null, + }); + expect(observed).toEqual(["head", "tail"]); + expect(releaseLock).toHaveBeenCalledTimes(1); + }); + + it("settles a pending downstream read when an external client signal starts drain", async () => { + const sourceState: { + controller: ReadableStreamDefaultController | null; + } = { controller: null }; + const source = new ReadableStream( + { + start(controller) { + sourceState.controller = controller; + }, + }, + { highWaterMark: 0 } + ); + const pump = createDemandDrivenResponsePump({ source, onChunk: vi.fn() }); + const pendingRead = pump.stream.getReader().read(); + const clientError = new Error("client signal aborted"); + + pump.startDrain(clientError); + const outcome = await Promise.race([ + pendingRead.then( + (result) => ({ kind: "resolved" as const, result }), + (error) => ({ kind: "rejected" as const, error }) + ), + nextTurn().then(() => ({ kind: "pending" as const })), + ]); + + sourceState.controller?.close(); + await pump.completion; + expect(outcome).toEqual({ kind: "rejected", error: clientError }); + }); + + it("hard-cancels an abort-insensitive source while preserving client-aborted state", async () => { + const cancel = vi.fn(); + const source = new ReadableStream( + { + cancel, + }, + { highWaterMark: 0 } + ); + const pump = createDemandDrivenResponsePump({ source, onChunk: vi.fn() }); + const clientError = new Error("client disconnected"); + const timeoutError = new Error("drain timeout"); + + pump.startDrain(clientError); + pump.cancelSource(timeoutError); + const completion = await pump.completion; + + expect(cancel).toHaveBeenCalledWith(timeoutError); + expect(completion).toEqual({ + streamEndedNormally: false, + clientAborted: true, + error: timeoutError, + }); + expect(pump.getState()).toBe("closed"); + }); + + it("settles hard cancellation even when the source cancel promise never resolves", async () => { + const cancel = vi.fn(() => new Promise(() => {})); + const source = new ReadableStream( + { + cancel, + }, + { highWaterMark: 0 } + ); + const releaseLock = trackReaderRelease(source); + const pump = createDemandDrivenResponsePump({ source, onChunk: vi.fn() }); + const timeoutError = new Error("drain timeout"); + + pump.startDrain(new Error("client disconnected")); + pump.cancelSource(timeoutError); + + await expect(pump.completion).resolves.toEqual({ + streamEndedNormally: false, + clientAborted: true, + error: timeoutError, + }); + expect(cancel).toHaveBeenCalledWith(timeoutError); + expect(releaseLock).toHaveBeenCalledTimes(1); + }); + + it("drains an already pending chunk without observing it twice", async () => { + const chunks = ["one", "two", "three"].map((chunk) => encoder.encode(chunk)); + const observed: string[] = []; + const source = new ReadableStream( + { + pull(controller) { + const chunk = chunks.shift(); + if (chunk) { + controller.enqueue(chunk); + } else { + controller.close(); + } + }, + }, + { highWaterMark: 0 } + ); + const pump = createDemandDrivenResponsePump({ + source, + onChunk: (chunk) => observed.push(decoder.decode(chunk)), + }); + await nextTurn(); + + await pump.stream.cancel("client disconnected"); + await pump.completion; + + expect(observed).toEqual(["one", "two", "three"]); + }); + + it("errors the downstream client then transfers source ownership to the drain", async () => { + const sourceState: { + controller: ReadableStreamDefaultController | null; + } = { controller: null }; + const source = new ReadableStream( + { + start(controller) { + sourceState.controller = controller; + }, + }, + { highWaterMark: 0 } + ); + const releaseLock = trackReaderRelease(source); + const pump = createDemandDrivenResponsePump({ source, onChunk: vi.fn() }); + const reader = pump.stream.getReader(); + const pendingRead = reader.read(); + const clientError = new Error("downstream deadline"); + + pump.errorClient(clientError); + + await expect(pendingRead).rejects.toBe(clientError); + expect(pump.getState()).toBe("draining"); + expect(pump.wasClientAborted()).toBe(false); + + sourceState.controller?.close(); + await expect(pump.completion).resolves.toEqual({ + streamEndedNormally: true, + clientAborted: false, + error: null, + }); + expect(releaseLock).toHaveBeenCalledTimes(1); + }); + + it("preserves a server-owned drain when a later client abort arrives", async () => { + const sourceState: { + controller: ReadableStreamDefaultController | null; + } = { controller: null }; + const source = new ReadableStream( + { + start(controller) { + sourceState.controller = controller; + }, + }, + { highWaterMark: 0 } + ); + const pump = createDemandDrivenResponsePump({ source, onChunk: vi.fn() }); + const reader = pump.stream.getReader(); + const pendingRead = reader.read(); + const idleTimeoutError = new Error("streaming_idle"); + + pump.errorClient(idleTimeoutError); + await expect(pendingRead).rejects.toBe(idleTimeoutError); + + pump.startDrain(new Error("late client abort")); + sourceState.controller?.error(idleTimeoutError); + + await expect(pump.completion).resolves.toEqual({ + streamEndedNormally: false, + clientAborted: false, + error: idleTimeoutError, + }); + }); + + it("forwards the original source error and settles only once during a cancel race", async () => { + const sourceState: { + controller: ReadableStreamDefaultController | null; + } = { controller: null }; + const source = new ReadableStream( + { + start(controller) { + sourceState.controller = controller; + }, + }, + { highWaterMark: 0 } + ); + const releaseLock = trackReaderRelease(source); + const pump = createDemandDrivenResponsePump({ source, onChunk: vi.fn() }); + const reader = pump.stream.getReader(); + const pendingRead = reader.read(); + const completionResults: DemandDrivenResponsePumpCompletion[] = []; + void pump.completion.then((completion) => completionResults.push(completion)); + const upstreamError = new Error("upstream failed"); + + const cancelPromise = reader.cancel("client disconnected"); + sourceState.controller?.error(upstreamError); + + await cancelPromise; + const completion = await pump.completion; + pump.startDrain("late drain"); + pump.errorClient(new Error("late client error")); + await nextTurn(); + + await expect(pendingRead).resolves.toEqual({ done: true, value: undefined }); + expect(completion).toEqual({ + streamEndedNormally: false, + clientAborted: true, + error: upstreamError, + }); + expect(completionResults).toEqual([completion]); + expect(releaseLock).toHaveBeenCalledTimes(1); + }); + + it("errors an active downstream reader when the source fails during the primed read", async () => { + const sourceState: { + controller: ReadableStreamDefaultController | null; + } = { controller: null }; + const source = new ReadableStream( + { + start(controller) { + sourceState.controller = controller; + }, + }, + { highWaterMark: 0 } + ); + const releaseLock = trackReaderRelease(source); + const pump = createDemandDrivenResponsePump({ source, onChunk: vi.fn() }); + const reader = pump.stream.getReader(); + const upstreamError = new Error("upstream failed"); + + sourceState.controller?.error(upstreamError); + + await expect(reader.read()).rejects.toBe(upstreamError); + await expect(pump.completion).resolves.toEqual({ + streamEndedNormally: false, + clientAborted: false, + error: upstreamError, + }); + expect(releaseLock).toHaveBeenCalledTimes(1); + }); +}); diff --git a/src/app/v1/_lib/proxy/demand-driven-response-pump.ts b/src/app/v1/_lib/proxy/demand-driven-response-pump.ts new file mode 100644 index 000000000..137db3dd9 --- /dev/null +++ b/src/app/v1/_lib/proxy/demand-driven-response-pump.ts @@ -0,0 +1,282 @@ +export type DemandDrivenResponsePumpState = "client-active" | "draining" | "finalizing" | "closed"; + +export interface DemandDrivenResponsePumpCompletion { + streamEndedNormally: boolean; + clientAborted: boolean; + error: Error | null; +} + +export interface DemandDrivenResponsePumpOptions { + source: ReadableStream; + onReadStart?: () => void; + onChunk: (chunk: Uint8Array) => void; + onClientCancel?: (reason: unknown) => void; +} + +export interface DemandDrivenResponsePump { + stream: ReadableStream; + completion: Promise; + startDrain: (reason?: unknown) => void; + cancelSource: (reason?: unknown) => void; + errorClient: (error: Error) => void; + getState: () => DemandDrivenResponsePumpState; + wasClientAborted: () => boolean; +} + +const PENDING_CHUNK_DEADLINE_MS = 60_000; + +function toError(error: unknown): Error { + return error instanceof Error ? error : new Error(String(error)); +} + +export function createDemandDrivenResponsePump( + options: DemandDrivenResponsePumpOptions +): DemandDrivenResponsePump { + const reader = options.source.getReader(); + let state: DemandDrivenResponsePumpState = "client-active"; + let clientAborted = false; + let clientController: ReadableStreamDefaultController | null = null; + let pendingChunk: Uint8Array | null = null; + let readInFlight: Promise | null = null; + let drainPromise: Promise | null = null; + let settled = false; + let readerReleased = false; + let pendingChunkDeadlineId: ReturnType | null = null; + let resolveCompletion: (completion: DemandDrivenResponsePumpCompletion) => void = () => {}; + const completion = new Promise((resolve) => { + resolveCompletion = resolve; + }); + + const releaseReader = () => { + if (readerReleased) return; + readerReleased = true; + try { + reader.releaseLock(); + } catch { + // A terminal result must still settle if the platform rejects a late release. + } + }; + + const clearPendingChunkDeadline = () => { + if (!pendingChunkDeadlineId) return; + clearTimeout(pendingChunkDeadlineId); + pendingChunkDeadlineId = null; + }; + + const settle = ( + streamEndedNormally: boolean, + error: Error | null, + sourceCancelReason?: Error + ) => { + if (settled) return; + settled = true; + state = "finalizing"; + pendingChunk = null; + clearPendingChunkDeadline(); + let cancelPromise: Promise | null = null; + const recordSourceCancelFailure = (cancelError: unknown) => { + const normalizedCancelError = toError(cancelError); + if (error && error.cause === undefined) error.cause = normalizedCancelError; + }; + if (sourceCancelReason) { + try { + cancelPromise = reader.cancel(sourceCancelReason); + } catch (cancelError) { + recordSourceCancelFailure( + cancelError instanceof Error ? cancelError : new Error(String(cancelError)) + ); + } + } + releaseReader(); + clientController = null; + state = "closed"; + resolveCompletion({ streamEndedNormally, clientAborted, error }); + void cancelPromise?.then(undefined, recordSourceCancelFailure); + }; + + const finishWithError = (error: unknown) => { + if (settled) return; + const normalized = toError(error); + if (state === "client-active") { + try { + clientController?.error(normalized); + } catch { + // The downstream may have cancelled concurrently. + } + } + settle(false, normalized); + }; + + const finishNormally = () => { + if (settled) return; + if (state === "client-active") { + try { + clientController?.close(); + } catch { + // The downstream may have cancelled concurrently. + } + } + settle(true, null); + }; + + let scheduleDrain = () => {}; + + const startDrain = (_reason?: unknown, markClientAborted = true) => { + if (settled || state === "finalizing" || state === "closed") return; + if (state === "draining") { + scheduleDrain(); + return; + } + if (markClientAborted) clientAborted = true; + state = "draining"; + try { + clientController?.error( + _reason == null ? new Error("Client disconnected") : toError(_reason) + ); + } catch (controllerError) { + if (!(controllerError instanceof TypeError)) throw controllerError; + // The ReadableStream cancel algorithm may have already detached the controller. + } + scheduleDrain(); + }; + + const cancelSource = (reason?: unknown) => { + if (settled) return; + const normalized = reason == null ? new Error("Source cancelled") : toError(reason); + settle(false, normalized, normalized); + }; + + const armPendingChunkDeadline = () => { + clearPendingChunkDeadline(); + pendingChunkDeadlineId = setTimeout(() => { + const error = new DOMException( + `Client response body was not consumed within ${PENDING_CHUNK_DEADLINE_MS}ms`, + "AbortError" + ); + startDrain(error, false); + cancelSource(error); + }, PENDING_CHUNK_DEADLINE_MS); + }; + + const ensureRead = (): Promise => { + if (settled || pendingChunk || readInFlight) { + return readInFlight ?? Promise.resolve(); + } + + let sourceRead: Promise>; + try { + options.onReadStart?.(); + sourceRead = reader.read(); + } catch (error) { + finishWithError(error); + return Promise.resolve(); + } + + const read = sourceRead + .then( + (result) => { + if (settled) return; + if (result.done) { + finishNormally(); + return; + } + + options.onChunk(result.value); + if (settled) return; + pendingChunk = result.value; + armPendingChunkDeadline(); + }, + (error) => finishWithError(error) + ) + .catch((error) => finishWithError(error)) + .finally(() => { + if (readInFlight === read) { + readInFlight = null; + } + if (state === "draining") { + scheduleDrain(); + } + }); + readInFlight = read; + return read; + }; + + scheduleDrain = () => { + if (drainPromise || settled || state !== "draining") return; + + drainPromise = (async () => { + while (!settled && state === "draining") { + if (readInFlight) { + await readInFlight; + continue; + } + if (pendingChunk) { + clearPendingChunkDeadline(); + pendingChunk = null; + continue; + } + await ensureRead(); + } + })().finally(() => { + drainPromise = null; + if (!settled && state === "draining") { + scheduleDrain(); + } + }); + }; + + const stream = new ReadableStream( + { + start(controller) { + clientController = controller; + }, + async pull() { + if (settled || state !== "client-active") return; + + await ensureRead(); + if (settled || state !== "client-active" || !pendingChunk) return; + + const chunk = pendingChunk; + clearPendingChunkDeadline(); + pendingChunk = null; + try { + clientController?.enqueue(chunk); + } catch (error) { + finishWithError(error); + return; + } + + void ensureRead(); + }, + cancel(reason) { + try { + options.onClientCancel?.(reason); + } finally { + startDrain(reason); + } + }, + }, + { highWaterMark: 0 } + ); + + void ensureRead(); + + return { + stream, + completion, + startDrain, + cancelSource, + errorClient(error) { + if (settled || state !== "client-active") return; + state = "draining"; + try { + clientController?.error(error); + } catch { + // The downstream may have cancelled concurrently. + } + scheduleDrain(); + }, + getState: () => state, + wasClientAborted: () => clientAborted, + }; +} diff --git a/src/app/v1/_lib/proxy/error-handler.ts b/src/app/v1/_lib/proxy/error-handler.ts index 13b4e4d6e..8f6aa9216 100644 --- a/src/app/v1/_lib/proxy/error-handler.ts +++ b/src/app/v1/_lib/proxy/error-handler.ts @@ -1,3 +1,4 @@ +import { findSafeDatabaseError } from "@/drizzle/admitted-client"; import { getCachedSystemSettings } from "@/lib/config/system-settings-cache"; import { isClaudeErrorFormat, @@ -8,8 +9,9 @@ import { import { emitProxyLangfuseTrace } from "@/lib/langfuse/emit-proxy-trace"; import { logger } from "@/lib/logger"; import { ProxyStatusTracker } from "@/lib/proxy-status-tracker"; +import { ERROR_CODES, getErrorMessageServer } from "@/lib/utils/error-messages"; import { sanitizeErrorTextForDetail } from "@/lib/utils/upstream-error-detection"; -import { updateMessageRequestDetails, updateMessageRequestDuration } from "@/repository/message"; +import { updateMessageRequestDetailsDurably } from "@/repository/message"; import type { SystemSettings } from "@/types/system-config"; import { deriveClientSafeUpstreamErrorMessage } from "./client-error-message"; import { attachSessionIdToErrorResponse } from "./error-session-id"; @@ -177,6 +179,7 @@ export class ProxyErrorHandler { let rateLimitMetadata: Record | null = null; let settingsResolved = false; let cachedSettings: SystemSettings | null = null; + const databaseError = findSafeDatabaseError(error); const getSettings = async (): Promise => { if (settingsResolved) return cachedSettings; @@ -186,15 +189,21 @@ export class ProxyErrorHandler { return cachedSettings; } catch (settingsError) { settingsResolved = true; + const settingsDatabaseError = findSafeDatabaseError(settingsError); logger.warn("ProxyErrorHandler: failed to load system settings, using defaults", { - error: settingsError instanceof Error ? settingsError.message : String(settingsError), + error: + settingsDatabaseError?.message ?? + (settingsError instanceof Error ? settingsError.message : String(settingsError)), + databaseCode: settingsDatabaseError?.code, + admissionPool: settingsDatabaseError?.pool, + admissionMaxOutstanding: settingsDatabaseError?.maxOutstanding, }); return null; } }; // 优先处理 RateLimitError(新增) - if (isRateLimitError(error)) { + if (!databaseError && isRateLimitError(error)) { clientErrorMessage = error.message; logErrorMessage = error.message; // 使用 helper 函数计算状态码 @@ -241,8 +250,24 @@ export class ProxyErrorHandler { logErrorMessage = "代理请求发生未知错误"; } + if (databaseError) { + // Drizzle wraps the admission cause with SQL text and bound parameters. + // Never let that wrapper cross the public, log, or observability boundary. + try { + const { getLocale } = await import("next-intl/server"); + clientErrorMessage = await getErrorMessageServer( + await getLocale(), + ERROR_CODES.DATABASE_ERROR + ); + } catch { + clientErrorMessage = "An error occurred"; + } + logErrorMessage = databaseError.message; + statusCode = databaseError.kind === "admission" ? 503 : 500; + } + // 后备方案:如果状态码仍是 500,尝试从 provider chain 中提取最后一次实际请求的状态码 - if (statusCode === 500) { + if (!databaseError && statusCode === 500) { const lastRequestStatusCode = ProxyErrorHandler.getLastRequestStatusCode(session); if (lastRequestStatusCode && lastRequestStatusCode !== 200) { statusCode = lastRequestStatusCode; @@ -270,18 +295,22 @@ export class ProxyErrorHandler { responseText, }); // 先发出 trace,再写数据库,避免 DB 持久化失败吞掉本次错误诊断。 - await ProxyErrorHandler.logErrorToDatabase( - session, - logErrorMessage, - finalResponse.status, - null - ); + if (databaseError) { + ProxyErrorHandler.endRequestTracking(session); + } else { + await ProxyErrorHandler.logErrorToDatabase( + session, + logErrorMessage, + finalResponse.status, + null + ); + } return finalResponse; }; // 检测是否有覆写配置(响应体或状态码) // 使用异步版本确保错误规则已加载 - if (error instanceof Error) { + if (error instanceof Error && !databaseError) { const override = await getErrorOverrideAsync(error); if (override) { // 运行时校验覆写状态码范围(400-599),防止数据库脏数据导致 Response 抛 RangeError @@ -512,7 +541,7 @@ export class ProxyErrorHandler { typeof upstreamRequestId === "string" && upstreamRequestId.trim() ? upstreamRequestId.trim() : undefined; - const settings = await getSettings(); + const settings = databaseError ? null : await getSettings(); const finalClientErrorMessage = resolveFinalClientErrorMessage({ error, currentFallbackMessage: clientErrorMessage, @@ -623,7 +652,6 @@ export class ProxyErrorHandler { } const duration = Date.now() - session.startTime; - await updateMessageRequestDuration(session.messageContext.id, duration); // 如果是限流错误,将元数据附加到错误消息中 let finalErrorMessage = errorMessage; @@ -632,7 +660,8 @@ export class ProxyErrorHandler { } // 保存错误信息和决策链 - await updateMessageRequestDetails(session.messageContext.id, { + await updateMessageRequestDetailsDurably(session.messageContext.id, { + durationMs: duration, errorMessage: finalErrorMessage, providerChain: session.getProviderChain(), statusCode: statusCode, @@ -643,6 +672,11 @@ export class ProxyErrorHandler { }); // 记录请求结束 + ProxyErrorHandler.endRequestTracking(session); + } + + private static endRequestTracking(session: ProxySession): void { + if (!session.messageContext) return; const tracker = ProxyStatusTracker.getInstance(); tracker.endRequest(session.messageContext.user.id, session.messageContext.id); } diff --git a/src/app/v1/_lib/proxy/errors.ts b/src/app/v1/_lib/proxy/errors.ts index 91e220263..647b4bfbd 100644 --- a/src/app/v1/_lib/proxy/errors.ts +++ b/src/app/v1/_lib/proxy/errors.ts @@ -6,12 +6,14 @@ * 2. 智能截断:JSON 完整保存,文本限制 500 字符 * 3. 可读性优先:纯文本格式化,便于排查问题 */ +import { isDbPoolAdmissionError } from "@/drizzle/admitted-client"; import { getEnvConfig } from "@/lib/config/env.schema"; import { type ErrorDetectionResult, errorRuleDetector } from "@/lib/error-rule-detector"; import { redactJsonString } from "@/lib/utils/message-redaction"; import { sanitizeErrorTextForDetail } from "@/lib/utils/upstream-error-detection"; import type { ErrorOverrideResponse } from "@/repository/error-rules"; import type { ProviderChainItem } from "@/types/message"; +import { RESERVED_INTERNAL_HEADERS } from "../responses-ws/internal-secret"; import type { ProxySession } from "./session"; /** Marker message for the synthetic terminal error emitted when every provider fails. */ @@ -554,6 +556,7 @@ export enum ErrorCategory { CLIENT_ABORT, // 客户端主动中断 → 不计入熔断器 + 不重试 + 直接返回 NON_RETRYABLE_CLIENT_ERROR, // 客户端输入错误(Prompt 超限、内容过滤、PDF 限制、Thinking 格式、参数缺失/额外参数、非法请求)→ 不计入熔断器 + 不重试 + 直接返回 RESOURCE_NOT_FOUND, // 上游 404 错误 → 不计入熔断器 + 直接切换供应商 + LOCAL_OVERLOAD, // 本地数据库等 admission 过载 → 不计入任何熔断器 + 不重试/切换供应商 } /** @@ -925,7 +928,7 @@ export function isEmptyResponseError(error: unknown): error is EmptyResponseErro * 此函数会确保错误规则已加载后再进行检测 * * @param error - 捕获的错误对象 - * @returns 错误分类(CLIENT_ABORT、NON_RETRYABLE_CLIENT_ERROR、PROVIDER_ERROR 或 SYSTEM_ERROR) + * @returns 错误分类(CLIENT_ABORT、LOCAL_OVERLOAD、NON_RETRYABLE_CLIENT_ERROR、PROVIDER_ERROR 或 SYSTEM_ERROR) */ export async function categorizeErrorAsync(error: Error): Promise { // 优先级 1: 客户端中断检测(优先级最高)- 使用统一的精确检测函数 @@ -933,6 +936,12 @@ export async function categorizeErrorAsync(error: Error): Promise return ErrorCategory.CLIENT_ABORT; // 客户端主动中断 } + // 优先级 1.25: 本地 DB admission 过载。Drizzle 会把底层错误包在 cause 中, + // 必须在网络/规则分类前识别,避免重试上游或惩罚 Provider/endpoint circuit。 + if (isDbPoolAdmissionError(error)) { + return ErrorCategory.LOCAL_OVERLOAD; + } + // 优先级 1.5: Native transport errors — must not be matched by error rules // These are always SYSTEM_ERROR regardless of message content if (isTransportError(error)) { @@ -1118,6 +1127,13 @@ const SENSITIVE_HEADERS = new Set([ "cookie", "set-cookie", ]); +const RESERVED_INTERNAL_HEADER_SET = new Set( + RESERVED_INTERNAL_HEADERS.map((header) => header.toLowerCase()) +); +function isReservedInternalHeader(name: string): boolean { + const lowerName = name.toLowerCase(); + return lowerName.startsWith("x-cch-") || RESERVED_INTERNAL_HEADER_SET.has(lowerName); +} const SENSITIVE_URL_PARAMS = new Set([ "key", @@ -1199,6 +1215,7 @@ export function sanitizeHeaders(headers: Headers | string): string { const collected: string[] = []; headers.forEach((value, key) => { const lowerKey = key.toLowerCase(); + if (isReservedInternalHeader(lowerKey)) return; if (SENSITIVE_HEADERS.has(lowerKey)) { const maskedValue = lowerKey === "authorization" ? maskAuthorizationValue(value) : maskSensitiveValue(value); @@ -1226,6 +1243,7 @@ export function sanitizeHeaders(headers: Headers | string): string { if (!name) return line; const lowerName = name.toLowerCase(); + if (isReservedInternalHeader(lowerName)) return null; if (!SENSITIVE_HEADERS.has(lowerName)) return line; const maskedValue = @@ -1233,7 +1251,7 @@ export function sanitizeHeaders(headers: Headers | string): string { return `${name}: ${maskedValue}`; }); - return sanitizedLines.join("\n"); + return sanitizedLines.filter((line): line is string => line !== null).join("\n"); } /** diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 4da548a5f..cb5dd36f8 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -4,6 +4,7 @@ import { pipeline as streamPipeline } from "node:stream"; import { createGunzip, constants as zlibConstants } from "node:zlib"; import type { Dispatcher } from "undici"; import { request as undiciRequest } from "undici"; +import { findDbPoolAdmissionError, findSafeDatabaseError } from "@/drizzle/admitted-client"; import { applyAnthropicProviderOverridesWithAudit } from "@/lib/anthropic/provider-overrides"; import { getCircuitState, @@ -62,7 +63,6 @@ import { buildProxyUrl } from "../url"; import { rectifyBillingHeader } from "./billing-header-rectifier"; import { bindClientAbortListener } from "./client-abort-listener"; import { deriveClientSafeUpstreamErrorMessage } from "./client-error-message"; -import { combineAbortSignals } from "./combine-abort-signals"; import { isStandardProxyEndpointPath } from "./endpoint-family-catalog"; import { resolveEndpointPolicy, shouldEnforceStrictEndpointPoolPolicy } from "./endpoint-policy"; import { @@ -1733,10 +1733,15 @@ export class ProxyForwarder { // ⭐ 1. 分类错误(供应商错误 vs 系统错误 vs 客户端中断) // 使用异步版本确保错误规则已加载 let errorCategory = await categorizeErrorAsync(lastError); + const databaseError = findSafeDatabaseError(lastError); + if (databaseError) { + errorCategory = ErrorCategory.LOCAL_OVERLOAD; + } const errorMessage = - lastError instanceof ProxyError + databaseError?.message ?? + (lastError instanceof ProxyError ? lastError.getDetailedErrorMessage() - : lastError.message; + : lastError.message); const isTimeoutError = lastError instanceof ProxyError && lastError.statusCode === 524; @@ -1746,7 +1751,7 @@ export class ProxyForwarder { ); } - if (activeEndpoint.endpointId != null) { + if (!databaseError && activeEndpoint.endpointId != null) { if (isTimeoutError || errorCategory === ErrorCategory.SYSTEM_ERROR) { await recordEndpointFailure(activeEndpoint.endpointId, lastError); } @@ -1761,7 +1766,7 @@ export class ProxyForwarder { totalProvidersAttempted, }); - await ProxyForwarder.clearSessionProviderBinding(session); + await ProxyForwarder.clearSessionProviderBinding(session, currentProvider.id); // 记录到决策链(标记为客户端中断) session.addProviderToChain(currentProvider, { @@ -1783,6 +1788,39 @@ export class ProxyForwarder { throw lastError; } + if (databaseError) { + const admission = findDbPoolAdmissionError(lastError); + logger.warn("ProxyForwarder: Local database operation failed", { + providerId: currentProvider.id, + providerName: currentProvider.name, + endpointId: activeEndpoint.endpointId, + pool: admission?.pool, + maxOutstanding: admission?.maxOutstanding, + attemptNumber: attemptCount, + }); + + session.addProviderToChain(currentProvider, { + ...endpointAudit, + reason: "system_error", + circuitState: getCircuitState(currentProvider.id), + attemptNumber: attemptCount, + errorMessage: databaseError.message, + errorDetails: { + system: { + errorType: + databaseError.kind === "admission" ? "DbPoolAdmissionError" : "DatabaseError", + errorName: + databaseError.kind === "admission" ? "DbPoolAdmissionError" : "DatabaseError", + errorMessage: databaseError.message, + errorCode: databaseError.code, + }, + request: buildRequestDetails(session), + }, + }); + + throw lastError; + } + // 2.5 Reactive rectifier:命中后对同供应商“整流 + 重试一次” const reactiveRectifierResult = await tryApplyReactiveRectifier({ provider: currentProvider, @@ -2337,8 +2375,11 @@ export class ProxyForwarder { }); } - // ⭐ 不暴露供应商详情,仅返回简单错误 - await ProxyForwarder.clearSessionProviderBinding(session); + // ⭐ 不暴露供应商详情,仅返回简单错误。CAS 清理本次请求实际尝试过的 + // 所有 provider,避免从 A 切换到 B 后只比较 B 而遗留 A 的 stale binding。 + const attemptedProviderIds = new Set(failedProviderIds); + if (session.provider?.id != null) attemptedProviderIds.add(session.provider.id); + await ProxyForwarder.clearSessionProviderBindings(session, attemptedProviderIds); throw ProxyForwarder.buildAllProvidersUnavailableError(lastError); // Service Unavailable } @@ -2975,23 +3016,31 @@ export class ProxyForwarder { } // 2. 组合双路信号:response + client - const signals = [responseController.signal]; - if (session.clientAbortSignal) { - signals.push(session.clientAbortSignal); - } - - // 优先 Node 20.3+ 原生 AbortSignal.any(V8 内部管理 listener,无需手动 cleanup); - // Next.js standalone 覆盖全局时 fallback 到 polyfill,由调用方在请求结束时调用 - // cleanupCombinedSignal 解绑源信号上的 listener,避免持有 session/请求体闭包。 - const { signal: combinedSignal, cleanup: cleanupCombinedSignal } = combineAbortSignals(signals); + const transportController = new AbortController(); + const abortTransportFrom = (source: AbortSignal) => { + if (!transportController.signal.aborted) { + transportController.abort(source.reason); + } + }; + const cleanupResponseTransportSignal = bindClientAbortListener(responseController.signal, () => + abortTransportFrom(responseController.signal) + ); + const cleanupClientTransportSignal = bindClientAbortListener(session.clientAbortSignal, () => { + const clientSignal = session.clientAbortSignal; + if (clientSignal) abortTransportFrom(clientSignal); + }); + const cleanupCombinedSignal = () => { + cleanupResponseTransportSignal(); + cleanupClientTransportSignal(); + }; logger.debug("ProxyForwarder: Combined abort signals", { - signalCount: signals.length, + signalCount: session.clientAbortSignal ? 2 : 1, }); const init: UndiciFetchOptions = { method: session.method, headers: processedHeaders, - signal: combinedSignal, // 使用组合信号 + signal: transportController.signal, // 使用组合信号 ...(requestBody ? { body: requestBody } : {}), }; @@ -3077,7 +3126,7 @@ export class ProxyForwarder { body: requestBodyJson, sessionId: getResponsesWsSessionId(session.headers), endpointId: responsesWsEndpointId, - abortSignal: combinedSignal, + abortSignal: transportController.signal, }); if ("response" in wsResult) { @@ -3669,6 +3718,11 @@ export class ProxyForwarder { } } + // Successful responses transfer abort ownership to ResponseHandler. + // Error bodies are consumed here, so their client listener must remain + // attached until fromUpstreamResponse() finishes or is aborted. + cleanupClientTransportSignal(); + // 将响应超时清理函数和 controller 引用附加到 session,供 response-handler 使用 // response-handler 会在读到首字节(流式)或完整响应(非流式)后调用此函数 const sessionWithTimeout = session as ProxySession & { @@ -3693,7 +3747,7 @@ export class ProxyForwarder { // Attach agent release callback for in-flight reference counting. // response-handler must call this in its finally block after the stream is fully consumed. - // 同时复用此回调作为 combineAbortSignals polyfill 的 cleanup 入口:response-handler 已经 + // 同时复用此回调作为 transport signal 的 cleanup 入口:response-handler 已经 // 保证在请求结束时(成功/异常)幂等地调用 releaseAgent,把 cleanup 合并到这里就不必再 // 改造 response-handler 的所有 finally 调用点。两个动作互不影响,cleanup 内部自带 cleaned // 标志,重复调用安全。 @@ -3808,7 +3862,9 @@ export class ProxyForwarder { const settleFailure = async (error: Error) => { if (settled) return; settled = true; - await ProxyForwarder.clearSessionProviderBinding(session); + const attemptedProviderIds = new Set(launchedProviderIds); + if (session.provider?.id != null) attemptedProviderIds.add(session.provider.id); + await ProxyForwarder.clearSessionProviderBindings(session, attemptedProviderIds); resolveResult?.({ error }); }; @@ -3852,6 +3908,7 @@ export class ProxyForwarder { const reader = attempt.reader; const response = attempt.response; const messageRequestId = session.messageContext?.id; + const messageRequestCreatedAtMs = session.messageContext?.createdAt.getTime(); if (!reader || !response || messageRequestId == null) { // 无可读响应或无请求行可归属 -> 无法计费,直接释放资源。 const cancel = reader?.cancel("hedge_loser_no_billing"); @@ -3923,6 +3980,7 @@ export class ProxyForwarder { await finalizeHedgeLoserBilling({ messageRequestId, + messageRequestCreatedAtMs: messageRequestCreatedAtMs ?? Date.now(), loserSession: attempt.session, provider: attempt.provider, attemptNumber: attempt.sequence, @@ -4226,8 +4284,10 @@ export class ProxyForwarder { let errorCategory = await categorizeErrorAsync(error); lastErrorCategory = errorCategory; const statusCode = error instanceof ProxyError ? error.statusCode : undefined; + const databaseError = findSafeDatabaseError(error); const errorMessage = - error instanceof ProxyError ? error.getDetailedErrorMessage() : error.message; + databaseError?.message ?? + (error instanceof ProxyError ? error.getDetailedErrorMessage() : error.message); let matchedRule: MatchedRuleDetails | undefined; let matchedRuleLogContext: Record = {}; @@ -4263,6 +4323,41 @@ export class ProxyForwarder { return; } + if (errorCategory === ErrorCategory.LOCAL_OVERLOAD) { + const admission = findDbPoolAdmissionError(error); + const safeAdmissionMessage = admission?.message ?? "Database pool admission exceeded"; + logger.warn("ProxyForwarder: Local database admission rejected during hedge", { + providerId: attempt.provider.id, + providerName: attempt.provider.name, + endpointId: attempt.endpointAudit.endpointId, + pool: admission?.pool, + maxOutstanding: admission?.maxOutstanding, + participantSequence: attempt.sequence, + attemptNumber: attempt.requestAttemptCount, + }); + + session.addProviderToChain(attempt.provider, { + ...attempt.endpointAudit, + reason: "system_error", + attemptNumber: attempt.sequence, + errorMessage: safeAdmissionMessage, + circuitState: getCircuitState(attempt.provider.id), + errorDetails: { + system: { + errorType: "DbPoolAdmissionError", + errorName: "DbPoolAdmissionError", + errorMessage: safeAdmissionMessage, + errorCode: admission?.code, + }, + request: buildRequestDetails(session), + }, + modelRedirect: getAttemptModelRedirect(attempt), + }); + abortAllAttempts(undefined, "database_pool_overload"); + await settleFailure(error); + return; + } + const reactiveRectifierResult = await tryApplyReactiveRectifier({ provider: attempt.provider, requestSession: attempt.session, @@ -4894,9 +4989,21 @@ export class ProxyForwarder { targetState.releaseAgent = sourceRuntime.releaseAgent; } - private static async clearSessionProviderBinding(session: ProxySession): Promise { + private static async clearSessionProviderBinding( + session: ProxySession, + expectedProviderId: number | null + ): Promise { if (!session.sessionId) return; - await SessionManager.clearSessionProvider(session.sessionId); + await SessionManager.clearSessionProvider(session.sessionId, expectedProviderId); + } + + private static async clearSessionProviderBindings( + session: ProxySession, + expectedProviderIds: Iterable + ): Promise { + for (const providerId of new Set(expectedProviderIds)) { + await ProxyForwarder.clearSessionProviderBinding(session, providerId); + } } private static markProviderFailed( diff --git a/src/app/v1/_lib/proxy/node-stream-to-web.test.ts b/src/app/v1/_lib/proxy/node-stream-to-web.test.ts index 5f52a068f..6ef589232 100644 --- a/src/app/v1/_lib/proxy/node-stream-to-web.test.ts +++ b/src/app/v1/_lib/proxy/node-stream-to-web.test.ts @@ -23,6 +23,139 @@ async function readAll(reader: ReadableStreamDefaultReader): Promise } describe("nodeStreamToWebStreamSafe", () => { + it("does not drain an unread source past the stream high-water marks", async () => { + const totalChunks = 32; + let producedChunks = 0; + + const node = new Readable({ + highWaterMark: 1, + read() { + if (producedChunks >= totalChunks) { + this.push(null); + return; + } + + producedChunks++; + this.push(Buffer.alloc(64 * 1024, producedChunks)); + }, + }); + + const web = nodeStreamToWebStreamSafe(node, 1, "test"); + + // Allow all currently scheduled Node stream work to run. A backpressured + // adapter should fill only the Node/Web high-water marks, not reach EOF. + await new Promise((resolve) => setImmediate(resolve)); + await new Promise((resolve) => setImmediate(resolve)); + + expect(producedChunks).toBeLessThanOrEqual(3); + + await web.cancel(); + }); + + it("uses a byte budget so small chunks do not pause after one enqueue", async () => { + const chunkSize = 256; + const highWaterMark = 64 * 1024; + const webHighWaterMark = highWaterMark * 2; + let producedChunks = 0; + + const node = new Readable({ + highWaterMark, + read() { + producedChunks++; + this.push(Buffer.alloc(chunkSize, producedChunks)); + }, + }); + + const web = nodeStreamToWebStreamSafe(node, 1, "test"); + await new Promise((resolve) => setImmediate(resolve)); + await new Promise((resolve) => setImmediate(resolve)); + + const producedBytes = producedChunks * chunkSize; + expect(producedBytes).toBeGreaterThan(webHighWaterMark + chunkSize * 8); + expect(producedBytes).toBeLessThanOrEqual(webHighWaterMark + highWaterMark + chunkSize * 2); + + await web.cancel(); + }); + + it("closes immediately when the source ended before conversion", async () => { + const node = new Readable({ + read() { + this.push(null); + }, + }); + node.resume(); + await new Promise((resolve) => node.once("end", resolve)); + + const web = nodeStreamToWebStreamSafe(node, 1, "test"); + const reader = web.getReader(); + const result = await Promise.race([ + reader.read(), + new Promise<"pending">((resolve) => setImmediate(() => resolve("pending"))), + ]); + + expect(result).toEqual({ done: true, value: undefined }); + }); + + it("rejects with the original error when the source errored before conversion", async () => { + const node = new Readable({ + read() { + // no-op + }, + }); + const boom = new Error("preexisting-upstream-error"); + node.once("error", () => {}); + node.destroy(boom); + await new Promise((resolve) => node.once("close", resolve)); + + const web = nodeStreamToWebStreamSafe(node, 1, "test"); + const reader = web.getReader(); + + await expect(reader.read()).rejects.toBe(boom); + expect(node.listenerCount("error")).toBe(0); + expect(node.listenerCount("close")).toBe(0); + }); + + it("rejects when the source was destroyed before conversion without reaching EOF", async () => { + const node = new Readable({ + read() { + // no-op + }, + }); + node.destroy(); + await new Promise((resolve) => node.once("close", resolve)); + + const web = nodeStreamToWebStreamSafe(node, 1, "test"); + const reader = web.getReader(); + + await expect(reader.read()).rejects.toThrow("closed before end"); + }); + + it("protects a delayed destroy error when destruction started before conversion", async () => { + const lateDestroyError = new Error("pre-aborted-late-destroy"); + const node = new Readable({ + read() { + // no-op + }, + destroy(_error, callback) { + setTimeout(() => callback(lateDestroyError), 30); + }, + }); + const uncaughtSpy = vi.fn(); + process.once("uncaughtException", uncaughtSpy); + + node.destroy(); + const web = nodeStreamToWebStreamSafe(node, 1, "test"); + const reader = web.getReader(); + + await expect(reader.read()).rejects.toThrow("closed before end"); + await new Promise((resolve) => setTimeout(resolve, 50)); + + process.removeListener("uncaughtException", uncaughtSpy); + expect(uncaughtSpy).not.toHaveBeenCalled(); + expect(node.listenerCount("error")).toBe(0); + expect(node.listenerCount("close")).toBe(0); + }); + it("forwards chunks then closes when the source ends normally", async () => { const node = Readable.from([Buffer.from("hello "), Buffer.from("world")]); @@ -62,6 +195,20 @@ describe("nodeStreamToWebStreamSafe", () => { expect(node.listenerCount("error")).toBe(0); }); + it("rejects when the source closes after conversion without reaching EOF", async () => { + const node = new Readable({ + read() { + // no-op, destroy after the adapter installs its listeners + }, + }); + + const web = nodeStreamToWebStreamSafe(node, 1, "test"); + const reader = web.getReader(); + node.destroy(); + + await expect(reader.read()).rejects.toThrow("closed before end"); + }); + it("destroys the source and detaches listeners on cancel(), and ignores subsequent events", async () => { const node = new Readable({ read() { @@ -140,15 +287,101 @@ describe("nodeStreamToWebStreamSafe", () => { }); node.destroy(); const destroySpy = vi.spyOn(node, "destroy"); + node.once("error", () => {}); const web = nodeStreamToWebStreamSafe(node, 1, "test"); const reader = web.getReader(); - await reader.cancel("client gone"); + await reader.cancel("client gone").catch(() => {}); // Wrapper must short-circuit when nodeStream.destroyed is true expect(destroySpy).not.toHaveBeenCalled(); }); + it("guards a delayed error when external destroy races with downstream cancel", async () => { + const lateDestroyError = new Error("external-destroy-late-error"); + const node = new Readable({ + read() { + // no-op + }, + destroy(_error, callback) { + setTimeout(() => callback(lateDestroyError), 30); + }, + }); + const web = nodeStreamToWebStreamSafe(node, 1, "test"); + const reader = web.getReader(); + const uncaughtSpy = vi.fn(); + process.once("uncaughtException", uncaughtSpy); + + node.destroy(); + expect(node.destroyed).toBe(true); + expect(node.errored).toBeNull(); + expect(node.closed).toBe(false); + + await reader.cancel("client gone"); + const adapterProtectedPendingDestroy = node.listenerCount("error") > 0; + if (!adapterProtectedPendingDestroy) { + // Keep the red test from leaking the expected late error into Vitest. + node.once("error", () => {}); + } + await new Promise((resolve) => setTimeout(resolve, 50)); + + process.removeListener("uncaughtException", uncaughtSpy); + expect(adapterProtectedPendingDestroy).toBe(true); + expect(uncaughtSpy).not.toHaveBeenCalled(); + expect(node.listenerCount("error")).toBe(0); + expect(node.listenerCount("close")).toBe(0); + }); + + it("protects a queued destroy(error) event when cancel races in the same tick", async () => { + const node = new Readable({ + read() { + // no-op + }, + }); + const web = nodeStreamToWebStreamSafe(node, 1, "test"); + const reader = web.getReader(); + + node.destroy(new Error("upstream-race")); + const cancelPromise = reader.cancel("client gone"); + + const adapterProtectedQueuedError = node.listenerCount("error") > 0; + if (!adapterProtectedQueuedError) { + // Keep the current broken implementation from surfacing an uncaught + // exception after the assertion has captured the missing protection. + node.once("error", () => {}); + } + + await cancelPromise; + await new Promise((resolve) => setImmediate(resolve)); + + expect(adapterProtectedQueuedError).toBe(true); + expect(node.listenerCount("error")).toBe(0); + expect(node.listenerCount("close")).toBe(0); + }); + + it("protects a delayed asynchronous destroy(error) event after cancel", async () => { + const lateDestroyError = new Error("late-destroy"); + const node = new Readable({ + read() { + // no-op + }, + destroy(_error, callback) { + setTimeout(() => callback(lateDestroyError), 30); + }, + }); + const uncaughtSpy = vi.fn(); + process.once("uncaughtException", uncaughtSpy); + + const web = nodeStreamToWebStreamSafe(node, 1, "test"); + await web.cancel(new Error("client gone")); + await new Promise((resolve) => setTimeout(resolve, 50)); + + process.removeListener("uncaughtException", uncaughtSpy); + expect(uncaughtSpy).not.toHaveBeenCalled(); + expect(node.listenerCount("error")).toBe(0); + expect(node.listenerCount("close")).toBe(0); + }); + it("treats back-to-back end + close as a single close on the web stream", async () => { const node = new Readable({ read() { diff --git a/src/app/v1/_lib/proxy/node-stream-to-web.ts b/src/app/v1/_lib/proxy/node-stream-to-web.ts index a2bf9ff94..82b3bbff8 100644 --- a/src/app/v1/_lib/proxy/node-stream-to-web.ts +++ b/src/app/v1/_lib/proxy/node-stream-to-web.ts @@ -37,107 +37,182 @@ export function nodeStreamToWebStreamSafe( onError = null; }; - return new ReadableStream({ - start(controller) { - logger.debug("ProxyForwarder: Starting Node-to-Web stream conversion", { - providerId, - providerName, - }); - - onData = (chunk: Buffer | Uint8Array) => { - if (settled) return; - chunkCount++; - totalBytes += chunk.length; - try { - const buf = chunk instanceof Uint8Array ? chunk : new Uint8Array(chunk); - controller.enqueue(buf); - } catch { - // controller 已关闭/出错时忽略 - } - }; - nodeStream.on("data", onData); + const installPendingDestroyErrorGuard = (stream: Readable) => { + let cleanupTimeout: NodeJS.Timeout | null = null; + const cleanup = () => { + stream.removeListener("error", swallow); + stream.removeListener("close", cleanup); + if (cleanupTimeout) { + clearTimeout(cleanupTimeout); + cleanupTimeout = null; + } + }; + const swallow = () => { + // ignore: the Web stream has already settled with the same terminal state + cleanup(); + }; - onEnd = () => { - if (settled) return; - settled = true; - logger.debug("ProxyForwarder: Node stream ended normally", { + stream.once("error", swallow); + stream.once("close", cleanup); + cleanupTimeout = setTimeout(cleanup, 60_000); + cleanupTimeout.unref?.(); + }; + + return new ReadableStream( + { + start(controller) { + logger.debug("ProxyForwarder: Starting Node-to-Web stream conversion", { providerId, providerName, - chunkCount, - totalBytes, }); - detach(nodeStream); - try { - controller.close(); - } catch { - // ignore + + // Web ReadableStream 通过 pull() 表达下游需求。先保持 Node 流暂停, + // 并在注册会启用 flowing mode 的 data 监听器前安装所有终态监听器。 + nodeStream.pause(); + + onEnd = () => { + if (settled) return; + settled = true; + logger.debug("ProxyForwarder: Node stream ended normally", { + providerId, + providerName, + chunkCount, + totalBytes, + }); + detach(nodeStream); + try { + controller.close(); + } catch { + // ignore + } + }; + nodeStream.on("end", onEnd); + + onClose = () => { + if (settled) return; + if (!nodeStream.readableEnded) { + onError?.(new Error("Upstream stream closed before end")); + return; + } + settled = true; + logger.debug("ProxyForwarder: Node stream closed", { + providerId, + providerName, + chunkCount, + totalBytes, + }); + detach(nodeStream); + try { + controller.close(); + } catch { + // ignore + } + }; + nodeStream.on("close", onClose); + + onError = (err: Error) => { + if (settled) return; + settled = true; + logger.warn("ProxyForwarder: Upstream stream error (signaling downstream)", { + providerId, + providerName, + error: err.message, + errorName: err.name, + }); + detach(nodeStream); + try { + controller.error(err); + } catch { + // ignore + } + }; + nodeStream.on("error", onError); + + onData = (chunk: Buffer | Uint8Array) => { + if (settled) return; + chunkCount++; + totalBytes += chunk.length; + try { + const buf = chunk instanceof Uint8Array ? chunk : new Uint8Array(chunk); + controller.enqueue(buf); + if (controller.desiredSize !== null && controller.desiredSize <= 0) { + nodeStream.pause(); + } + } catch { + // controller 已关闭/出错时忽略 + } + }; + nodeStream.on("data", onData); + + const preexistingError = nodeStream.errored; + if (preexistingError) { + // Before close, destroy(error) may have recorded the error while its + // event is still queued. Once close has fired, no delayed destroy event + // remains and retaining the bounded guard would only hold the stream. + if (!nodeStream.closed) { + installPendingDestroyErrorGuard(nodeStream); + } + onError(preexistingError); + return; } - }; - nodeStream.on("end", onEnd); - onClose = () => { - if (settled) return; - settled = true; - logger.debug("ProxyForwarder: Node stream closed", { - providerId, - providerName, - chunkCount, - totalBytes, - }); - detach(nodeStream); - try { - controller.close(); - } catch { - // ignore + if (nodeStream.readableAborted || (nodeStream.closed && !nodeStream.readableEnded)) { + if (!nodeStream.closed) { + // destroy() marks the stream aborted before an asynchronous _destroy + // callback can emit its error. Keep a bounded listener until that + // callback reaches error/close so the settled Web stream cannot turn + // a request-local failure into an uncaught process error. + installPendingDestroyErrorGuard(nodeStream); + } + onError(new Error("Upstream stream closed before end")); + return; + } + + // end/close 可能在包装前已经发出;监听器安装完成后复查终态, + // 避免 Web reader 永久等待一个不会再次触发的事件。 + if (nodeStream.readableEnded) { + onEnd(); + } + }, + + pull() { + if (!settled && !nodeStream.destroyed) { + nodeStream.resume(); } - }; - nodeStream.on("close", onClose); + }, - onError = (err: Error) => { + cancel(reason) { + // 重复 cancel 应是 no-op:避免重复 detach、重复注册 swallow 监听 if (settled) return; settled = true; - logger.warn("ProxyForwarder: Upstream stream error (signaling downstream)", { - providerId, - providerName, - error: err.message, - errorName: err.name, - }); detach(nodeStream); + + if (nodeStream.destroyed) { + // destroy(error) may queue an error even after closed becomes true. + // External destroy() can also set destroyed before an asynchronous + // _destroy callback supplies its error, so guard while close is pending. + if (nodeStream.errored || !nodeStream.closed) { + installPendingDestroyErrorGuard(nodeStream); + } + return; + } + + installPendingDestroyErrorGuard(nodeStream); + try { - controller.error(err); + nodeStream.destroy( + reason instanceof Error ? reason : reason ? new Error(String(reason)) : undefined + ); } catch { // ignore } - }; - nodeStream.on("error", onError); - }, - - cancel(reason) { - // 重复 cancel 应是 no-op:避免重复 detach、重复注册 swallow 监听 - if (settled) return; - settled = true; - detach(nodeStream); - if (nodeStream.destroyed) return; - - // destroy(reason) 在 reason 为 Error 时会 re-emit "error",而我们已经 - // detach 了错误监听。注册一次 swallow 监听吞掉它,避免触发 uncaughtException。 - // 但 destroy() 不带 reason / 带非 Error reason 时 不会 emit error,此时 - // once("error") 会成为常驻泄露 —— 用 once("close") 兜底清理。 - const swallow = () => { - // ignore: web 流已 cancel,下游没有 reader 关心了 - }; - nodeStream.once("error", swallow); - nodeStream.once("close", () => { - nodeStream.removeListener("error", swallow); - }); - - try { - nodeStream.destroy( - reason instanceof Error ? reason : reason ? new Error(String(reason)) : undefined - ); - } catch { - // ignore - } + }, }, - }); + { + highWaterMark: Math.max(1, nodeStream.readableHighWaterMark * 2), + size(chunk) { + return chunk.byteLength; + }, + } + ); } diff --git a/src/app/v1/_lib/proxy/provider-selector.ts b/src/app/v1/_lib/proxy/provider-selector.ts index 8b9fde2b9..4dc42b9f0 100644 --- a/src/app/v1/_lib/proxy/provider-selector.ts +++ b/src/app/v1/_lib/proxy/provider-selector.ts @@ -484,7 +484,7 @@ export class ProxyProviderResolver { sessionId: session.sessionId, providerId, }); - await SessionManager.clearSessionProvider(session.sessionId); + await SessionManager.clearSessionProvider(session.sessionId, providerId); return null; } @@ -494,7 +494,7 @@ export class ProxyProviderResolver { providerId: provider.id, providerName: provider.name, }); - await SessionManager.clearSessionProvider(session.sessionId); + await SessionManager.clearSessionProvider(session.sessionId, providerId); return null; } @@ -508,7 +508,7 @@ export class ProxyProviderResolver { activeTimeEnd: provider.activeTimeEnd, timezone: systemTimezone, }); - await SessionManager.clearSessionProvider(session.sessionId); + await SessionManager.clearSessionProvider(session.sessionId, providerId); return null; } @@ -549,7 +549,7 @@ export class ProxyProviderResolver { providerType: provider.providerType, originalFormat: session.originalFormat, }); - await SessionManager.clearSessionProvider(session.sessionId); + await SessionManager.clearSessionProvider(session.sessionId, providerId); return null; } @@ -568,7 +568,7 @@ export class ProxyProviderResolver { // 清除过时绑定,避免 SET NX 死锁 // 当 session 内请求模型发生变化时,旧绑定已无意义, // 清除后新的成功请求可通过 SET NX 重新绑定匹配的 provider - await SessionManager.clearSessionProvider(session.sessionId); + await SessionManager.clearSessionProvider(session.sessionId, providerId); logger.info("ProviderSelector: Cleared stale provider binding (model mismatch)", { sessionId: session.sessionId, staleProviderId: provider.id, @@ -624,7 +624,7 @@ export class ProxyProviderResolver { ], }, }); - await SessionManager.clearSessionProvider(session.sessionId); + await SessionManager.clearSessionProvider(session.sessionId, providerId); return null; } diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index 1a8406b6d..d633c8cbd 100644 --- a/src/app/v1/_lib/proxy/response-handler.ts +++ b/src/app/v1/_lib/proxy/response-handler.ts @@ -3,6 +3,7 @@ import { resolveAnthropicStreamActualResponseModel, } from "@/app/v1/_lib/proxy/anthropic-actual-response-model"; import { ResponseFixer } from "@/app/v1/_lib/proxy/response-fixer"; +import { findSafeDatabaseError } from "@/drizzle/admitted-client"; import { AsyncTaskManager } from "@/lib/async-task-manager"; import { getEnvConfig } from "@/lib/config/env.schema"; import { getCachedSystemSettings } from "@/lib/config/system-settings-cache"; @@ -37,8 +38,8 @@ import { import { addMessageRequestHedgeLoserCost, updateMessageRequestCostWithBreakdown, - updateMessageRequestDetails, - updateMessageRequestDuration, + updateMessageRequestDetailsDurably, + updateMessageRequestDetailsIfUnfinalized, updateMessageRequestWinnerCost, } from "@/repository/message"; import type { HedgeLoserBilling, StoredCostBreakdown } from "@/types/cost-breakdown"; @@ -49,6 +50,10 @@ import { GeminiAdapter } from "../gemini/adapter"; import type { GeminiResponse } from "../gemini/types"; import { extractActualResponseModelForProvider } from "./actual-response-model"; import { bindClientAbortListener } from "./client-abort-listener"; +import { + createDemandDrivenResponsePump, + type DemandDrivenResponsePump, +} from "./demand-driven-response-pump"; import { isClientAbortError, isTransportError } from "./errors"; import type { ProxySession } from "./session"; import { @@ -81,16 +86,259 @@ function resolveNonStreamTaskStaleTimeoutMs(provider: Provider): number { : Number.POSITIVE_INFINITY; } -function resolveStreamTaskStaleTimeoutMs(provider: Provider): number { - if (provider.streamingIdleTimeoutMs <= 0) { - return Number.POSITIVE_INFINITY; +function resolveStreamTaskStaleTimeoutMs(): number { + // Streaming liveness is owned by first-byte, Provider-idle, and client-drain + // timers. The generic watchdog cannot distinguish a stalled Provider from a + // healthy pending chunk that is deliberately waiting for downstream demand. + return Number.POSITIVE_INFINITY; +} + +const STREAM_FINALIZATION_MAX_MS = 120_000; +const STREAM_FAILURE_PERSISTENCE_MAX_MS = 5_000; +const NON_STREAM_TERMINAL_PERSISTENCE_ERROR = Symbol("non_stream_terminal_persistence_error"); + +type MessageRequestTerminalDetails = Parameters[1]; +type NonStreamTerminalPersistenceError = Error & { + [NON_STREAM_TERMINAL_PERSISTENCE_ERROR]: true; +}; + +function markNonStreamTerminalPersistenceError(error: unknown): NonStreamTerminalPersistenceError { + const markedError = + error instanceof Error + ? error + : new Error(error === undefined ? "Unknown error" : String(error)); + Object.defineProperty(markedError, NON_STREAM_TERMINAL_PERSISTENCE_ERROR, { + configurable: false, + enumerable: false, + value: true, + writable: false, + }); + return markedError as NonStreamTerminalPersistenceError; +} + +function isNonStreamTerminalPersistenceError( + error: unknown +): error is NonStreamTerminalPersistenceError { + return ( + error instanceof Error && + NON_STREAM_TERMINAL_PERSISTENCE_ERROR in error && + (error as NonStreamTerminalPersistenceError)[NON_STREAM_TERMINAL_PERSISTENCE_ERROR] === true + ); +} + +async function persistNonStreamTerminalDetails(options: { + taskId: string; + messageRequestId: number; + durationMs: number; + details: MessageRequestTerminalDetails; + onCommitted?: () => void | Promise; +}): Promise { + const completeTerminalDetails = { + ...options.details, + durationMs: options.durationMs, + }; + try { + const committed = await updateMessageRequestDetailsDurably( + options.messageRequestId, + completeTerminalDetails, + options.onCommitted ? { onCommitted: options.onCommitted } : undefined + ); + if (committed) return true; + } catch (primaryError) { + const databaseError = findSafeDatabaseError(primaryError); + logger.error("ResponseHandler: Durable non-stream terminal persistence failed", { + taskId: options.taskId, + messageId: options.messageRequestId, + statusCode: options.details.statusCode, + error: + databaseError?.message ?? + (primaryError instanceof Error ? primaryError.message : String(primaryError)), + errorCode: databaseError?.code, + errorPool: databaseError?.pool, + }); } - if (provider.firstByteTimeoutStreamingMs > 0) { - return Math.max(provider.firstByteTimeoutStreamingMs, provider.streamingIdleTimeoutMs); + try { + return await awaitTerminalPersistenceWithOwnership({ + promise: updateMessageRequestDetailsIfUnfinalized( + options.messageRequestId, + completeTerminalDetails, + options.onCommitted ? { onCommitted: options.onCommitted } : undefined + ), + taskId: options.taskId, + operation: "nonstream-terminal-fallback", + timeoutMs: STREAM_FAILURE_PERSISTENCE_MAX_MS, + }); + } catch (fallbackError) { + const databaseError = findSafeDatabaseError(fallbackError); + logger.error("ResponseHandler: Conditional non-stream terminal fallback failed", { + taskId: options.taskId, + messageId: options.messageRequestId, + statusCode: options.details.statusCode, + error: + databaseError?.message ?? + (fallbackError instanceof Error ? fallbackError.message : String(fallbackError)), + errorCode: databaseError?.code, + errorPool: databaseError?.pool, + }); + throw markNonStreamTerminalPersistenceError(fallbackError); } +} - return Number.POSITIVE_INFINITY; +function raceWithTimeout(promise: Promise, timeoutMs: number, message: string): Promise { + let timeoutId: NodeJS.Timeout | null = null; + const timeoutPromise = new Promise((_resolve, reject) => { + timeoutId = setTimeout(() => reject(new Error(message)), timeoutMs); + timeoutId.unref?.(); + }); + + return Promise.race([promise, timeoutPromise]).finally(() => { + if (timeoutId) { + clearTimeout(timeoutId); + } + }); +} + +function raceWithDeadline( + promise: Promise, + deadlineAtMs: number, + message: string +): Promise { + const operation = Promise.resolve(promise); + const remainingMs = deadlineAtMs - Date.now(); + if (remainingMs <= 0) { + void operation.catch(() => { + // The caller has already exhausted its deadline; absorb late rejection. + }); + return Promise.reject(new Error(message)); + } + + return raceWithTimeout(operation, remainingMs, message); +} + +let terminalPersistenceTailSequence = 0; + +function awaitTerminalPersistenceWithOwnership(options: { + promise: Promise; + taskId: string; + operation: string; + timeoutMs: number; +}): Promise { + const persistence = Promise.resolve(options.promise); + let started = false; + const controller = AsyncTaskManager.register( + `${options.taskId}-${options.operation}-${++terminalPersistenceTailSequence}`, + async () => { + started = true; + try { + await persistence; + } catch { + // The request owner observes and projects the original rejection. This + // tail task only keeps late persistence joinable during shutdown. + } + }, + { + taskType: "terminal-persistence-tail", + staleTimeoutMs: Number.POSITIVE_INFINITY, + } + ); + + if (controller.signal.aborted && !started) { + return persistence; + } + + return raceWithTimeout( + persistence, + options.timeoutMs, + `${options.operation}_persistence_timeout` + ); +} + +function schedulePostTerminalSideEffects(options: { + taskId: string; + providerId: number; + sessionId: string | null; + commit: (signal: AbortSignal) => Promise; +}): Promise { + const effectTaskId = `${options.taskId}-post-terminal-effects`; + const completion = Promise.withResolvers(); + let started = false; + const run = async (signal: AbortSignal) => { + started = true; + try { + if (signal.aborted) { + logger.info("[ResponseHandler] Post-terminal side effects cancelled before start", { + taskId: options.taskId, + providerId: options.providerId, + sessionId: options.sessionId, + }); + return; + } + + let commitPromise: Promise; + try { + commitPromise = Promise.resolve(options.commit(signal)); + } catch (error) { + commitPromise = Promise.reject(error); + } + + let warningTimer: ReturnType | null = setTimeout(() => { + logger.warn("[ResponseHandler] Post-terminal side effects are still pending", { + taskId: options.taskId, + providerId: options.providerId, + sessionId: options.sessionId, + maxWaitMs: STREAM_FINALIZATION_MAX_MS, + }); + }, STREAM_FINALIZATION_MAX_MS); + warningTimer.unref?.(); + try { + await commitPromise; + } catch (error) { + logger.warn("[ResponseHandler] Post-terminal side effects failed", { + taskId: options.taskId, + providerId: options.providerId, + sessionId: options.sessionId, + error: error instanceof Error ? error.message : String(error), + }); + } finally { + if (warningTimer) { + clearTimeout(warningTimer); + warningTimer = null; + } + } + } finally { + completion.resolve(); + } + }; + const controller = AsyncTaskManager.register(effectTaskId, run, { + taskType: "post-terminal-side-effects", + staleTimeoutMs: STREAM_FINALIZATION_MAX_MS, + }); + + if (controller.signal.aborted && !started) { + // shutdownAll may have closed the registry before a late SQL commit is + // observed. Execute inline so the writer callback remains the owner and + // shutdown can join it before Redis/DB dependencies are closed. + void run(new AbortController().signal); + } + + return completion.promise; +} + +async function runPostTerminalSideEffects( + effects: ReadonlyArray<() => Promise>, + signal: AbortSignal +): Promise { + if (signal.aborted) return; + await Promise.allSettled( + effects.map((effect) => { + try { + return effect(); + } catch (error) { + return Promise.reject(error); + } + }) + ); } // 流式统计只需要头部元信息和尾部 usage/final event。按字节保存窗口,避免 @@ -409,9 +657,9 @@ async function consumeBeforeResponseBodySnapshot(session: ProxySession): Promise } } -function discardBeforeResponseBodySnapshot(session: ProxySession): void { +function discardBeforeResponseBodySnapshot(session: ProxySession): boolean { const source = takeBeforeResponseBodySnapshotSource(session); - if (!source?.body) return; + if (!source?.body) return false; void source.body.cancel().catch((error) => { logger.warn("[ResponseHandler] Failed to discard before-response snapshot body", { @@ -420,6 +668,7 @@ function discardBeforeResponseBodySnapshot(session: ProxySession): void { error, }); }); + return true; } export type UsageMetrics = { @@ -866,6 +1115,8 @@ type FinalizeDeferredStreamingResult = { usageMetrics: UsageMetrics | null; providerType: Provider["providerType"] | undefined; }; + /** Circuit and Session side effects, committed after durable terminal details. */ + commitSideEffects?: () => Promise; }; /** @@ -885,22 +1136,22 @@ type FinalizeDeferredStreamingResult = { * @param clientAborted - 标记是否为客户端主动中断(用于内部状态码映射,避免把中断记为 200 completed) * @param abortReason - 非自然结束时的原因码(用于内部记录/熔断归因;不会影响客户端响应) */ -async function finalizeDeferredStreamingFinalizationIfNeeded( +function finalizeDeferredStreamingFinalizationIfNeeded( session: ProxySession, allContent: string, upstreamStatusCode: number, streamEndedNormally: boolean, clientAborted: boolean, abortReason?: string -): Promise { +): FinalizeDeferredStreamingResult { const meta = consumeDeferredStreamingFinalization(session); const provider = session.provider; + const providerIdForPersistence = meta?.providerId ?? provider?.id ?? null; const clearSessionBinding = async () => { if (!session.sessionId) return; - await SessionManager.clearSessionProvider(session.sessionId); + await SessionManager.clearSessionProvider(session.sessionId, providerIdForPersistence); }; - const providerIdForPersistence = meta?.providerId ?? provider?.id ?? null; const isHedgeWinner = meta?.isHedgeWinner === true; const billHedgeLosers = meta?.billHedgeLosers === true; @@ -915,12 +1166,7 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( : ({ isError: false } as const); let clientAbortGateUsage: FinalizeDeferredStreamingResult["clientAbortGateUsage"]; const clientAbortCompleteSuccess = (() => { - if ( - streamEndedNormally || - !clientAborted || - upstreamStatusCode < 200 || - upstreamStatusCode >= 300 - ) { + if (!clientAborted || upstreamStatusCode < 200 || upstreamStatusCode >= 300) { return false; } @@ -965,9 +1211,16 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( } else if (clientAbortCompleteSuccess) { effectiveStatusCode = upstreamStatusCode; errorMessage = null; + } else if (streamEndedNormally && upstreamStatusCode >= 400) { + effectiveStatusCode = upstreamStatusCode; + const upstreamError = detectUpstreamErrorFromSseOrJsonText(allContent); + errorMessage = upstreamError.isError ? upstreamError.code : `HTTP ${upstreamStatusCode}`; + } else if (clientAborted) { + effectiveStatusCode = 499; + errorMessage = "CLIENT_ABORTED"; } else if (!streamEndedNormally) { - effectiveStatusCode = clientAborted ? 499 : 502; - errorMessage = clientAborted ? "CLIENT_ABORTED" : (abortReason ?? "STREAM_ABORTED"); + effectiveStatusCode = 502; + errorMessage = abortReason ?? "STREAM_ABORTED"; } else { // streamEndedNormally=true effectiveStatusCode = upstreamStatusCode; @@ -983,14 +1236,10 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( } const shouldClearSessionBindingOnFailure = - (!streamEndedNormally && !clientAbortCompleteSuccess) || + ((clientAborted || !streamEndedNormally) && !clientAbortCompleteSuccess) || detected.isError || (upstreamStatusCode >= 400 && errorMessage !== null); - if ((!meta || !provider) && shouldClearSessionBindingOnFailure) { - await clearSessionBinding(); - } - // 未启用延迟结算 / provider 缺失: // - 只返回“内部状态码 + 错误原因”,由调用方写入统计; // - 不在这里更新熔断/绑定(meta 缺失意味着 Forwarder 没有启用延迟结算;provider 缺失意味着无法归因)。 @@ -1002,6 +1251,7 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( isHedgeWinner, billHedgeLosers, clientAbortGateUsage, + commitSideEffects: shouldClearSessionBindingOnFailure ? clearSessionBinding : undefined, }; } @@ -1019,27 +1269,14 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( canonicalProviderId: meta.providerId, }); - // 尝试用 meta.providerId 找回正确的 Provider 对象,保证 providerChain 的审计数据一致 - try { - const providers = await session.getProvidersSnapshot(); - const resolved = providers.find((p) => p.id === meta.providerId); - if (resolved) { - providerForChain = resolved; - } else { - logger.warn("[ResponseHandler] Deferred streaming meta provider not found in snapshot", { - sessionId: session.sessionId ?? null, - metaProviderId: meta.providerId, - currentProviderId: provider.id, - }); - } - } catch (resolveError) { - logger.warn("[ResponseHandler] Failed to resolve meta provider from snapshot", { - sessionId: session.sessionId ?? null, - metaProviderId: meta.providerId, - currentProviderId: provider.id, - error: resolveError, - }); - } + // The deferred metadata is the canonical attempt identity. Build the audit + // entry synchronously so Provider snapshot I/O cannot block durable outcome + // persistence or retain the response body beyond the finalization deadline. + providerForChain = { + ...provider, + id: meta.providerId, + name: meta.providerName, + }; } // 未自然结束:不更新 session 绑定(避免把会话粘到不稳定 provider),但要避免把它误记为 200 completed。 @@ -1047,28 +1284,7 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( // 同时,为了让故障转移/熔断能正确工作: // - 客户端主动中断:不计入熔断器(这通常不是供应商问题) // - 非客户端中断:计入 provider/endpoint 熔断失败(与 timeout 路径保持一致) - if (!streamEndedNormally && !clientAbortCompleteSuccess) { - await clearSessionBinding(); - - if (!clientAborted && session.getEndpointPolicy().allowCircuitBreakerAccounting) { - try { - // 动态导入:避免 proxy 模块与熔断器模块之间潜在的循环依赖。 - const { recordFailure } = await import("@/lib/circuit-breaker"); - await recordFailure(meta.providerId, new Error(errorMessage ?? "STREAM_ABORTED")); - } catch (cbError) { - logger.warn("[ResponseHandler] Failed to record streaming failure in circuit breaker", { - providerId: meta.providerId, - sessionId: session.sessionId ?? null, - error: cbError, - }); - } - - // NOTE: Do NOT call recordEndpointFailure here. Stream aborts are key-level - // errors (auth, rate limit, bad key). The endpoint itself delivered HTTP 200 - // successfully. Only forwarder-level failures (timeout, network error) and - // probe failures should penalize the endpoint circuit breaker. - } - + if ((clientAborted || !streamEndedNormally) && !clientAbortCompleteSuccess) { session.addProviderToChain(providerForChain, { endpointId: meta.endpointId, endpointUrl: meta.endpointUrl, @@ -1078,6 +1294,26 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( errorMessage: errorMessage ?? undefined, }); + const commitSideEffects = async () => { + await clearSessionBinding(); + + if (!clientAborted && session.getEndpointPolicy().allowCircuitBreakerAccounting) { + try { + const { recordFailure } = await import("@/lib/circuit-breaker"); + await recordFailure(meta.providerId, new Error(errorMessage ?? "STREAM_ABORTED")); + } catch (cbError) { + logger.warn("[ResponseHandler] Failed to record streaming failure in circuit breaker", { + providerId: meta.providerId, + sessionId: session.sessionId ?? null, + error: cbError, + }); + } + + // Stream aborts are key-level errors. The endpoint delivered HTTP 200, + // so only the Provider circuit is updated here. + } + }; + return { effectiveStatusCode, errorMessage, @@ -1085,12 +1321,11 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( isHedgeWinner, billHedgeLosers, clientAbortGateUsage, + commitSideEffects, }; } if (detected.isError) { - await clearSessionBinding(); - logger.warn("[ResponseHandler] SSE completed but body indicates error (fake 200)", { providerId: meta.providerId, providerName: meta.providerName, @@ -1104,23 +1339,6 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( const chainReason = effectiveStatusCode === 404 ? "resource_not_found" : "retry_failed"; - // 计入熔断器:让后续请求能正确触发故障转移/熔断。 - // - // 注意:404 语义在 forwarder 中属于 RESOURCE_NOT_FOUND,不计入熔断器(避免把“资源/模型不存在”当作供应商故障)。 - if (effectiveStatusCode !== 404 && session.getEndpointPolicy().allowCircuitBreakerAccounting) { - try { - // 动态导入:避免 proxy 模块与熔断器模块之间潜在的循环依赖。 - const { recordFailure } = await import("@/lib/circuit-breaker"); - await recordFailure(meta.providerId, new Error(detected.code)); - } catch (cbError) { - logger.warn("[ResponseHandler] Failed to record fake-200 error in circuit breaker", { - providerId: meta.providerId, - sessionId: session.sessionId ?? null, - error: cbError, - }); - } - } - // NOTE: Do NOT call recordEndpointFailure here. Fake-200 errors are key-level // issues (invalid key, auth failure). The endpoint returned HTTP 200 successfully; // the error is in the response content, not endpoint connectivity. @@ -1139,6 +1357,27 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( errorMessage: detected.detail ? `${detected.code}: ${detected.detail}` : detected.code, }); + const commitSideEffects = async () => { + await clearSessionBinding(); + + // 404 is RESOURCE_NOT_FOUND and must not penalize the Provider circuit. + if ( + effectiveStatusCode !== 404 && + session.getEndpointPolicy().allowCircuitBreakerAccounting + ) { + try { + const { recordFailure } = await import("@/lib/circuit-breaker"); + await recordFailure(meta.providerId, new Error(detected.code)); + } catch (cbError) { + logger.warn("[ResponseHandler] Failed to record fake-200 error in circuit breaker", { + providerId: meta.providerId, + sessionId: session.sessionId ?? null, + error: cbError, + }); + } + } + }; + return { effectiveStatusCode, errorMessage, @@ -1146,13 +1385,12 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( isHedgeWinner, billHedgeLosers, clientAbortGateUsage, + commitSideEffects, }; } // ========== 非200状态码处理(流自然结束但HTTP状态码表示错误)========== if (upstreamStatusCode >= 400 && errorMessage !== null) { - await clearSessionBinding(); - logger.warn("[ResponseHandler] SSE completed but HTTP status indicates error", { providerId: meta.providerId, providerName: meta.providerName, @@ -1163,21 +1401,6 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( const chainReason = effectiveStatusCode === 404 ? "resource_not_found" : "retry_failed"; - // 计入熔断器:让后续请求能正确触发故障转移/熔断。 - // 注意:与 forwarder 口径保持一致:404 不计入熔断器(资源不存在不是供应商故障)。 - if (effectiveStatusCode !== 404 && session.getEndpointPolicy().allowCircuitBreakerAccounting) { - try { - const { recordFailure } = await import("@/lib/circuit-breaker"); - await recordFailure(meta.providerId, new Error(errorMessage)); - } catch (cbError) { - logger.warn("[ResponseHandler] Failed to record non-200 error in circuit breaker", { - providerId: meta.providerId, - sessionId: session.sessionId ?? null, - error: cbError, - }); - } - } - // NOTE: Do NOT call recordEndpointFailure here. Non-200 HTTP errors (401, 429, // etc.) are typically key/auth-level errors. The endpoint was reachable and // responded; only forwarder-level failures should penalize the endpoint breaker. @@ -1192,6 +1415,26 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( errorMessage: errorMessage, }); + const commitSideEffects = async () => { + await clearSessionBinding(); + + if ( + effectiveStatusCode !== 404 && + session.getEndpointPolicy().allowCircuitBreakerAccounting + ) { + try { + const { recordFailure } = await import("@/lib/circuit-breaker"); + await recordFailure(meta.providerId, new Error(errorMessage)); + } catch (cbError) { + logger.warn("[ResponseHandler] Failed to record non-200 error in circuit breaker", { + providerId: meta.providerId, + sessionId: session.sessionId ?? null, + error: cbError, + }); + } + } + }; + return { effectiveStatusCode, errorMessage, @@ -1199,38 +1442,51 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( isHedgeWinner, billHedgeLosers, clientAbortGateUsage, + commitSideEffects, }; } // ========== 真正成功(SSE 完整结束且未命中错误判定)========== - if (meta.endpointId != null) { + // Build the durable audit chain before persistence, but defer external + // circuit/Session mutations until billing and terminal details are committed. + // A slow Redis binding must not turn an already completed, billable request + // into a fallback 500 or prevent lease settlement. + if (!meta.isHedgeWinner) { + session.addProviderToChain(providerForChain, { + endpointId: meta.endpointId, + endpointUrl: meta.endpointUrl, + reason: meta.isFirstAttempt ? "request_success" : "retry_success", + attemptNumber: meta.attemptNumber, + statusCode: meta.upstreamStatusCode, + }); + } + + const commitSideEffects = async () => { + if (meta.endpointId != null) { + try { + const { recordEndpointSuccess } = await import("@/lib/endpoint-circuit-breaker"); + await recordEndpointSuccess(meta.endpointId); + } catch (endpointError) { + logger.warn("[ResponseHandler] Failed to record endpoint success (stream finalized)", { + endpointId: meta.endpointId, + providerId: meta.providerId, + error: endpointError, + }); + } + } + try { - const { recordEndpointSuccess } = await import("@/lib/endpoint-circuit-breaker"); - await recordEndpointSuccess(meta.endpointId); - } catch (endpointError) { - logger.warn("[ResponseHandler] Failed to record endpoint success (stream finalized)", { - endpointId: meta.endpointId, + const { recordSuccess } = await import("@/lib/circuit-breaker"); + await recordSuccess(meta.providerId); + } catch (cbError) { + logger.warn("[ResponseHandler] Failed to record streaming success in circuit breaker", { providerId: meta.providerId, - error: endpointError, + error: cbError, }); } - } - - try { - const { recordSuccess } = await import("@/lib/circuit-breaker"); - await recordSuccess(meta.providerId); - } catch (cbError) { - logger.warn("[ResponseHandler] Failed to record streaming success in circuit breaker", { - providerId: meta.providerId, - error: cbError, - }); - } - // Hedge winner: commitWinner() already performed session binding and chain logging. - // Skip duplicate operations to avoid double entries in the provider chain. - if (!meta.isHedgeWinner) { - // 成功后绑定 session 到供应商(智能绑定策略) - if (session.sessionId) { + // Hedge winner: commitWinner() already performed session binding and chain logging. + if (!meta.isHedgeWinner && session.sessionId) { const result = await SessionManager.updateSessionBindingSmart( session.sessionId, meta.providerId, @@ -1262,7 +1518,6 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( }); } - // 统一更新两个数据源(确保监控数据一致) if (session.shouldTrackSessionObservability()) { void SessionManager.updateSessionProvider(session.sessionId, { providerId: meta.providerId, @@ -1270,30 +1525,20 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( }).catch((err) => { logger.error( "[ResponseHandler] Failed to update session provider info (stream finalized)", - { - error: err, - } + { error: err } ); }); } } - session.addProviderToChain(providerForChain, { - endpointId: meta.endpointId, - endpointUrl: meta.endpointUrl, - reason: meta.isFirstAttempt ? "request_success" : "retry_success", + logger.info("[ResponseHandler] Streaming request finalized as success", { + providerId: meta.providerId, + providerName: meta.providerName, attemptNumber: meta.attemptNumber, + totalProvidersAttempted: meta.totalProvidersAttempted, statusCode: meta.upstreamStatusCode, }); - } - - logger.info("[ResponseHandler] Streaming request finalized as success", { - providerId: meta.providerId, - providerName: meta.providerName, - attemptNumber: meta.attemptNumber, - totalProvidersAttempted: meta.totalProvidersAttempted, - statusCode: meta.upstreamStatusCode, - }); + }; return { effectiveStatusCode, @@ -1302,6 +1547,7 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( isHedgeWinner, billHedgeLosers, clientAbortGateUsage, + commitSideEffects, }; } @@ -1310,7 +1556,8 @@ export class ProxyResponseHandler { const snapshotSession = session as ProxySession & { detailSnapshotResponseBeforeSource?: Response | null; }; - if (session.sessionId && session.shouldPersistSessionDebugArtifacts()) { + const isStreamingResponse = response.headers.get("content-type")?.includes("text/event-stream"); + if (!isStreamingResponse && session.sessionId && session.shouldPersistSessionDebugArtifacts()) { snapshotSession.detailSnapshotResponseBeforeSource = response.clone(); } @@ -1412,7 +1659,7 @@ export class ProxyResponseHandler { statsAbortController, taskId ); - const statsPromise = (async () => { + const runStatsTask = async () => { try { const responseText = await readResponseTextWithTaskActivity(responseForStats, taskId); @@ -1457,29 +1704,35 @@ export class ProxyResponseHandler { // 非200状态码处理:解析错误响应并计入熔断器 let errorMessageForFinalize: string | undefined; + let commitProviderFailure: (() => Promise) | undefined; if (statusCode >= 400) { const detected = detectUpstreamErrorFromSseOrJsonText(responseText); errorMessageForFinalize = detected.isError ? detected.code : `HTTP ${statusCode}`; - - // 计入熔断器 - if (session.getEndpointPolicy().allowCircuitBreakerAccounting) { - try { - const { recordFailure } = await import("@/lib/circuit-breaker"); - await recordFailure(provider.id, new Error(errorMessageForFinalize)); - } catch (cbError) { - logger.warn( - "ResponseHandler: Failed to record non-200 error in circuit breaker (passthrough)", - { - providerId: provider.id, - error: cbError, - } - ); - } + const isResourceNotFound = statusCode === 404; + + if ( + !isResourceNotFound && + session.getEndpointPolicy().allowCircuitBreakerAccounting + ) { + commitProviderFailure = async () => { + try { + const { recordFailure } = await import("@/lib/circuit-breaker"); + await recordFailure(provider.id, new Error(errorMessageForFinalize)); + } catch (cbError) { + logger.warn( + "ResponseHandler: Failed to record non-200 error in circuit breaker (passthrough)", + { + providerId: provider.id, + error: cbError, + } + ); + } + }; } // 记录到决策链 session.addProviderToChain(provider, { - reason: "retry_failed", + reason: isResourceNotFound ? "resource_not_found" : "retry_failed", attemptNumber: 1, statusCode: statusCode, errorMessage: errorMessageForFinalize, @@ -1488,6 +1741,20 @@ export class ProxyResponseHandler { // 使用共享的统计处理方法 const duration = Date.now() - session.startTime; + let providerFailureScheduled = false; + const scheduleProviderFailure = () => { + if (!commitProviderFailure || providerFailureScheduled) return; + providerFailureScheduled = true; + return schedulePostTerminalSideEffects({ + taskId, + providerId: provider.id, + sessionId: session.sessionId, + commit: async (signal) => { + if (signal.aborted) return; + await commitProviderFailure(); + }, + }); + }; const finalizedUsage = await finalizeRequestStats( session, responseText, @@ -1495,9 +1762,9 @@ export class ProxyResponseHandler { duration, errorMessageForFinalize, undefined, - false // Gemini 非流式透传 + false, // Gemini 非流式透传 + scheduleProviderFailure ); - emitProxyLangfuseTrace(session, { responseHeaders: response.headers, responseText, @@ -1512,32 +1779,120 @@ export class ProxyResponseHandler { if (session.sessionId && session.shouldPersistSessionDebugArtifacts()) { await discardBeforeResponseBodySnapshot(session); } - if (!isClientAbortError(error as Error)) { + const clientAborted = isClientAbortError(error as Error); + if (!clientAborted) { logger.error( "[ResponseHandler] Gemini non-stream passthrough stats task failed:", error ); } + + let finalizedStatusCode = statusCode >= 400 ? statusCode : 502; + if (clientAborted) { + finalizedStatusCode = 499; + } + const isResourceNotFound = finalizedStatusCode === 404; + const errorDetails = buildProcessingErrorDetails(error); + if (!clientAborted) { + session.addProviderToChain(provider, { + reason: isResourceNotFound ? "resource_not_found" : "retry_failed", + attemptNumber: 1, + statusCode: finalizedStatusCode, + errorMessage: errorDetails.errorMessage, + }); + } + + const postTerminalSideEffects: Array<() => Promise> = []; + if (session.sessionId) { + const sessionId = session.sessionId; + postTerminalSideEffects.push(async () => { + await SessionManager.clearSessionProvider(sessionId, provider.id); + }); + } + if ( + !clientAborted && + !isResourceNotFound && + session.getEndpointPolicy().allowCircuitBreakerAccounting + ) { + postTerminalSideEffects.push(async () => { + try { + const { recordFailure } = await import("@/lib/circuit-breaker"); + await recordFailure(provider.id, error as Error); + } catch (cbError) { + logger.warn( + "ResponseHandler: Failed to record Gemini non-stream body failure in circuit breaker", + { + providerId: provider.id, + error: cbError, + } + ); + } + }); + } + let postTerminalSideEffectsScheduled = false; + const scheduleCommittedSideEffects = () => { + if (postTerminalSideEffects.length === 0 || postTerminalSideEffectsScheduled) return; + postTerminalSideEffectsScheduled = true; + return schedulePostTerminalSideEffects({ + taskId, + providerId: provider.id, + sessionId: session.sessionId, + commit: (signal) => runPostTerminalSideEffects(postTerminalSideEffects, signal), + }); + }; + + if (messageContext) { + const duration = Date.now() - session.startTime; + const tracker = ProxyStatusTracker.getInstance(); + try { + await persistNonStreamTerminalDetails({ + taskId, + messageRequestId: messageContext.id, + durationMs: duration, + details: { + statusCode: finalizedStatusCode, + ...errorDetails, + ttfbMs: session.ttfbMs ?? duration, + providerChain: session.getProviderChain(), + model: session.getCurrentModel() ?? undefined, + providerId: session.provider?.id, + context1mApplied: session.getContext1mApplied(), + swapCacheTtlApplied: session.provider?.swapCacheTtlBilling ?? false, + specialSettings: session.getSpecialSettings() ?? undefined, + }, + onCommitted: scheduleCommittedSideEffects, + }); + } finally { + tracker.endRequest(messageContext.user.id, messageContext.id); + } + } } finally { cleanupTaskAbortBinding(); releaseSessionAgent(session); } - })(); + }; - AsyncTaskManager.register(taskId, statsPromise, { - taskType: "non-stream-passthrough-stats", - abortController: statsAbortController, - staleTimeoutMs: resolveNonStreamTaskStaleTimeoutMs(provider), - }); - statsPromise.catch((error) => { - if (session.sessionId && session.shouldPersistSessionDebugArtifacts()) { - void discardBeforeResponseBodySnapshot(session); + AsyncTaskManager.register( + taskId, + () => { + const statsPromise = runStatsTask(); + statsPromise.catch((error) => { + if (session.sessionId && session.shouldPersistSessionDebugArtifacts()) { + void discardBeforeResponseBodySnapshot(session); + } + logger.error( + "[ResponseHandler] Gemini non-stream passthrough stats task uncaught error:", + error + ); + }); + return statsPromise; + }, + { + taskType: "non-stream-passthrough-stats", + abortController: statsAbortController, + staleTimeoutMs: resolveNonStreamTaskStaleTimeoutMs(provider), } - logger.error( - "[ResponseHandler] Gemini non-stream passthrough stats task uncaught error:", - error - ); - }); + ); if (session.sessionId && session.shouldPersistSessionDebugArtifacts()) { const responseAfterMetaTask = SessionManager.storeSessionResponsePhaseSnapshot?.( @@ -1605,39 +1960,76 @@ export class ProxyResponseHandler { abortController.abort(); }); - const processingPromise = (async () => { - const finalizeNonStreamAbort = async (): Promise => { - const finalizedStatusCode = session.clientAbortSignal?.aborted ? 499 : statusCode; - if (messageContext) { - const duration = Date.now() - session.startTime; - await updateMessageRequestDuration(messageContext.id, duration); - await updateMessageRequestDetails(messageContext.id, { - statusCode: finalizedStatusCode, - ttfbMs: session.ttfbMs ?? duration, - providerChain: session.getProviderChain(), - model: session.getCurrentModel() ?? undefined, // 更新重定向后的模型 - providerId: session.provider?.id, // 更新最终供应商ID(重试切换后) - context1mApplied: session.getContext1mApplied(), - swapCacheTtlApplied: session.provider?.swapCacheTtlBilling ?? false, - }); - const tracker = ProxyStatusTracker.getInstance(); - tracker.endRequest(messageContext.user.id, messageContext.id); - } - + const runProcessingTask = async () => { + const finalizeNonStreamAbort = async ( + options: { + statusCode?: number; + error?: unknown; + postTerminalSideEffects?: Array<() => Promise>; + } = {} + ): Promise => { + const finalizedStatusCode = + options.statusCode ?? (session.clientAbortSignal?.aborted ? 499 : statusCode); + const errorDetails = + options.error === undefined ? undefined : buildProcessingErrorDetails(options.error); + const postTerminalSideEffects = [...(options.postTerminalSideEffects ?? [])]; if (session.sessionId) { - await SessionManager.clearSessionProvider(session.sessionId); - - const sessionUsagePayload: SessionUsageUpdate = { - status: finalizedStatusCode >= 200 && finalizedStatusCode < 300 ? "completed" : "error", - statusCode: finalizedStatusCode, - }; + const sessionId = session.sessionId; + postTerminalSideEffects.push(async () => { + await SessionManager.clearSessionProvider(sessionId, provider.id); + + const sessionUsagePayload: SessionUsageUpdate = { + status: + finalizedStatusCode >= 200 && finalizedStatusCode < 300 ? "completed" : "error", + statusCode: finalizedStatusCode, + ...(errorDetails?.errorMessage + ? { errorMessage: errorDetails.errorMessage } + : undefined), + }; - if (session.shouldTrackSessionObservability()) { - void SessionManager.updateSessionUsage(session.sessionId, sessionUsagePayload).catch( - (error: unknown) => { + if (session.shouldTrackSessionObservability()) { + try { + await SessionManager.updateSessionUsage(sessionId, sessionUsagePayload); + } catch (error) { logger.error("[ResponseHandler] Failed to update session usage:", error); } - ); + } + }); + } + let postTerminalSideEffectsScheduled = false; + const scheduleCommittedSideEffects = () => { + if (postTerminalSideEffects.length === 0 || postTerminalSideEffectsScheduled) return; + postTerminalSideEffectsScheduled = true; + return schedulePostTerminalSideEffects({ + taskId, + providerId: provider.id, + sessionId: session.sessionId, + commit: (signal) => runPostTerminalSideEffects(postTerminalSideEffects, signal), + }); + }; + if (messageContext) { + const duration = Date.now() - session.startTime; + const terminalDetails: MessageRequestTerminalDetails = { + statusCode: finalizedStatusCode, + ...errorDetails, + ttfbMs: session.ttfbMs ?? duration, + providerChain: session.getProviderChain(), + model: session.getCurrentModel() ?? undefined, // 更新重定向后的模型 + providerId: session.provider?.id, // 更新最终供应商ID(重试切换后) + context1mApplied: session.getContext1mApplied(), + swapCacheTtlApplied: session.provider?.swapCacheTtlBilling ?? false, + }; + const tracker = ProxyStatusTracker.getInstance(); + try { + await persistNonStreamTerminalDetails({ + taskId, + messageRequestId: messageContext.id, + durationMs: duration, + details: terminalDetails, + onCommitted: scheduleCommittedSideEffects, + }); + } finally { + tracker.endRequest(messageContext.user.id, messageContext.id); } } }; @@ -1657,6 +2049,9 @@ export class ProxyResponseHandler { providerId: provider.id, finalizeError, }); + if (isNonStreamTerminalPersistenceError(finalizeError)) { + throw finalizeError; + } } return; } @@ -1672,6 +2067,7 @@ export class ProxyResponseHandler { sessionWithCleanup.clearResponseTimeout(); } let usageMetrics: UsageMetrics | null = null; + const postTerminalSideEffects: Array<() => Promise> = []; const usageResult = parseUsageFromResponseText(responseText, provider.providerType); usageMetrics = usageResult.usageMetrics; @@ -1709,18 +2105,30 @@ export class ProxyResponseHandler { } // Codex: Extract prompt_cache_key and update session binding - if (provider.providerType === "codex" && session.sessionId && provider.id) { + if ( + provider.providerType === "codex" && + statusCode >= 200 && + statusCode < 300 && + session.sessionId && + provider.id + ) { try { const responseData = JSON.parse(responseText) as Record; const promptCacheKey = SessionManager.extractCodexPromptCacheKey(responseData); if (promptCacheKey) { - void SessionManager.updateSessionWithCodexCacheKey( - session.sessionId, - promptCacheKey, - provider.id, - session.authState?.key?.id ?? session.messageContext?.key?.id ?? null - ).catch((err) => { - logger.error("[ResponseHandler] Failed to update Codex session:", err); + const sessionId = session.sessionId; + const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; + postTerminalSideEffects.push(async () => { + try { + await SessionManager.updateSessionWithCodexCacheKey( + sessionId, + promptCacheKey, + provider.id, + keyId + ); + } catch (err) { + logger.error("[ResponseHandler] Failed to update Codex session:", err); + } }); } } catch (parseError) { @@ -1863,39 +2271,52 @@ export class ProxyResponseHandler { }); } - // 非200状态码处理:解析错误响应并计入熔断器 + // 非200状态码处理:先构造审计链,durable details 后再更新熔断器。 + let terminalErrorMessage: string | undefined; if (statusCode >= 400) { const detected = detectUpstreamErrorFromSseOrJsonText(responseText); const errorMessageForDb = detected.isError ? detected.code : `HTTP ${statusCode}`; + terminalErrorMessage = errorMessageForDb; + const isResourceNotFound = statusCode === 404; - // 计入熔断器 - if (session.getEndpointPolicy().allowCircuitBreakerAccounting) { - try { - const { recordFailure } = await import("@/lib/circuit-breaker"); - await recordFailure(provider.id, new Error(errorMessageForDb)); - } catch (cbError) { - logger.warn("ResponseHandler: Failed to record non-200 error in circuit breaker", { - providerId: provider.id, - error: cbError, - }); - } + if (!isResourceNotFound && session.getEndpointPolicy().allowCircuitBreakerAccounting) { + postTerminalSideEffects.push(async () => { + try { + const { recordFailure } = await import("@/lib/circuit-breaker"); + await recordFailure(provider.id, new Error(errorMessageForDb)); + } catch (cbError) { + logger.warn("ResponseHandler: Failed to record non-200 error in circuit breaker", { + providerId: provider.id, + error: cbError, + }); + } + }); } // 记录到决策链 session.addProviderToChain(provider, { - reason: "retry_failed", + reason: isResourceNotFound ? "resource_not_found" : "retry_failed", attemptNumber: 1, statusCode: statusCode, errorMessage: errorMessageForDb, }); } + let postTerminalSideEffectsScheduled = false; + const scheduleCommittedSideEffects = () => { + if (postTerminalSideEffects.length === 0 || postTerminalSideEffectsScheduled) return; + postTerminalSideEffectsScheduled = true; + return schedulePostTerminalSideEffects({ + taskId, + providerId: provider.id, + sessionId: session.sessionId, + commit: (signal) => runPostTerminalSideEffects(postTerminalSideEffects, signal), + }); + }; + if (messageContext) { const duration = Date.now() - session.startTime; - await updateMessageRequestDuration(messageContext.id, duration); - - // 保存扩展信息(status code, tokens, provider chain) - await updateMessageRequestDetails(messageContext.id, { + const terminalDetails: MessageRequestTerminalDetails = { statusCode: statusCode, inputTokens: usageMetrics?.input_tokens, outputTokens: usageMetrics?.output_tokens, @@ -1906,6 +2327,7 @@ export class ProxyResponseHandler { cacheCreation1hInputTokens: usageMetrics?.cache_creation_1h_input_tokens, cacheTtlApplied: usageMetrics?.cache_ttl ?? null, providerChain: session.getProviderChain(), + ...(terminalErrorMessage ? { errorMessage: terminalErrorMessage } : {}), model: session.getCurrentModel() ?? undefined, // 更新重定向后的模型 actualResponseModel: extractActualResponseModelForProvider( provider.providerType, @@ -1916,11 +2338,19 @@ export class ProxyResponseHandler { context1mApplied: session.getContext1mApplied(), swapCacheTtlApplied: session.provider?.swapCacheTtlBilling ?? false, specialSettings: session.getSpecialSettings() ?? undefined, - }); - - // 记录请求结束 + }; const tracker = ProxyStatusTracker.getInstance(); - tracker.endRequest(messageContext.user.id, messageContext.id); + try { + await persistNonStreamTerminalDetails({ + taskId, + messageRequestId: messageContext.id, + durationMs: duration, + details: terminalDetails, + onCommitted: scheduleCommittedSideEffects, + }); + } finally { + tracker.endRequest(messageContext.user.id, messageContext.id); + } } logger.debug("ResponseHandler: Non-stream response processed", { @@ -1941,6 +2371,9 @@ export class ProxyResponseHandler { isStreaming: false, }); } catch (error) { + if (isNonStreamTerminalPersistenceError(error)) { + throw error; + } if (session.sessionId && session.shouldPersistSessionDebugArtifacts()) { await discardBeforeResponseBodySnapshot(session); } @@ -1958,7 +2391,6 @@ export class ProxyResponseHandler { !session.clientAbortSignal?.aborted; if (isResponseTimeout) { - // ⚠️ 响应超时:计入熔断器并记录错误日志 logger.error("ResponseHandler: Response timeout during non-stream body read", { taskId, providerId: provider.id, @@ -1966,44 +2398,48 @@ export class ProxyResponseHandler { errorName: err.name, }); - // 计入熔断器(动态导入避免循环依赖) - if (session.getEndpointPolicy().allowCircuitBreakerAccounting) { - try { - const { recordFailure } = await import("@/lib/circuit-breaker"); - await recordFailure(provider.id, err); - logger.debug("ResponseHandler: Response timeout recorded in circuit breaker", { - providerId: provider.id, - }); - } catch (cbError) { - logger.warn("ResponseHandler: Failed to record timeout in circuit breaker", { - providerId: provider.id, - error: cbError, - }); - } + const finalizedStatusCode = statusCode >= 400 ? statusCode : 502; + const isResourceNotFound = finalizedStatusCode === 404; + const postTerminalSideEffects: Array<() => Promise> = []; + if (!isResourceNotFound && session.getEndpointPolicy().allowCircuitBreakerAccounting) { + postTerminalSideEffects.push(async () => { + try { + const { recordFailure } = await import("@/lib/circuit-breaker"); + await recordFailure(provider.id, err); + logger.debug("ResponseHandler: Response timeout recorded in circuit breaker", { + providerId: provider.id, + }); + } catch (cbError) { + logger.warn("ResponseHandler: Failed to record timeout in circuit breaker", { + providerId: provider.id, + error: cbError, + }); + } + }); } - // 注意:无法重试,因为客户端已收到 HTTP 200 - // 错误已记录,熔断器已更新,不抛出异常(避免影响后台任务) - - // 更新数据库记录(避免 orphan record) - await persistRequestFailure({ - session, - messageContext, - statusCode: statusCode && statusCode >= 400 ? statusCode : 502, - error: err, - taskId, - phase: "non-stream", + session.addProviderToChain(provider, { + reason: isResourceNotFound ? "resource_not_found" : "retry_failed", + attemptNumber: 1, + statusCode: finalizedStatusCode, + errorMessage: formatProcessingError(err), }); - // 执行清理逻辑 try { - await finalizeNonStreamAbort(); + await finalizeNonStreamAbort({ + statusCode: finalizedStatusCode, + error: err, + postTerminalSideEffects, + }); } catch (finalizeError) { logger.error("ResponseHandler: Failed to finalize aborted non-stream response", { taskId, providerId: provider.id, finalizeError, }); + if (isNonStreamTerminalPersistenceError(finalizeError)) { + throw finalizeError; + } } } else { // 客户端主动中断:正常日志,不抛出错误 @@ -2025,6 +2461,9 @@ export class ProxyResponseHandler { providerId: provider.id, finalizeError, }); + if (isNonStreamTerminalPersistenceError(finalizeError)) { + throw finalizeError; + } } } } else { @@ -2045,30 +2484,41 @@ export class ProxyResponseHandler { cleanupClientAbortListener(); releaseSessionAgent(session); } - })(); + }; // 注册任务并添加全局错误捕获 - AsyncTaskManager.register(taskId, processingPromise, { - taskType: "non-stream-processing", - abortController, - staleTimeoutMs: resolveNonStreamTaskStaleTimeoutMs(provider), - }); - processingPromise.catch(async (error) => { - logger.error("ResponseHandler: Uncaught error in non-stream processing", { - taskId, - error, - }); + AsyncTaskManager.register( + taskId, + () => { + const processingPromise = runProcessingTask(); + return processingPromise.catch(async (error) => { + logger.error("ResponseHandler: Uncaught error in non-stream processing", { + taskId, + error, + }); - // 更新数据库记录(避免 orphan record) - await persistRequestFailure({ - session, - messageContext, - statusCode: statusCode && statusCode >= 400 ? statusCode : 500, - error, - taskId, - phase: "non-stream", - }); - }); + if (isNonStreamTerminalPersistenceError(error)) { + throw error; + } + + // 更新数据库记录(避免 orphan record) + await persistRequestFailure({ + session, + messageContext, + statusCode: statusCode && statusCode >= 400 ? statusCode : 500, + error, + taskId, + phase: "non-stream", + }); + throw error; + }); + }, + { + taskType: "non-stream-processing", + abortController, + staleTimeoutMs: resolveNonStreamTaskStaleTimeoutMs(provider), + } + ); return finalResponse; } @@ -2093,35 +2543,70 @@ export class ProxyResponseHandler { (provider.providerType === "gemini" || provider.providerType === "gemini-cli"); if (isGeminiPassthrough) { - // 完全透传:clone 用于后台统计,返回原始 response - logger.debug( - "[ResponseHandler] Gemini stream passthrough (clone for stats, return original)", - { - originalFormat: session.originalFormat, - providerType: provider.providerType, - model: session.request.model, - statusCode: response.status, - reason: "Client receives untouched response, stats read from clone", - } - ); + logger.debug("[ResponseHandler] Gemini stream passthrough (demand-driven stats)", { + originalFormat: session.originalFormat, + providerType: provider.providerType, + model: session.request.model, + statusCode: response.status, + reason: "Client receives untouched chunks observed by the authoritative pump", + }); + discardBeforeResponseBodySnapshot(session); // 注意:不要在“仅收到响应头”时清除首字节超时。 // 背景:部分上游可能会快速返回 200 + SSE headers,但随后长时间不发送任何 body 数据。 // 若在 headers 阶段就 clearResponseTimeout,会导致首字节超时失效,客户端与服务端都会表现为一直“请求中”。 // 透传场景下,我们在后台 stats 读取到第一块数据时再清除超时(与非透传路径口径一致)。 - const responseForStats = response.clone(); + let observePassthroughChunk = (_value: Uint8Array) => {}; + let observePassthroughReadStart = () => {}; + let observePassthroughDrainStart = () => {}; + let abortPassthroughTransport = (_reason: Error) => {}; + let passthroughPump: DemandDrivenResponsePump; + let passthroughDrainTimeoutId: ReturnType | null = null; + const clearPassthroughDrainTimeout = () => { + if (passthroughDrainTimeoutId) { + clearTimeout(passthroughDrainTimeoutId); + passthroughDrainTimeoutId = null; + } + }; + const startPassthroughDrain = (reason?: unknown) => { + passthroughPump.startDrain(reason); + observePassthroughDrainStart(); + if (passthroughDrainTimeoutId) return; + passthroughDrainTimeoutId = setTimeout(() => { + passthroughDrainTimeoutId = null; + const drainTimeoutError = new Error("client_abort_drain_timeout"); + abortPassthroughTransport(drainTimeoutError); + passthroughPump.cancelSource(drainTimeoutError); + }, CLIENT_ABORT_DRAIN_MAX_MS); + passthroughDrainTimeoutId.unref?.(); + }; + passthroughPump = createDemandDrivenResponsePump({ + source: response.body, + onReadStart: () => observePassthroughReadStart(), + onChunk: (value) => observePassthroughChunk(value), + onClientCancel: (reason) => { + startPassthroughDrain(reason); + }, + }); + const cleanupPassthroughClientAbortListener = bindClientAbortListener( + session.clientAbortSignal, + () => { + const reason = session.clientAbortSignal?.reason; + startPassthroughDrain(reason); + } + ); const statusCode = response.status; const taskId = `stream-passthrough-${messageContext.id}`; - const streamTaskStaleTimeoutMs = resolveStreamTaskStaleTimeoutMs(provider); + const streamTaskStaleTimeoutMs = resolveStreamTaskStaleTimeoutMs(); const statsAbortController = new AbortController(); const cleanupTaskAbortBinding = bindTaskAbortToUpstreamResponse( session, statsAbortController, taskId ); - const statsPromise = (async () => { + const runStatsTask = async () => { const sessionWithCleanup = session as typeof session & { clearResponseTimeout?: () => void; }; @@ -2129,15 +2614,32 @@ export class ProxyResponseHandler { responseController?: AbortController; }; - let reader: ReadableStreamDefaultReader | null = null; const streamTextAccumulator = new BoundedStreamTextAccumulator(); let lastStreamTextSnapshot: BoundedStreamTextSnapshot | null = null; const getCollectedChunkCount = () => lastStreamTextSnapshot?.chunkCount ?? streamTextAccumulator.chunkCount; let isFirstChunk = true; let streamEndedNormally = false; + let terminalFinalizationStarted = false; let responseTimeoutCleared = false; + let pumpClientAborted = false; let abortReason: string | undefined; + let transportReleased = false; + let commitSideEffectsScheduled = false; + let latestCommitSideEffects: (() => Promise) | undefined; + const scheduleCommitSideEffects = (effect: (() => Promise) | undefined) => { + if (!effect || commitSideEffectsScheduled) return; + commitSideEffectsScheduled = true; + return schedulePostTerminalSideEffects({ + taskId, + providerId: provider.id, + sessionId: session.sessionId, + commit: async (signal) => { + if (signal.aborted) return; + await effect(); + }, + }); + }; // 静默期 Watchdog:透传也需要支持中途卡住(无新数据推送) const idleTimeoutMs = @@ -2175,6 +2677,23 @@ export class ProxyResponseHandler { }, idleTimeoutMs); }; + observePassthroughReadStart = () => { + if (!isFirstChunk) startIdleTimer(); + }; + observePassthroughDrainStart = startIdleTimer; + abortPassthroughTransport = (reason) => { + try { + sessionWithController.responseController?.abort(reason); + } catch { + // ignore + } + try { + statsAbortController.abort(reason); + } catch { + // ignore + } + }; + const clearResponseTimeoutOnce = (firstChunkSize?: number) => { if (responseTimeoutCleared) return; if (!sessionWithCleanup.clearResponseTimeout) return; @@ -2203,58 +2722,62 @@ export class ProxyResponseHandler { return flushAndSnapshot().text; }; - try { - const body = responseForStats.body; - if (!body) return; - reader = body.getReader(); - - // 注意:即使 STORE_SESSION_RESPONSE_BODY=false(不写入 Redis),这里也会在内存中累积完整流内容: - // - 用于解析 usage/cost 与内部结算(例如“假 200”检测) - // 因此该开关仅影响“是否持久化”,不用于控制流式内存占用。 - while (true) { - if (session.clientAbortSignal?.aborted) break; - - const { done, value } = await reader.read(); - if (done) { - const wasResponseControllerAborted = - sessionWithController.responseController?.signal.aborted ?? false; - const clientAborted = session.clientAbortSignal?.aborted ?? false; - - // abort -> nodeStreamToWebStreamSafe 可能会把错误吞掉并 close(),导致 done=true; - // 这里必须结合 abort signal 判断是否为“自然结束”。 - if (wasResponseControllerAborted || clientAborted) { - streamEndedNormally = false; - if (!abortReason) { - abortReason = clientAborted ? "CLIENT_ABORTED" : "STREAM_RESPONSE_TIMEOUT"; - } - } else { - streamEndedNormally = true; - } - break; - } - - const chunkSize = value?.byteLength ?? 0; - if (value && chunkSize > 0) { - if (isFirstChunk) { - isFirstChunk = false; - session.recordTtfb(); - clearResponseTimeoutOnce(chunkSize); - } - - streamTextAccumulator.pushBytes(value); - AsyncTaskManager.touch(taskId); - } + observePassthroughChunk = (value) => { + clearIdleTimer(); + if (isFirstChunk) { + isFirstChunk = false; + session.recordTtfb(); + clearResponseTimeoutOnce(value.byteLength); + } + streamTextAccumulator.pushBytes(value); + AsyncTaskManager.touch(taskId); + }; - // 首块数据到达后才启动 idle timer(避免与首字节超时职责重叠) - if (!isFirstChunk) { - startIdleTimer(); + const releaseTransportResources = () => { + if (transportReleased) return; + transportReleased = true; + clearPassthroughDrainTimeout(); + cleanupPassthroughClientAbortListener(); + cleanupTaskAbortBinding(); + clearIdleTimer(); + try { + const wasResponseControllerAborted = + sessionWithController.responseController?.signal.aborted ?? false; + const clientAborted = session.clientAbortSignal?.aborted ?? false; + const shouldClearTimeout = + responseTimeoutCleared || + streamEndedNormally || + wasResponseControllerAborted || + clientAborted; + if (shouldClearTimeout) { + clearResponseTimeoutOnce(); } + } catch (error) { + logger.warn( + "[ResponseHandler] Gemini passthrough: Failed to clear response timeout", + { + taskId, + providerId: provider.id, + providerName: provider.name, + error: error instanceof Error ? error.message : String(error), + } + ); } + releaseSessionAgent(session); + }; + + try { + const pumpCompletion = await passthroughPump.completion; + streamEndedNormally = pumpCompletion.streamEndedNormally; + pumpClientAborted = pumpCompletion.clientAborted; + if (pumpCompletion.error) throw pumpCompletion.error; clearIdleTimer(); const streamSnapshot = flushAndSnapshot(); const allContent = streamSnapshot.text; - const clientAborted = session.clientAbortSignal?.aborted ?? false; + const clientAborted = + pumpClientAborted || (session.clientAbortSignal?.aborted ?? false); + releaseTransportResources(); // 存储响应体到 Redis(5分钟过期) if ( @@ -2302,6 +2825,7 @@ export class ProxyResponseHandler { // 使用共享的统计处理方法 const duration = Date.now() - session.startTime; + terminalFinalizationStarted = true; const finalized = await finalizeDeferredStreamingFinalizationIfNeeded( session, allContent, @@ -2310,6 +2834,7 @@ export class ProxyResponseHandler { clientAborted, abortReason ); + latestCommitSideEffects = finalized.commitSideEffects; const finalizedUsage = await finalizeRequestStats( session, allContent, @@ -2317,9 +2842,9 @@ export class ProxyResponseHandler { duration, finalized.errorMessage ?? undefined, finalized.providerIdForPersistence ?? undefined, - true // Gemini 流式透传(NDJSON 无 data:/event: 前缀,必须显式告知) + true, // Gemini 流式透传(NDJSON 无 data:/event: 前缀,必须显式告知) + () => scheduleCommitSideEffects(latestCommitSideEffects) ); - emitProxyLangfuseTrace(session, { responseHeaders: response.headers, responseText: allContent, @@ -2332,7 +2857,8 @@ export class ProxyResponseHandler { }); } catch (error) { const err = error instanceof Error ? error : new Error(String(error)); - const clientAborted = session.clientAbortSignal?.aborted ?? false; + const clientAborted = + passthroughPump.wasClientAborted() || (session.clientAbortSignal?.aborted ?? false); const isResponseControllerAborted = sessionWithController.responseController?.signal.aborted ?? false; const isIdleTimeout = !!err.message?.includes("streaming_idle"); @@ -2362,6 +2888,10 @@ export class ProxyResponseHandler { }); try { + if (terminalFinalizationStarted) { + throw err; + } + terminalFinalizationStarted = true; clearIdleTimer(); const allContent = flushAndJoin(); const duration = Date.now() - session.startTime; @@ -2374,104 +2904,71 @@ export class ProxyResponseHandler { clientAborted, abortReason ); + latestCommitSideEffects = finalized.commitSideEffects; await finalizeRequestStats( session, allContent, finalized.effectiveStatusCode, duration, - finalized.errorMessage ?? abortReason, + finalized.errorMessage ?? undefined, finalized.providerIdForPersistence ?? undefined, - true // 流式透传错误兜底也是流式上下文 + true, // 流式透传错误兜底也是流式上下文 + () => scheduleCommitSideEffects(latestCommitSideEffects) ); } catch (finalizeError) { + const fallbackStatusCode = + statusCode >= 400 + ? statusCode + : streamEndedNormally + ? 500 + : clientAborted + ? 499 + : 502; await persistRequestFailure({ session, messageContext, - statusCode: statusCode && statusCode >= 400 ? statusCode : 502, + statusCode: fallbackStatusCode, error: finalizeError, taskId, phase: "stream", + detailsWriter: updateMessageRequestDetailsIfUnfinalized, + onCommitted: () => scheduleCommitSideEffects(latestCommitSideEffects), + awaitPersistence: (promise: Promise) => + awaitTerminalPersistenceWithOwnership({ + promise, + taskId, + operation: "gemini-stream-fallback", + timeoutMs: STREAM_FAILURE_PERSISTENCE_MAX_MS, + }), }); } } finally { - cleanupTaskAbortBinding(); - clearIdleTimer(); - // 兜底:在流结束/中断后清理首字节超时,避免定时器泄漏 - // 注意:不应在流仍可能继续时清理(否则会让首字节超时失效) - try { - const wasResponseControllerAborted = - sessionWithController.responseController?.signal.aborted ?? false; - const clientAborted = session.clientAbortSignal?.aborted ?? false; - const shouldClearTimeout = - responseTimeoutCleared || - streamEndedNormally || - wasResponseControllerAborted || - clientAborted; - if (shouldClearTimeout) { - clearResponseTimeoutOnce(); - } - } catch (e) { - logger.warn( - "[ResponseHandler] Gemini passthrough: Failed to clear response timeout", - { - taskId, - providerId: provider.id, - providerName: provider.name, - error: e instanceof Error ? e.message : String(e), - } - ); - } - try { - // 取消 tee 分支,避免 stats 任务提前退出时 backpressure 影响客户端透传 - const cancelPromise = reader?.cancel(); - if (cancelPromise) { - cancelPromise.catch((err) => { - logger.warn( - "[ResponseHandler] Gemini passthrough: Failed to cancel stats reader", - { - taskId, - providerId: provider.id, - providerName: provider.name, - error: err instanceof Error ? err.message : String(err), - } - ); - }); - } - } catch (e) { - logger.warn("[ResponseHandler] Gemini passthrough: Failed to cancel stats reader", { - taskId, - providerId: provider.id, - providerName: provider.name, - error: e instanceof Error ? e.message : String(e), - }); - } - try { - // 取消 reader lock - reader?.releaseLock(); - } catch (e) { - logger.warn("[ResponseHandler] Gemini passthrough: Failed to release reader lock", { - taskId, - providerId: provider.id, - providerName: provider.name, - error: e instanceof Error ? e.message : String(e), - }); - } - releaseSessionAgent(session); + releaseTransportResources(); } - })(); + }; - AsyncTaskManager.register(taskId, statsPromise, { - taskType: "stream-passthrough-stats", - abortController: statsAbortController, - staleTimeoutMs: streamTaskStaleTimeoutMs, - }); - statsPromise.catch((error) => { - if (session.sessionId && session.shouldPersistSessionDebugArtifacts()) { - void discardBeforeResponseBodySnapshot(session); + AsyncTaskManager.register( + taskId, + () => { + const statsPromise = runStatsTask(); + statsPromise.catch((error) => { + if (session.sessionId && session.shouldPersistSessionDebugArtifacts()) { + void discardBeforeResponseBodySnapshot(session); + } + logger.error( + "[ResponseHandler] Gemini passthrough stats task uncaught error:", + error + ); + }); + return statsPromise; + }, + { + taskType: "stream-passthrough-stats", + abortController: statsAbortController, + staleTimeoutMs: streamTaskStaleTimeoutMs, } - logger.error("[ResponseHandler] Gemini passthrough stats task uncaught error:", error); - }); + ); if (session.sessionId && session.shouldPersistSessionDebugArtifacts()) { const responseAfterMetaTask = SessionManager.storeSessionResponsePhaseSnapshot?.( @@ -2491,8 +2988,11 @@ export class ProxyResponseHandler { }); } - discardBeforeResponseBodySnapshot(session); - return response; + return new Response(passthroughPump.stream, { + status: response.status, + statusText: response.statusText, + headers: cleanResponseHeaders(response.headers), + }); } else { // ❌ 需要转换:客户端不是 Gemini 格式(如 OpenAI/Claude) logger.debug("[ResponseHandler] Transforming Gemini stream to client format", { @@ -2544,21 +3044,6 @@ export class ProxyResponseHandler { } } - // 使用 TransformStream 包装流,以便在 idle timeout 时能关闭客户端流 - // 这解决了 tee() 后 internalStream abort 不影响 clientStream 的问题 - let streamController: TransformStreamDefaultController | null = null; - const controllableStream = processedStream.pipeThrough( - new TransformStream({ - start(controller) { - streamController = controller; // 保存 controller 引用 - }, - transform(chunk, controller) { - controller.enqueue(chunk); // 透传数据 - }, - }) - ); - - const [clientStream, internalStream] = controllableStream.tee(); const statusCode = response.status; // 使用 AsyncTaskManager 管理后台处理任务 @@ -2573,8 +3058,9 @@ export class ProxyResponseHandler { provider.streamingIdleTimeoutMs > 0 ? provider.streamingIdleTimeoutMs : Number.POSITIVE_INFINITY; - const streamTaskStaleTimeoutMs = resolveStreamTaskStaleTimeoutMs(provider); + const streamTaskStaleTimeoutMs = resolveStreamTaskStaleTimeoutMs(); const clientAbortDrainTimeoutMs = CLIENT_ABORT_DRAIN_MAX_MS; + let responsePump: DemandDrivenResponsePump | null = null; // 提升 idleTimeoutId 到外部作用域,以便客户端断开时能清除 let idleTimeoutId: NodeJS.Timeout | null = null; @@ -2608,8 +3094,11 @@ export class ProxyResponseHandler { // 1. 关闭客户端流(让客户端收到连接关闭通知,避免悬挂) try { - if (streamController) { - streamController.error(new Error("Streaming idle timeout")); + if (responsePump) { + const idleTimeoutError = new Error("streaming_idle"); + idleTimeoutError.name = "AbortError"; + responsePump.errorClient(idleTimeoutError); + responsePump.cancelSource(idleTimeoutError); logger.debug("ResponseHandler: Client stream closed due to idle timeout", { taskId, providerId: provider.id, @@ -2647,16 +3136,21 @@ export class ProxyResponseHandler { abortController.abort(new Error("streaming_idle")); }, idleTimeoutMs); }; - const cleanupClientAbortListener = bindClientAbortListener(session.clientAbortSignal, () => { + let cleanupClientAbortListener = () => {}; + let clientDetachHandled = false; + const handleClientAbort = (reason?: unknown) => { + if (responsePump?.getState() === "closed") return; + responsePump?.startDrain(reason ?? "client_detached"); + if (clientDetachHandled) return; + clientDetachHandled = true; logger.debug("ResponseHandler: Client disconnected, cleaning up", { taskId, providerId: provider.id, messageId: messageContext.id, }); - // Do not cancel internal accounting on pure client disconnect. If the - // upstream stream has already completed, the tee'd internal branch can - // still drain buffered final usage and record the request as successful. - // Idle/response timeout paths still abort via abortController. + // Do not cancel internal accounting on pure client disconnect. Transfer + // ownership to the bounded background drain so terminal usage can still + // be recorded. Idle/response timeout paths still abort upstream. clearClientAbortDrainTimer(); if (!idleTimeoutId) { startIdleTimer(); @@ -2682,34 +3176,113 @@ export class ProxyResponseHandler { }); } - abortController.abort(new Error("client_abort_drain_timeout")); - }, clientAbortDrainTimeoutMs); - }); + const drainTimeoutError = new Error("client_abort_drain_timeout"); + abortController.abort(drainTimeoutError); + responsePump?.cancelSource(drainTimeoutError); + }, clientAbortDrainTimeoutMs); + }; + + // 统计/结算只保留有界的“头 + 尾”文本快照,避免长流式响应把进程堆撑满。 + let usageForCost: UsageMetrics | null = null; + let isFirstChunk = true; // 标记是否为第一块数据 + + // 不在首次读取前启动 idle timer(避免与首字节超时职责重叠) + // idle timer 仅在首块数据到达后启动,用于检测流中途静默。 + // 客户端断开后例外:后台 drain 也会启动 idle timer,避免 pre-body + // 静默一直等到 60s drain 总上限。 + + const flushAndJoin = (): string => { + const snapshot = streamTextAccumulator.finish(); + lastStreamTextSnapshot = snapshot; + return snapshot.text; + }; + + let responseTimeoutCleared = false; + const clearResponseTimeoutOnce = (): boolean => { + if (responseTimeoutCleared) return false; + const sessionWithCleanup = session as typeof session & { + clearResponseTimeout?: () => void; + }; + if (!sessionWithCleanup.clearResponseTimeout) return false; + responseTimeoutCleared = true; + sessionWithCleanup.clearResponseTimeout(); + return true; + }; - const processingPromise = (async () => { - const reader = internalStream.getReader(); - // 统计/结算只保留有界的“头 + 尾”文本快照,避免长流式响应把进程堆撑满。 - let usageForCost: UsageMetrics | null = null; - let isFirstChunk = true; // 标记是否为第一块数据 - - // 不在首次读取前启动 idle timer(避免与首字节超时职责重叠) - // idle timer 仅在首块数据到达后启动,用于检测流中途静默。 - // 客户端断开后例外:后台 drain 也会启动 idle timer,避免 pre-body - // 静默一直等到 60s drain 总上限。 - - const flushAndJoin = (): string => { - const snapshot = streamTextAccumulator.finish(); - lastStreamTextSnapshot = snapshot; - return snapshot.text; + const getResponseControllerAbortError = (): Error | null => { + const sessionWithController = session as typeof session & { + responseController?: AbortController; }; + const signal = sessionWithController.responseController?.signal; + if (!signal?.aborted) return null; + if (signal.reason instanceof Error && isClientAbortError(signal.reason)) { + return signal.reason; + } - const finalizeStream = async ( - allContent: string, - streamEndedNormally: boolean, - clientAborted: boolean, - abortReason?: string - ): Promise => { - const finalized = await finalizeDeferredStreamingFinalizationIfNeeded( + const error = new Error( + signal.reason instanceof Error ? signal.reason.message : "Response timeout" + ); + error.name = "AbortError"; + return error; + }; + + let terminalDetailsPersisted = false; + let streamCommitSideEffectsScheduled = false; + let latestStreamCommitSideEffects: Array<() => Promise> = []; + const scheduleStreamCommitSideEffects = () => { + if (latestStreamCommitSideEffects.length === 0 || streamCommitSideEffectsScheduled) return; + streamCommitSideEffectsScheduled = true; + const committedEffects = [...latestStreamCommitSideEffects]; + return schedulePostTerminalSideEffects({ + taskId, + providerId: provider.id, + sessionId: session.sessionId, + commit: (signal) => runPostTerminalSideEffects(committedEffects, signal), + }); + }; + let streamFailurePersistencePromise: Promise | null = null; + const persistStreamFailureOnce = ( + options: Parameters[0] + ): Promise => { + if (terminalDetailsPersisted) return Promise.resolve(); + if (!streamFailurePersistencePromise) { + streamFailurePersistencePromise = persistRequestFailure({ + ...options, + detailsWriter: updateMessageRequestDetailsIfUnfinalized, + onCommitted: options.onCommitted ?? scheduleStreamCommitSideEffects, + awaitPersistence: (promise: Promise) => + awaitTerminalPersistenceWithOwnership({ + promise, + taskId, + operation: "stream-failure-fallback", + timeoutMs: STREAM_FAILURE_PERSISTENCE_MAX_MS, + }), + }) + .then(() => undefined) + .catch((error) => { + logger.error("ResponseHandler: Stream failure fallback threw", { + taskId, + messageId: messageContext.id, + error, + }); + }); + } + return streamFailurePersistencePromise ?? Promise.resolve(); + }; + + let streamFinalizationPromise: Promise | null = null; + const finalizeStream = ( + allContent: string, + streamEndedNormally: boolean, + clientAborted: boolean, + abortReason?: string + ): Promise => { + if (streamFinalizationPromise) return streamFinalizationPromise; + streamFinalizationPromise = (async () => { + const finalizationDeadlineAtMs = Date.now() + STREAM_FINALIZATION_MAX_MS; + const awaitFinalization = (promise: Promise): Promise => + raceWithDeadline(promise, finalizationDeadlineAtMs, "stream_finalization_timeout"); + const finalized = finalizeDeferredStreamingFinalizationIfNeeded( session, allContent, statusCode, @@ -2717,6 +3290,9 @@ export class ProxyResponseHandler { clientAborted, abortReason ); + latestStreamCommitSideEffects = finalized.commitSideEffects + ? [finalized.commitSideEffects] + : []; const effectiveStatusCode = finalized.effectiveStatusCode; const streamErrorMessage = finalized.errorMessage; const providerIdForPersistence = finalized.providerIdForPersistence; @@ -2729,7 +3305,7 @@ export class ProxyResponseHandler { session.shouldPersistSessionDebugArtifacts() && !streamSnapshot?.truncated ) { - const beforeBody = (await consumeBeforeResponseBodySnapshot(session)) ?? allContent; + const beforeBody = allContent; void SessionManager.storeSessionResponse( session.sessionId, allContent, @@ -2770,7 +3346,6 @@ export class ProxyResponseHandler { } const duration = Date.now() - session.startTime; - await updateMessageRequestDuration(messageContext.id, duration); const tracker = ProxyStatusTracker.getInstance(); tracker.endRequest(messageContext.user.id, messageContext.id); @@ -2783,9 +3358,8 @@ export class ProxyResponseHandler { usageForCost = usageResult.usageMetrics; const actualServiceTier = parseServiceTierFromResponseText(allContent); - const codexPriorityBillingDecision = await resolveCodexPriorityBillingDecision( - session, - actualServiceTier + const codexPriorityBillingDecision = await awaitFinalization( + resolveCodexPriorityBillingDecision(session, actualServiceTier) ); if (!isNonBillingUsageEndpoint(session)) { ensureCodexServiceTierResultSpecialSetting(session, codexPriorityBillingDecision); @@ -2802,8 +3376,21 @@ export class ProxyResponseHandler { maybeSetCodexContext1m(session, provider, usageForCost?.input_tokens); - // Codex: Extract prompt_cache_key from SSE events and update session binding - if (provider.providerType === "codex" && session.sessionId && provider.id) { + let codexCacheBinding: + | { + sessionId: string; + promptCacheKey: string; + providerId: number; + keyId: number | null; + } + | undefined; + if ( + provider.providerType === "codex" && + effectiveStatusCode >= 200 && + effectiveStatusCode < 300 && + session.sessionId && + provider.id + ) { try { const sseEvents = parseSSEData(allContent); for (const event of sseEvents) { @@ -2812,14 +3399,12 @@ export class ProxyResponseHandler { event.data as Record ); if (promptCacheKey) { - void SessionManager.updateSessionWithCodexCacheKey( - session.sessionId, + codexCacheBinding = { + sessionId: session.sessionId, promptCacheKey, - provider.id, - session.authState?.key?.id ?? session.messageContext?.key?.id ?? null - ).catch((err) => { - logger.error("[ResponseHandler] Failed to update Codex session (stream):", err); - }); + providerId: provider.id, + keyId: session.authState?.key?.id ?? session.messageContext?.key?.id ?? null, + }; break; // Only need first prompt_cache_key } } @@ -2829,36 +3414,42 @@ export class ProxyResponseHandler { } } - const billableUsageForCost = await resolveBillableUsageMetricsForCost( - session, - provider, - usageForCost, - effectiveStatusCode, - allContent + const billableUsageForCost = await awaitFinalization( + resolveBillableUsageMetricsForCost( + session, + provider, + usageForCost, + effectiveStatusCode, + allContent + ) ); const billing = sessionBillingInputs(session, provider, priorityServiceTierApplied); - const costUpdateResult = await updateRequestCostFromUsage( - messageContext.id, - session, - billableUsageForCost, - billing, - // Any hedge-path winner with loser billing on uses the loser-sum-aware write. - // Gate on billHedgeLosers (not the racy isHedgeWinner/launchedProviderCount): - // an alternative can still be mid-launch when the initial provider commits, so - // isHedgeWinner may read false even though a loser will bill — using it would - // let the winner's replacement clobber that loser's additive write. - finalized.billHedgeLosers + const costUpdateResult = await awaitFinalization( + updateRequestCostFromUsage( + messageContext.id, + session, + billableUsageForCost, + billing, + // Any hedge-path winner with loser billing on uses the loser-sum-aware write. + // Gate on billHedgeLosers (not the racy isHedgeWinner/launchedProviderCount): + // an alternative can still be mid-launch when the initial provider commits, so + // isHedgeWinner may read false even though a loser will bill — using it would + // let the winner's replacement clobber that loser's additive write. + finalized.billHedgeLosers + ) ); if (costUpdateResult.longContextPricingApplied) { ensureLongContextPricingAudit(session, costUpdateResult.longContextPricing); } // 追踪消费到 Redis(用于限流) - await trackCostToRedis(session, billableUsageForCost, billing, { - resolvedPricing: costUpdateResult.resolvedPricing, - longContextPricing: costUpdateResult.longContextPricing, - }); + await awaitFinalization( + trackCostToRedis(session, billableUsageForCost, billing, { + resolvedPricing: costUpdateResult.resolvedPricing, + longContextPricing: costUpdateResult.longContextPricing, + }) + ); // Calculate cost for session tracking (with multiplier) and Langfuse (raw) let costUsdStr: string | undefined; @@ -2867,7 +3458,9 @@ export class ProxyResponseHandler { if (billableUsageForCost) { try { if (session.request.model) { - const resolvedPricing = await session.getResolvedPricingByBillingSource(provider); + const resolvedPricing = await awaitFinalization( + session.getResolvedPricingByBillingSource(provider) + ); if (resolvedPricing) { ensurePricingResolutionSpecialSetting(session, resolvedPricing); const longContextPricing = @@ -2980,26 +3573,53 @@ export class ProxyResponseHandler { ? anthropicModelDetection.actualResponseModel : extractActualResponseModelForProvider(provider.providerType, true, allContent); + const postTerminalSideEffects = [...latestStreamCommitSideEffects]; + if (codexCacheBinding) { + const { sessionId, promptCacheKey, providerId, keyId } = codexCacheBinding; + postTerminalSideEffects.push(async () => { + try { + await SessionManager.updateSessionWithCodexCacheKey( + sessionId, + promptCacheKey, + providerId, + keyId + ); + } catch (err) { + logger.error("[ResponseHandler] Failed to update Codex session (stream):", err); + } + }); + } + latestStreamCommitSideEffects = postTerminalSideEffects; + // 保存扩展信息(status code, tokens, provider chain) - await updateMessageRequestDetails(messageContext.id, { - statusCode: effectiveStatusCode, - inputTokens: usageForCost?.input_tokens, - outputTokens: usageForCost?.output_tokens, - ttfbMs: session.ttfbMs, - cacheCreationInputTokens: usageForCost?.cache_creation_input_tokens, - cacheReadInputTokens: usageForCost?.cache_read_input_tokens, - cacheCreation5mInputTokens: usageForCost?.cache_creation_5m_input_tokens, - cacheCreation1hInputTokens: usageForCost?.cache_creation_1h_input_tokens, - cacheTtlApplied: usageForCost?.cache_ttl ?? null, - providerChain: session.getProviderChain(), - ...(streamErrorMessage ? { errorMessage: streamErrorMessage } : {}), - model: currentRequestedModel ?? undefined, // 更新重定向后的模型 - actualResponseModel: finalActualResponseModel, - providerId: providerIdForPersistence ?? session.provider?.id, // 更新最终供应商ID(重试切换后) - context1mApplied: session.getContext1mApplied(), - swapCacheTtlApplied: provider.swapCacheTtlBilling ?? false, - specialSettings: session.getSpecialSettings() ?? undefined, - }); + terminalDetailsPersisted = await awaitFinalization( + updateMessageRequestDetailsDurably( + messageContext.id, + { + statusCode: effectiveStatusCode, + durationMs: duration, + inputTokens: usageForCost?.input_tokens, + outputTokens: usageForCost?.output_tokens, + ttfbMs: session.ttfbMs, + cacheCreationInputTokens: usageForCost?.cache_creation_input_tokens, + cacheReadInputTokens: usageForCost?.cache_read_input_tokens, + cacheCreation5mInputTokens: usageForCost?.cache_creation_5m_input_tokens, + cacheCreation1hInputTokens: usageForCost?.cache_creation_1h_input_tokens, + cacheTtlApplied: usageForCost?.cache_ttl ?? null, + providerChain: session.getProviderChain(), + ...(streamErrorMessage ? { errorMessage: streamErrorMessage } : {}), + model: currentRequestedModel ?? undefined, // 更新重定向后的模型 + actualResponseModel: finalActualResponseModel, + providerId: providerIdForPersistence ?? session.provider?.id, // 更新最终供应商ID(重试切换后) + context1mApplied: session.getContext1mApplied(), + swapCacheTtlApplied: provider.swapCacheTtlBilling ?? false, + specialSettings: session.getSpecialSettings() ?? undefined, + }, + { + onCommitted: scheduleStreamCommitSideEffects, + } + ) + ); emitProxyLangfuseTrace(session, { responseHeaders: response.headers, @@ -3013,65 +3633,93 @@ export class ProxyResponseHandler { sseEventCount: getCollectedChunkCount(), errorMessage: streamErrorMessage ?? undefined, }); - }; + })(); + return streamFinalizationPromise; + }; - try { - let streamEndedNormally = false; - while (true) { - // 检查取消信号 - if (abortController.signal.aborted) { - logger.info("ResponseHandler: Stream processing cancelled", { - taskId, - providerId: provider.id, - providerName: provider.name, - chunksCollected: getCollectedChunkCount(), - }); - break; // 提前终止 - } + const observeChunk = (value: Uint8Array) => { + const chunkSize = value.length; + clearIdleTimer(); + streamTextAccumulator.pushBytes(value); + AsyncTaskManager.touch(taskId); - const { value, done } = await reader.read(); - if (done) { - streamEndedNormally = true; - break; - } - if (value) { - const chunkSize = value.length; - streamTextAccumulator.pushBytes(value); - AsyncTaskManager.touch(taskId); + logger.trace("ResponseHandler: Upstream stream chunk received", { + taskId, + providerId: provider.id, + chunksCollected: getCollectedChunkCount(), + lastChunkSize: chunkSize, + idleTimeoutMs: idleTimeoutMs === Infinity ? "disabled" : idleTimeoutMs, + }); - // 每次收到数据后重置静默期计时器(首次收到数据时启动) - startIdleTimer(); - logger.trace("ResponseHandler: Idle timer reset (data received)", { - taskId, - providerId: provider.id, - chunksCollected: getCollectedChunkCount(), - lastChunkSize: chunkSize, - idleTimeoutMs: idleTimeoutMs === Infinity ? "disabled" : idleTimeoutMs, - }); + if (isFirstChunk) { + session.recordTtfb(); + isFirstChunk = false; + if (clearResponseTimeoutOnce()) { + logger.debug("ResponseHandler: First chunk received, response timeout cleared", { + taskId, + providerId: provider.id, + firstChunkSize: chunkSize, + }); + } + } + }; - // 流式:读到第一块数据后立即清除响应超时定时器 - if (isFirstChunk) { - session.recordTtfb(); - isFirstChunk = false; - const sessionWithCleanup = session as typeof session & { - clearResponseTimeout?: () => void; - }; - if (sessionWithCleanup.clearResponseTimeout) { - sessionWithCleanup.clearResponseTimeout(); - logger.debug("ResponseHandler: First chunk received, response timeout cleared", { - taskId, - providerId: provider.id, - firstChunkSize: chunkSize, - }); - } - } - } + responsePump = createDemandDrivenResponsePump({ + source: processedStream, + onReadStart() { + // A pending chunk is deliberately not considered Provider idle. The + // pump invokes this only when it actually starts the next source read. + if (!isFirstChunk) { + startIdleTimer(); + } + }, + onChunk: observeChunk, + onClientCancel: handleClientAbort, + }); + const activeResponsePump = responsePump; + const cleanupResponseControllerAbortListener = bindClientAbortListener( + ( + session as typeof session & { + responseController?: AbortController; + } + ).responseController?.signal, + () => { + const responseControllerAbortError = getResponseControllerAbortError(); + if (responseControllerAbortError) { + activeResponsePump.errorClient(responseControllerAbortError); + activeResponsePump.cancelSource(responseControllerAbortError); + } + } + ); + cleanupClientAbortListener = bindClientAbortListener(session.clientAbortSignal, () => + handleClientAbort(session.clientAbortSignal?.reason) + ); + + const runProcessingTask = async () => { + try { + const pumpCompletion = await activeResponsePump.completion; + cleanupTaskAbortBinding(); + releaseSessionAgent(session); + cleanupResponseControllerAbortListener(); + cleanupClientAbortListener(); + cleanupClientAbortListener = () => {}; + clearClientAbortDrainTimer(); + clearIdleTimer(); + clearResponseTimeoutOnce(); + const responseControllerAbortError = getResponseControllerAbortError(); + if (responseControllerAbortError) { + throw responseControllerAbortError; } + if (pumpCompletion.error) { + throw pumpCompletion.error; + } + const streamEndedNormally = + pumpCompletion.streamEndedNormally && !abortController.signal.aborted; // 流式读取完成:清除静默期计时器 clearIdleTimer(); const allContent = flushAndJoin(); - const clientAborted = session.clientAbortSignal?.aborted ?? false; + const clientAborted = pumpCompletion.clientAborted; try { await finalizeStream(allContent, streamEndedNormally, clientAborted); } catch (finalizeError) { @@ -3086,7 +3734,7 @@ export class ProxyResponseHandler { }); // 回退:避免 finalizeStream 失败导致 request record 未被更新 - await persistRequestFailure({ + await persistStreamFailureOnce({ session, messageContext, statusCode: statusCode && statusCode >= 400 ? statusCode : 500, @@ -3101,7 +3749,10 @@ export class ProxyResponseHandler { const sessionWithController = session as typeof session & { responseController?: AbortController; }; - const clientAborted = session.clientAbortSignal?.aborted ?? false; + const pumpClientAborted = activeResponsePump.wasClientAborted(); + // The pump records which terminal cause won. Reading the raw signal here + // would let a later client disconnect overwrite an earlier Provider timeout/error. + const clientAborted = pumpClientAborted; const isResponseControllerAborted = sessionWithController.responseController?.signal.aborted ?? false; @@ -3136,7 +3787,7 @@ export class ProxyResponseHandler { }); // 回退:至少保证 DB 记录能落下,避免 orphan record - await persistRequestFailure({ + await persistStreamFailureOnce({ session, messageContext, statusCode: statusCode && statusCode >= 400 ? statusCode : 502, @@ -3175,7 +3826,7 @@ export class ProxyResponseHandler { }); // 回退:至少保证 DB 记录能落下,避免 orphan record - await persistRequestFailure({ + await persistStreamFailureOnce({ session, messageContext, statusCode: statusCode && statusCode >= 400 ? statusCode : 502, @@ -3208,7 +3859,7 @@ export class ProxyResponseHandler { }); // 回退:至少保证 DB 记录能落下,避免 orphan record - await persistRequestFailure({ + await persistStreamFailureOnce({ session, messageContext, statusCode: 502, @@ -3240,9 +3891,49 @@ export class ProxyResponseHandler { messageId: messageContext.id, finalizeError, }); + await persistStreamFailureOnce({ + session, + messageContext, + statusCode: 499, + error: "CLIENT_ABORTED", + taskId, + phase: "stream", + }); } } } else if (isTransportError(err)) { + if (pumpClientAborted) { + logger.warn("ResponseHandler: Transport closed after client detached", { + taskId, + providerId: provider.id, + providerName: provider.name, + messageId: messageContext.id, + chunksCollected: getCollectedChunkCount(), + errorName: err.name, + errorCode: (err as NodeJS.ErrnoException).code, + }); + + try { + const allContent = flushAndJoin(); + await finalizeStream(allContent, false, true); + } catch (finalizeError) { + logger.error("ResponseHandler: Failed to finalize client-detached transport", { + taskId, + messageId: messageContext.id, + finalizeError, + }); + await persistStreamFailureOnce({ + session, + messageContext, + statusCode: 499, + error: "CLIENT_ABORTED", + taskId, + phase: "stream", + }); + } + return; + } + // 上游流传输错误(SocketError, ECONNRESET 等):与 upstream abort 相同处理 // 参见 #916 — controller.error(err) 传播的 transport error logger.error("ResponseHandler: Upstream stream transport error", { @@ -3266,7 +3957,7 @@ export class ProxyResponseHandler { finalizeError, }); - await persistRequestFailure({ + await persistStreamFailureOnce({ session, messageContext, statusCode: 502, @@ -3290,7 +3981,7 @@ export class ProxyResponseHandler { }); // 回退:至少保证 DB 记录能落下,避免 orphan record - await persistRequestFailure({ + await persistStreamFailureOnce({ session, messageContext, statusCode: statusCode && statusCode >= 400 ? statusCode : 500, @@ -3303,44 +3994,44 @@ export class ProxyResponseHandler { } finally { // 确保资源释放 cleanupTaskAbortBinding(); + cleanupResponseControllerAbortListener(); cleanupClientAbortListener(); clearClientAbortDrainTimer(); clearIdleTimer(); // 清除静默期计时器(防止泄漏) - try { - reader.releaseLock(); - } catch (releaseError) { - logger.warn("Failed to release reader lock", { - taskId, - releaseError, - }); - } releaseSessionAgent(session); } - })(); + }; // 注册任务并添加全局错误捕获 - AsyncTaskManager.register(taskId, processingPromise, { - taskType: "stream-processing", - abortController, - staleTimeoutMs: streamTaskStaleTimeoutMs, - }); - processingPromise.catch(async (error) => { - logger.error("ResponseHandler: Uncaught error in stream processing", { - taskId, - messageId: messageContext.id, - error, - }); + AsyncTaskManager.register( + taskId, + () => { + const processingPromise = runProcessingTask(); + return processingPromise.catch(async (error) => { + logger.error("ResponseHandler: Uncaught error in stream processing", { + taskId, + messageId: messageContext.id, + error, + }); - // 更新数据库记录(避免 orphan record) - await persistRequestFailure({ - session, - messageContext, - statusCode: statusCode && statusCode >= 400 ? statusCode : 500, - error, - taskId, - phase: "stream", - }); - }); + // 更新数据库记录(避免 orphan record) + await persistStreamFailureOnce({ + session, + messageContext, + statusCode: statusCode && statusCode >= 400 ? statusCode : 500, + error, + taskId, + phase: "stream", + }); + throw error; + }); + }, + { + taskType: "stream-processing", + abortController, + staleTimeoutMs: streamTaskStaleTimeoutMs, + } + ); // ⭐ 修复 Bun 运行时的 Transfer-Encoding 重复问题 // 清理上游的传输 headers,让 Response API 自动管理 @@ -3363,7 +4054,7 @@ export class ProxyResponseHandler { }); } - return new Response(clientStream, { + return new Response(activeResponsePump.stream, { status: response.status, statusText: response.statusText, headers: finalStreamHeaders, @@ -4161,6 +4852,8 @@ async function updateRequestCostFromUsage( */ export async function finalizeHedgeLoserBilling(params: { messageRequestId: number; + /** Original request timestamp for Redis rolling-window alignment. */ + messageRequestCreatedAtMs: number; /** Loser's session — used for pricing/multiplier resolution and Redis cost tracking. */ loserSession: ProxySession; provider: Provider; @@ -4193,6 +4886,7 @@ export async function finalizeHedgeLoserBilling(params: { }): Promise { const { messageRequestId, + messageRequestCreatedAtMs, loserSession, provider, attemptNumber, @@ -4310,7 +5004,11 @@ export async function finalizeHedgeLoserBilling(params: { priorityServiceTierApplied, groupCostMultiplier, }, - { resolvedPricing, longContextPricing } + { resolvedPricing, longContextPricing }, + { + eventId: `${messageRequestId}:hedge-loser:${provider.id}:${attemptNumber}`, + createdAtMs: messageRequestCreatedAtMs, + } ); logger.info("[HedgeLoserBilling] Billed hedge loser", { @@ -4355,7 +5053,8 @@ export async function finalizeRequestStats( * 导致 extractActualResponseModelForProvider 走 non-stream JSON.parse 失败 * - 如果不传则回退为 isSSEText 嗅探(仅兼容保留) */ - isStreaming?: boolean + isStreaming?: boolean, + onCommitted?: () => void | Promise ): Promise { const { messageContext, provider } = session; if (!provider || !messageContext) { @@ -4427,8 +5126,9 @@ export async function finalizeRequestStats( }); } - await updateMessageRequestDetails(messageContext.id, { + const terminalDetails = { statusCode: statusCode, + durationMs: duration, ...(errorMessage ? { errorMessage } : {}), ttfbMs: session.ttfbMs ?? duration, providerChain: session.getProviderChain(), @@ -4442,7 +5142,14 @@ export async function finalizeRequestStats( context1mApplied: session.getContext1mApplied(), swapCacheTtlApplied: session.provider?.swapCacheTtlBilling ?? false, specialSettings: session.getSpecialSettings() ?? undefined, - }); + }; + if (onCommitted) { + await updateMessageRequestDetailsDurably(messageContext.id, terminalDetails, { + onCommitted, + }); + } else { + await updateMessageRequestDetailsDurably(messageContext.id, terminalDetails); + } return null; } @@ -4530,8 +5237,9 @@ export async function finalizeRequestStats( } // 7. 更新请求详情 - await updateMessageRequestDetails(messageContext.id, { + const terminalDetails = { statusCode: statusCode, + durationMs: duration, inputTokens: normalizedUsage.input_tokens, outputTokens: normalizedUsage.output_tokens, ttfbMs: session.ttfbMs ?? duration, @@ -4552,7 +5260,12 @@ export async function finalizeRequestStats( context1mApplied: session.getContext1mApplied(), swapCacheTtlApplied: provider.swapCacheTtlBilling ?? false, specialSettings: session.getSpecialSettings() ?? undefined, - }); + }; + if (onCommitted) { + await updateMessageRequestDetailsDurably(messageContext.id, terminalDetails, { onCommitted }); + } else { + await updateMessageRequestDetailsDurably(messageContext.id, terminalDetails); + } if (session.sessionId && session.requestSequence != null) { if (session.shouldTrackSessionObservability()) { @@ -4579,6 +5292,11 @@ type BillingComputeInputs = { groupCostMultiplier: number; }; +type CostTrackingEventContext = { + eventId: string | number; + createdAtMs: number; +}; + function sessionBillingInputs( session: ProxySession, provider: Provider, @@ -4600,9 +5318,10 @@ async function trackCostToRedis( pricingOverrides?: { resolvedPricing?: Awaited> | null; longContextPricing?: ResolvedLongContextPricing | null; - } + }, + eventContext?: CostTrackingEventContext ): Promise { - if (!usage || !session.sessionId) return; + if (!usage) return; if (isNonBillingUsageEndpoint(session)) return; try { @@ -4611,7 +5330,11 @@ async function trackCostToRedis( const key = session.authState?.key; const user = session.authState?.user; - if (!messageContext || !provider || !key || !user) return; + if (!provider || !key || !user) return; + + const eventId = eventContext?.eventId ?? messageContext?.id; + const createdAtMs = eventContext?.createdAtMs ?? messageContext?.createdAt.getTime(); + if (eventId == null || createdAtMs == null || !Number.isFinite(createdAtMs)) return; const modelName = session.request.model; if (!modelName) return; @@ -4644,71 +5367,42 @@ async function trackCostToRedis( const costFloat = parseFloat(cost.toString()); // 追踪到 Redis(使用 session.sessionId) - await RateLimitService.trackCost( - key.id, - provider.id, - session.sessionId, // 直接使用 session.sessionId - costFloat, - { - userId: user.id, - key5hResetMode: key.limit5hResetMode, - keyResetTime: key.dailyResetTime, - keyResetMode: key.dailyResetMode, - provider5hResetMode: provider.limit5hResetMode, - providerResetTime: provider.dailyResetTime, - providerResetMode: provider.dailyResetMode, - user5hResetMode: user.limit5hResetMode, - requestId: messageContext.id, - createdAtMs: messageContext.createdAt.getTime(), - } - ); - - // 新增:追踪用户层每日消费 - await RateLimitService.trackUserDailyCost( - user.id, - costFloat, - user.dailyResetTime, - user.dailyResetMode, - { - requestId: messageContext.id, - createdAtMs: messageContext.createdAt.getTime(), - } - ); + await RateLimitService.trackCost(key.id, provider.id, session.sessionId ?? "", costFloat, { + userId: user.id, + key5hResetMode: key.limit5hResetMode, + keyResetTime: key.dailyResetTime, + keyResetMode: key.dailyResetMode, + provider5hResetMode: provider.limit5hResetMode, + providerResetTime: provider.dailyResetTime, + providerResetMode: provider.dailyResetMode, + user5hResetMode: user.limit5hResetMode, + userResetTime: user.dailyResetTime, + userResetMode: user.dailyResetMode, + requestId: eventId, + createdAtMs, + }); - // Decrement lease budgets for all windows (fire-and-forget) - void Promise.all([ - RateLimitService.decrementLeaseBudget(key.id, "key", "5h", costFloat, { - resetMode: key.limit5hResetMode, - }), - RateLimitService.decrementLeaseBudget(key.id, "key", "daily", costFloat, { - resetMode: key.dailyResetMode, - }), - RateLimitService.decrementLeaseBudget(key.id, "key", "weekly", costFloat), - RateLimitService.decrementLeaseBudget(key.id, "key", "monthly", costFloat), - RateLimitService.decrementLeaseBudget(user.id, "user", "5h", costFloat, { - resetMode: user.limit5hResetMode, - }), - RateLimitService.decrementLeaseBudget(user.id, "user", "daily", costFloat, { - resetMode: user.dailyResetMode, - }), - RateLimitService.decrementLeaseBudget(user.id, "user", "weekly", costFloat), - RateLimitService.decrementLeaseBudget(user.id, "user", "monthly", costFloat), - RateLimitService.decrementLeaseBudget(provider.id, "provider", "5h", costFloat, { - resetMode: provider.limit5hResetMode, - }), - RateLimitService.decrementLeaseBudget(provider.id, "provider", "daily", costFloat, { - resetMode: provider.dailyResetMode, - }), - RateLimitService.decrementLeaseBudget(provider.id, "provider", "weekly", costFloat), - RateLimitService.decrementLeaseBudget(provider.id, "provider", "monthly", costFloat), - ]).catch((error) => { - logger.warn("[ResponseHandler] Failed to decrement lease budgets:", { - error: error instanceof Error ? error.message : String(error), - }); + await RateLimitService.settleLeaseBudgets({ + requestId: eventId, + cost: costFloat, + entities: { + key: { + id: key.id, + resetModes: { "5h": key.limit5hResetMode, daily: key.dailyResetMode }, + }, + user: { + id: user.id, + resetModes: { "5h": user.limit5hResetMode, daily: user.dailyResetMode }, + }, + provider: { + id: provider.id, + resetModes: { "5h": provider.limit5hResetMode, daily: provider.dailyResetMode }, + }, + }, }); // 刷新 session 时间戳(滑动窗口) - if (session.shouldTrackSessionObservability()) { + if (session.sessionId && session.shouldTrackSessionObservability()) { void SessionTracker.refreshSession(session.sessionId, key.id, provider.id, user.id).catch( (error) => { logger.error("[ResponseHandler] Failed to refresh session tracker:", error); @@ -4722,6 +5416,41 @@ async function trackCostToRedis( } } +function buildProcessingErrorDetails(error: unknown): { + errorMessage: string; + errorStack?: string; + errorCause?: string; +} { + const databaseError = findSafeDatabaseError(error); + if (databaseError) { + return { errorMessage: databaseError.message }; + } + + const maxErrorStackLength = 8192; + const maxErrorCauseLength = 4096; + const errorMessage = formatProcessingError(error); + + let errorStack = error instanceof Error ? error.stack : undefined; + if (errorStack && errorStack.length > maxErrorStackLength) { + errorStack = `${errorStack.substring(0, maxErrorStackLength)}\n...[truncated]`; + } + + let errorCause: string | undefined; + if (error instanceof Error && (error as NodeJS.ErrnoException).cause) { + try { + const cause = (error as NodeJS.ErrnoException).cause; + errorCause = JSON.stringify(cause, Object.getOwnPropertyNames(cause as object)); + } catch { + errorCause = String((error as NodeJS.ErrnoException).cause); + } + if (errorCause && errorCause.length > maxErrorCauseLength) { + errorCause = `${errorCause.substring(0, maxErrorCauseLength)}...[truncated]`; + } + } + + return { errorMessage, errorStack, errorCause }; +} + /** * 持久化请求失败信息到数据库 * - 用于后台异步任务中的错误处理,确保 orphan records 被正确更新 @@ -4734,53 +5463,33 @@ async function persistRequestFailure(options: { error: unknown; taskId: string; phase: "stream" | "non-stream"; -}): Promise { + awaitPersistence?: (promise: Promise) => Promise; + detailsWriter?: typeof updateMessageRequestDetailsIfUnfinalized; + onCommitted?: () => void | Promise; +}): Promise { const { session, messageContext, statusCode, error, taskId, phase } = options; + const awaitPersistence = options.awaitPersistence ?? ((promise: Promise) => promise); + const detailsWriter = options.detailsWriter ?? updateMessageRequestDetailsIfUnfinalized; if (!messageContext) { logger.warn("ResponseHandler: Cannot persist failure without messageContext", { taskId, phase, }); - return; + return false; } const tracker = ProxyStatusTracker.getInstance(); - const errorMessage = formatProcessingError(error); + const { errorMessage, errorStack, errorCause } = buildProcessingErrorDetails(error); const duration = Date.now() - session.startTime; - - // 提取完整错误信息用于排查(限制长度防止异常大的错误信息) - const MAX_ERROR_STACK_LENGTH = 8192; // 8KB,足够容纳大多数堆栈信息 - const MAX_ERROR_CAUSE_LENGTH = 4096; // 4KB,足够容纳 JSON 序列化的错误原因 - - let errorStack = error instanceof Error ? error.stack : undefined; - if (errorStack && errorStack.length > MAX_ERROR_STACK_LENGTH) { - errorStack = `${errorStack.substring(0, MAX_ERROR_STACK_LENGTH)}\n...[truncated]`; - } - - let errorCause: string | undefined; - if (error instanceof Error && (error as NodeJS.ErrnoException).cause) { - try { - // 序列化错误原因链,保留所有属性 - const cause = (error as NodeJS.ErrnoException).cause; - errorCause = JSON.stringify(cause, Object.getOwnPropertyNames(cause as object)); - } catch { - // 如果序列化失败,使用简单字符串 - errorCause = String((error as NodeJS.ErrnoException).cause); - } - // 截断过长的错误原因 - if (errorCause && errorCause.length > MAX_ERROR_CAUSE_LENGTH) { - errorCause = `${errorCause.substring(0, MAX_ERROR_CAUSE_LENGTH)}...[truncated]`; - } - } + let committed = false; try { - // 更新请求持续时间 - await updateMessageRequestDuration(messageContext.id, duration); - - // 更新错误详情和 provider chain - await updateMessageRequestDetails(messageContext.id, { + // duration 与 terminal status 必须属于同一个 CAS patch,避免 ordinary + // metadata 在 overflow 或进程退出时丢失而留下永久 active 记录。 + const terminalDetails = { statusCode, + durationMs: duration, errorMessage, errorStack, errorCause, @@ -4791,7 +5500,11 @@ async function persistRequestFailure(options: { context1mApplied: session.getContext1mApplied(), swapCacheTtlApplied: session.provider?.swapCacheTtlBilling ?? false, specialSettings: session.getSpecialSettings() ?? undefined, - }); + }; + const persistence = options.onCommitted + ? detailsWriter(messageContext.id, terminalDetails, { onCommitted: options.onCommitted }) + : detailsWriter(messageContext.id, terminalDetails); + committed = Boolean(await awaitPersistence(persistence)); if (session.sessionId && session.requestSequence != null) { if (session.shouldTrackSessionObservability()) { @@ -4814,12 +5527,16 @@ async function persistRequestFailure(options: { } ); } catch (dbError) { + const databaseError = findSafeDatabaseError(dbError); logger.error("ResponseHandler: Failed to persist request failure", { taskId, phase, messageId: messageContext.id, error: errorMessage, - dbError, + databaseError: + databaseError?.message ?? (dbError instanceof Error ? dbError.message : String(dbError)), + databaseErrorCode: databaseError?.code, + databaseErrorPool: databaseError?.pool, }); } finally { // 确保无论数据库操作成功与否,都清理追踪状态 @@ -4846,6 +5563,7 @@ async function persistRequestFailure(options: { sseEventCount: phase === "stream" ? 0 : undefined, errorMessage, }); + return committed; } /** diff --git a/src/app/v1beta/[...route]/route.ts b/src/app/v1beta/[...route]/route.ts index 73b785b39..2c70283fe 100644 --- a/src/app/v1beta/[...route]/route.ts +++ b/src/app/v1beta/[...route]/route.ts @@ -4,6 +4,7 @@ import { handle } from "hono/vercel"; import { registerCors } from "@/app/v1/_lib/cors"; import { handleAvailableModels } from "@/app/v1/_lib/models/available-models"; import { handleProxyRequest } from "@/app/v1/_lib/proxy-handler"; +import { withDataDbScope } from "@/drizzle/db"; export const runtime = "nodejs"; @@ -19,10 +20,14 @@ app.get("/models", handleAvailableModels); // 格式检测会自动识别 Gemini 请求体中的 contents 字段 app.all("*", handleProxyRequest); -export const GET = handle(app); -export const POST = handle(app); -export const PUT = handle(app); -export const DELETE = handle(app); -export const PATCH = handle(app); -export const OPTIONS = handle(app); -export const HEAD = handle(app); +const routeHandler = withDataDbScope(handle(app)); + +export { + routeHandler as GET, + routeHandler as POST, + routeHandler as PUT, + routeHandler as DELETE, + routeHandler as PATCH, + routeHandler as OPTIONS, + routeHandler as HEAD, +}; diff --git a/src/drizzle/admitted-client.ts b/src/drizzle/admitted-client.ts new file mode 100644 index 000000000..e871c15e8 --- /dev/null +++ b/src/drizzle/admitted-client.ts @@ -0,0 +1,360 @@ +export const DB_POOL_ADMISSION_ERROR_CODE = "DB_POOL_ADMISSION_EXCEEDED"; + +export interface DbPoolAdmissionErrorDetails { + code: typeof DB_POOL_ADMISSION_ERROR_CODE; + pool: string; + maxOutstanding: number; + message: string; +} + +export interface SafeDatabaseErrorDetails { + kind: "admission" | "query"; + code?: string; + pool?: string; + maxOutstanding?: number; + message: string; +} + +export class DbPoolAdmissionError extends Error { + readonly code = DB_POOL_ADMISSION_ERROR_CODE; + + constructor( + readonly pool: string, + readonly maxOutstanding: number + ) { + super(`Database pool ${pool} exceeded ${maxOutstanding} outstanding operations`); + this.name = "DbPoolAdmissionError"; + } +} + +export function findDbPoolAdmissionError(error: unknown): DbPoolAdmissionErrorDetails | null { + const visited = new Set(); + let current: unknown = error; + + for (let depth = 0; depth < 8; depth += 1) { + if ((typeof current !== "object" && typeof current !== "function") || current === null) { + return null; + } + if (visited.has(current)) return null; + visited.add(current); + + const candidate = current as { + code?: unknown; + pool?: unknown; + maxOutstanding?: unknown; + message?: unknown; + cause?: unknown; + }; + if (candidate.code === DB_POOL_ADMISSION_ERROR_CODE) { + return { + code: DB_POOL_ADMISSION_ERROR_CODE, + pool: typeof candidate.pool === "string" ? candidate.pool : "unknown", + maxOutstanding: + typeof candidate.maxOutstanding === "number" ? candidate.maxOutstanding : -1, + message: + typeof candidate.message === "string" ? candidate.message : DB_POOL_ADMISSION_ERROR_CODE, + }; + } + current = candidate.cause; + } + + return null; +} + +export function isDbPoolAdmissionError(error: unknown): boolean { + return findDbPoolAdmissionError(error) !== null; +} + +export function findSafeDatabaseError(error: unknown): SafeDatabaseErrorDetails | null { + const admission = findDbPoolAdmissionError(error); + if (admission) { + return { + kind: "admission", + code: admission.code, + pool: admission.pool, + maxOutstanding: admission.maxOutstanding, + message: `Database pool admission exceeded (pool=${admission.pool}, maxOutstanding=${admission.maxOutstanding})`, + }; + } + + const visited = new Set(); + let current: unknown = error; + let databaseCode: string | undefined; + + for (let depth = 0; depth < 8; depth += 1) { + if ((typeof current !== "object" && typeof current !== "function") || current === null) { + break; + } + if (visited.has(current)) break; + visited.add(current); + + const candidate = current as { + name?: unknown; + code?: unknown; + query?: unknown; + params?: unknown; + cause?: unknown; + }; + if (databaseCode === undefined && typeof candidate.code === "string") { + databaseCode = candidate.code; + } + if ( + candidate.name === "DrizzleQueryError" || + (typeof candidate.query === "string" && Array.isArray(candidate.params)) + ) { + let cause = candidate.cause; + const causeVisited = new Set(); + for (let causeDepth = 0; causeDepth < 8; causeDepth += 1) { + if ((typeof cause !== "object" && typeof cause !== "function") || cause === null) break; + if (causeVisited.has(cause)) break; + causeVisited.add(cause); + const causeCandidate = cause as { code?: unknown; cause?: unknown }; + if (databaseCode === undefined && typeof causeCandidate.code === "string") { + databaseCode = causeCandidate.code; + break; + } + cause = causeCandidate.cause; + } + return { + kind: "query", + code: databaseCode, + message: "Database query failed", + }; + } + current = candidate.cause; + } + + return null; +} + +interface AdmittedClientOptions { + pool: string; + maxOutstanding: number; +} + +interface UnsafeAndBeginClient { + unsafe: (...args: unknown[]) => unknown; + begin: (...args: unknown[]) => unknown; +} + +function isPromiseLike(value: unknown): value is PromiseLike { + return ( + (typeof value === "object" || typeof value === "function") && + value !== null && + typeof (value as { then?: unknown }).then === "function" + ); +} + +function isAsyncIterable(value: unknown): value is AsyncIterable { + return ( + (typeof value === "object" || typeof value === "function") && + value !== null && + typeof (value as { [Symbol.asyncIterator]?: unknown })[Symbol.asyncIterator] === "function" + ); +} + +function wrapAsyncIterable( + iterable: AsyncIterable, + release: () => void +): AsyncIterable { + return { + [Symbol.asyncIterator]() { + const iterator = iterable[Symbol.asyncIterator](); + return { + async next(...args: [] | [undefined]) { + try { + const result = await iterator.next(...args); + if (result.done) release(); + return result; + } catch (error) { + release(); + throw error; + } + }, + async return(value?: unknown) { + try { + if (iterator.return) return await iterator.return(value); + return { done: true as const, value }; + } finally { + release(); + } + }, + async throw(error?: unknown) { + try { + if (iterator.throw) return await iterator.throw(error); + throw error; + } finally { + release(); + } + }, + }; + }, + }; +} + +function wrapPendingQuery(pending: unknown, release: () => void): unknown { + if (!isPromiseLike(pending)) { + release(); + return pending; + } + + let trackedPromise: Promise | null = null; + let proxy: object; + + const track = () => { + if (!trackedPromise) { + trackedPromise = Promise.resolve(pending).then( + (value) => { + release(); + return value; + }, + (error) => { + release(); + throw error; + } + ); + } + return trackedPromise; + }; + + proxy = new Proxy(pending as object, { + get(target, property) { + if (property === "then") return track().then.bind(track()); + if (property === "catch") return track().catch.bind(track()); + if (property === "finally") return track().finally.bind(track()); + + const value = Reflect.get(target, property, target); + if (typeof value !== "function") return value; + + return (...args: unknown[]) => { + let result: unknown; + try { + result = Reflect.apply(value, target, args); + } catch (error) { + release(); + throw error; + } + + if (property === "cancel") { + // postgres.js Query is lazy: observing an unexecuted query through + // then/catch would start it again after cancel(). A queued query is + // rejected synchronously by cancel(), so its token can be released + // immediately. An already-executed query can be observed safely. + const executed = (target as { executed?: unknown }).executed === true; + if (executed) { + void track().catch(() => undefined); + } else { + release(); + } + return result === target ? proxy : result; + } + + if (result === target) { + if (property === "execute" || property === "forEach") { + void track().catch(() => undefined); + } + return proxy; + } + if (isPromiseLike(result)) { + return Promise.resolve(result).then( + (resolved) => { + release(); + return resolved; + }, + (error) => { + release(); + throw error; + } + ); + } + if (isAsyncIterable(result)) return wrapAsyncIterable(result, release); + return result; + }; + }, + }); + + return proxy; +} + +export function createAdmittedSqlClient( + client: TClient, + options: AdmittedClientOptions +): TClient { + const rawClient = client as unknown as UnsafeAndBeginClient; + const originalUnsafe = rawClient.unsafe.bind(client); + const originalBegin = rawClient.begin.bind(client); + let outstanding = 0; + + const acquire = () => { + if (outstanding >= options.maxOutstanding) { + throw new DbPoolAdmissionError(options.pool, options.maxOutstanding); + } + outstanding += 1; + let released = false; + return () => { + if (released) return; + released = true; + outstanding -= 1; + }; + }; + + const admittedUnsafe = (...args: unknown[]) => { + const release = acquire(); + try { + return wrapPendingQuery(originalUnsafe(...args), release); + } catch (error) { + release(); + throw error; + } + }; + + const admittedBegin = (...args: unknown[]) => { + const release = acquire(); + let result: unknown; + try { + result = originalBegin(...args); + } catch (error) { + release(); + throw error; + } + + if (!isPromiseLike(result)) { + release(); + return result; + } + return Promise.resolve(result).then( + (value) => { + release(); + return value; + }, + (error) => { + release(); + throw error; + } + ); + }; + + return new Proxy(client, { + apply(target, thisArg, argArray) { + const release = acquire(); + try { + return wrapPendingQuery( + Reflect.apply( + target as unknown as (...args: unknown[]) => unknown, + thisArg, + argArray + ), + release + ); + } catch (error) { + release(); + throw error; + } + }, + get(target, property, receiver) { + if (property === "unsafe") return admittedUnsafe; + if (property === "begin") return admittedBegin; + return Reflect.get(target, property, receiver); + }, + }); +} diff --git a/src/drizzle/db.ts b/src/drizzle/db.ts index 4e3078a24..b5d4089de 100644 --- a/src/drizzle/db.ts +++ b/src/drizzle/db.ts @@ -1,45 +1,165 @@ -import 'server-only'; +import "server-only"; -import { drizzle, type PostgresJsDatabase } from 'drizzle-orm/postgres-js'; -import postgres from 'postgres'; -import { getEnvConfig } from '@/lib/config/env.schema'; -import * as schema from './schema'; +import { AsyncLocalStorage } from "node:async_hooks"; +import { drizzle, type PostgresJsDatabase } from "drizzle-orm/postgres-js"; +import postgres from "postgres"; +import { getEnvConfig } from "@/lib/config/env.schema"; +import { createAdmittedSqlClient } from "./admitted-client"; +import * as schema from "./schema"; -let dbInstance: PostgresJsDatabase | null = null; +type DbLane = "data" | "control" | "writer"; +type DbPoolLifecycleState = "open" | "closing" | "closed"; +type DatabaseInstance = PostgresJsDatabase; +type SqlClient = ReturnType; -function createDbInstance(): PostgresJsDatabase { +interface PoolBudget { + data: number; + control: number; + writer: number; +} + +interface PoolInstance { + client: SqlClient; + db: DatabaseInstance; +} + +const APPLICATION_NAMES: Record = { + data: "claude-code-hub:data", + control: "claude-code-hub:control", + writer: "claude-code-hub:writer", +}; +const MIN_OUTSTANDING_PER_POOL = 32; +const OUTSTANDING_PER_CONNECTION = 8; + +const globalForDbScope = globalThis as typeof globalThis & { + __CCH_DATA_DB_SCOPE__?: AsyncLocalStorage; +}; +const dataDbScope = globalForDbScope.__CCH_DATA_DB_SCOPE__ ?? new AsyncLocalStorage(); +globalForDbScope.__CCH_DATA_DB_SCOPE__ = dataDbScope; + +let poolInstances: Partial> = {}; +let closePromise: Promise | null = null; +let poolLifecycleState: DbPoolLifecycleState = "open"; + +function splitPoolBudget(total: number): PoolBudget { + if (total === 1) return { data: 0, control: 1, writer: 0 }; + if (total === 2) return { data: 1, control: 1, writer: 0 }; + + const writer = 1; + const control = Math.min(total - 2, Math.max(1, Math.round(total * 0.2))); + return { data: total - control - writer, control, writer }; +} + +function getPoolBudget(): PoolBudget { + const env = getEnvConfig(); + const defaultTotal = env.NODE_ENV === "production" ? 20 : 10; + return splitPoolBudget(env.DB_POOL_MAX ?? defaultTotal); +} + +function resolvePhysicalLane(lane: DbLane, budget: PoolBudget): DbLane { + if (budget[lane] > 0) return lane; + if (lane === "writer") return budget.control > 0 ? "control" : "data"; + return "control"; +} + +function createDbInstance(lane: DbLane, max: number): PoolInstance { const env = getEnvConfig(); const connectionString = env.DSN; if (!connectionString) { - throw new Error('DSN environment variable is not set'); + throw new Error("DSN environment variable is not set"); } - // postgres.js 默认 max=10,在高并发下容易出现查询排队 - // 这里采用“生产环境默认更大、同时可通过 env 覆盖”的策略,兼容单机与 k8s 多副本 - const defaultMax = env.NODE_ENV === 'production' ? 20 : 10; const client = postgres(connectionString, { - max: env.DB_POOL_MAX ?? defaultMax, + max, idle_timeout: env.DB_POOL_IDLE_TIMEOUT ?? 20, connect_timeout: env.DB_POOL_CONNECT_TIMEOUT ?? 10, + connection: { + application_name: APPLICATION_NAMES[lane], + statement_timeout: env.DB_STATEMENT_TIMEOUT_MS, + lock_timeout: env.DB_LOCK_TIMEOUT_MS, + }, }); - return drizzle(client, { schema }); + const admittedClient = createAdmittedSqlClient(client, { + pool: lane, + maxOutstanding: Math.max(MIN_OUTSTANDING_PER_POOL, max * OUTSTANDING_PER_CONNECTION), + }); + + return { + client, + db: drizzle(admittedClient, { schema }), + }; } -export function getDb(): PostgresJsDatabase { - if (!dbInstance) { - dbInstance = createDbInstance(); +function getPool(lane: DbLane): DatabaseInstance { + if (poolLifecycleState !== "open") { + throw new Error(`Database pools are ${poolLifecycleState}`); } - return dbInstance; + const budget = getPoolBudget(); + const physicalLane = resolvePhysicalLane(lane, budget); + const existing = poolInstances[physicalLane]; + if (existing) return existing.db; + + const created = createDbInstance(physicalLane, budget[physicalLane]); + poolInstances[physicalLane] = created; + return created.db; } -export const db = new Proxy({} as PostgresJsDatabase, { - get(_target, prop, receiver) { - const instance = getDb(); - const value = Reflect.get(instance, prop, receiver); +export function runWithDataDbScope(callback: () => T): T { + return dataDbScope.run("data", callback); +} + +export function withDataDbScope( + handler: (...args: TArgs) => TResult +): (...args: TArgs) => TResult { + return (...args) => runWithDataDbScope(() => handler(...args)); +} + +export function getDb(): DatabaseInstance { + return getPool(dataDbScope.getStore() === "data" ? "data" : "control"); +} + +export function getMessageWriterDb(): DatabaseInstance { + return getPool("writer"); +} + +export function closeDbPools(): Promise { + if (closePromise) return closePromise; + + poolLifecycleState = "closing"; + const pools = Object.values(poolInstances); + let resolveClose!: () => void; + let rejectClose!: (reason?: unknown) => void; + const publishedClosePromise = new Promise((resolve, reject) => { + resolveClose = resolve; + rejectClose = reject; + }); + closePromise = publishedClosePromise; - return typeof value === 'function' ? value.bind(instance) : value; + void (async () => { + try { + const results = await Promise.allSettled( + pools.map(({ client }) => client.end({ timeout: 5 })) + ); + const failure = results.find( + (result): result is PromiseRejectedResult => result.status === "rejected" + ); + if (failure) throw failure.reason; + } finally { + poolInstances = {}; + poolLifecycleState = "closed"; + } + })().then(resolveClose, rejectClose); + + return publishedClosePromise; +} + +export const db = new Proxy({} as DatabaseInstance, { + get(_target, property) { + const instance = getDb(); + const value = Reflect.get(instance, property, instance); + return typeof value === "function" ? value.bind(instance) : value; }, }); diff --git a/src/instrumentation.ts b/src/instrumentation.ts index 1767b1bfb..3da834224 100644 --- a/src/instrumentation.ts +++ b/src/instrumentation.ts @@ -3,6 +3,7 @@ * 在服务器启动时自动执行数据库迁移 */ +import { findSafeDatabaseError } from "@/drizzle/admitted-client"; import { startCacheCleanup } from "@/lib/cache/session-cache"; import { getBenignBrokenPipeCode } from "@/lib/lifecycle/benign-errors"; import { logger } from "@/lib/logger"; @@ -33,7 +34,8 @@ const instrumentationState = globalThis as unknown as { * 处理器,并在崩溃时尝试写入 Node 诊断报告(report.*.json)。 * * 必要的 node 启动参数(在 Dockerfile 中配置): - * --report-on-fatalerror --report-uncaught-exception --report-directory=/app/reports + * --report-on-fatalerror --report-uncaught-exception --report-exclude-env + * --report-directory=/app/reports * * 这两个 process.on(...) 不会与现有的 SIGTERM / SIGINT 处理器冲突。 */ @@ -43,15 +45,27 @@ export function registerCrashDiagnostics(): void { } instrumentationState.__CCH_CRASH_HANDLERS_REGISTERED__ = true; + const toSafeCrashError = (error: unknown): Error => { + const databaseError = findSafeDatabaseError(error); + if (databaseError) return new Error(databaseError.message); + return error instanceof Error ? error : new Error(String(error)); + }; + const writeReport = (trigger: string, err: unknown): string | undefined => { try { const report = ( process as NodeJS.Process & { - report?: { writeReport: (filename?: string, err?: unknown) => string }; + report?: { + excludeEnv?: boolean; + writeReport: (filename?: string, err?: unknown) => string; + }; } ).report; if (report?.writeReport) { - return report.writeReport(`report.${trigger}.${Date.now()}.json`, err as Error); + // Diagnostic reports are durable artifacts. Exclude the entire + // environment rather than maintaining an incomplete secret allowlist. + report.excludeEnv = true; + return report.writeReport(`report.${trigger}.${Date.now()}.json`, toSafeCrashError(err)); } } catch { // 写入诊断报告失败不应再抛出 @@ -95,12 +109,13 @@ export function registerCrashDiagnostics(): void { return; } - const reportPath = writeReport("uncaughtException", err); - writeFatalStderr("uncaughtException", err, reportPath); + const safeError = toSafeCrashError(err); + const reportPath = writeReport("uncaughtException", safeError); + writeFatalStderr("uncaughtException", safeError, reportPath); logger.fatal("[Lifecycle] uncaughtException", { - error: err.message, - errorName: err.name, - stack: err.stack, + error: safeError.message, + errorName: safeError.name, + stack: safeError.stack, reportPath, }); // 与 Node 默认行为一致:捕获后退出,避免后续运行在不一致状态 @@ -126,7 +141,7 @@ export function registerCrashDiagnostics(): void { return; } - const err = reason instanceof Error ? reason : new Error(String(reason)); + const err = toSafeCrashError(reason); const reportPath = writeReport("unhandledRejection", err); writeFatalStderr("unhandledRejection", err, reportPath); logger.fatal("[Lifecycle] unhandledRejection", { @@ -344,12 +359,17 @@ export async function register() { } // Ledger backfill: fire-and-forget after migration (non-blocking, idempotent) - import("@/lib/ledger-backfill") - .then(({ backfillUsageLedger }) => - backfillUsageLedger().then((result) => { - logger.info("[Instrumentation] Ledger backfill complete", result); - }) - ) + Promise.all([import("@/lib/async-task-manager"), import("@/lib/ledger-backfill")]) + .then(([{ AsyncTaskManager }, { backfillUsageLedger }]) => { + AsyncTaskManager.register( + "startup-ledger-backfill", + async (signal) => { + const result = await backfillUsageLedger(signal); + logger.info("[Instrumentation] Ledger backfill complete", result); + }, + { taskType: "startup-ledger-backfill", staleTimeoutMs: Number.POSITIVE_INFINITY } + ); + }) .catch((err) => { logger.warn("[Instrumentation] Ledger backfill failed (non-fatal)", { error: err instanceof Error ? err.message : String(err), @@ -429,6 +449,23 @@ export async function register() { // 初始化通知任务队列(如果启用) const { scheduleNotifications } = await import("@/lib/notification/notification-queue"); await scheduleNotifications(); + ( + globalThis as typeof globalThis & { + __CCH_STOP_BACKGROUND_QUEUES__?: () => Promise; + } + ).__CCH_STOP_BACKGROUND_QUEUES__ = async () => { + const [{ stopCleanupQueue }, { stopNotificationQueue }] = await Promise.all([ + import("@/lib/log-cleanup/cleanup-queue"), + import("@/lib/notification/notification-queue"), + ]); + const results = await Promise.allSettled([stopCleanupQueue(), stopNotificationQueue()]); + const failures = results.flatMap((result) => + result.status === "rejected" ? [result.reason] : [] + ); + if (failures.length > 0) { + throw new AggregateError(failures, "Failed to stop background queues"); + } + }; // 初始化智能探测调度器(如果启用) const { startProbeScheduler, isSmartProbingEnabled } = await import( @@ -493,12 +530,17 @@ export async function register() { await runMigrations(); // Ledger backfill: fire-and-forget after migration (non-blocking, idempotent) - import("@/lib/ledger-backfill") - .then(({ backfillUsageLedger }) => - backfillUsageLedger().then((result) => { - logger.info("[Instrumentation] Ledger backfill complete", result); - }) - ) + Promise.all([import("@/lib/async-task-manager"), import("@/lib/ledger-backfill")]) + .then(([{ AsyncTaskManager }, { backfillUsageLedger }]) => { + AsyncTaskManager.register( + "startup-ledger-backfill", + async (signal) => { + const result = await backfillUsageLedger(signal); + logger.info("[Instrumentation] Ledger backfill complete", result); + }, + { taskType: "startup-ledger-backfill", staleTimeoutMs: Number.POSITIVE_INFINITY } + ); + }) .catch((err) => { logger.warn("[Instrumentation] Ledger backfill failed (non-fatal)", { error: err instanceof Error ? err.message : String(err), diff --git a/src/lib/async-task-manager.ts b/src/lib/async-task-manager.ts index 912cf4140..4c5ae8c47 100644 --- a/src/lib/async-task-manager.ts +++ b/src/lib/async-task-manager.ts @@ -17,6 +17,7 @@ import { logger } from "./logger"; */ interface TaskInfo { + taskId: string; promise: Promise; abortController: AbortController; createdAt: number; @@ -31,11 +32,19 @@ interface RegisterTaskOptions { staleTimeoutMs?: number; } +type AsyncTaskFactory = (signal: AbortSignal) => Promise; + +type AsyncTaskLifecycleState = "open" | "draining" | "closed"; + const DEFAULT_STALE_TASK_TIMEOUT_MS = 10 * 60 * 1000; class AsyncTaskManagerClass { + // tasks 仅指向每个 taskId 的最新 generation;pendingTasks 跟踪所有尚未 settled 的 generation。 private tasks: Map = new Map(); + private pendingTasks: Set = new Set(); private cleanupInterval: NodeJS.Timeout | null = null; + private shutdownPromise: Promise | null = null; + private lifecycleState: AsyncTaskLifecycleState = "open"; // Lazily initialize Node-only hooks on first use to avoid side effects at import time. private initialized = false; @@ -65,7 +74,7 @@ class AsyncTaskManagerClass { // 耗尽路径(例如脚本类调用方未触发 SIGTERM)。 process.once("beforeExit", () => { logger.info("[AsyncTaskManager] beforeExit reached, cancelling remaining tasks", { - activeTaskCount: this.tasks.size, + activeTaskCount: this.pendingTasks.size, }); this.cleanupAll(); }); @@ -80,40 +89,47 @@ class AsyncTaskManagerClass { * 注册一个异步任务 * * @param taskId 任务唯一标识 - * @param promise 异步任务 Promise + * @param factory 通过 admission 后才启动的异步任务 factory * @param taskType 任务类型(用于日志) * @returns AbortController(可用于取消任务) */ register( taskId: string, - promise: Promise, + factory: AsyncTaskFactory, taskTypeOrOptions: string | RegisterTaskOptions = "unknown" ): AbortController { - this.initializeIfNeeded(); - const options = typeof taskTypeOrOptions === "string" ? { taskType: taskTypeOrOptions } : taskTypeOrOptions; const taskType = options.taskType ?? "unknown"; + const abortController = options.abortController ?? new AbortController(); - // 如果任务已存在,先取消旧任务 - const oldTaskInfo = this.tasks.get(taskId); - if (oldTaskInfo) { - logger.warn("[AsyncTaskManager] Task already exists, cancelling old task", { - taskId, - taskType, - }); - this.cancel(taskId); - this.cleanup(taskId, oldTaskInfo); + if ( + this.lifecycleState === "closed" || + (this.lifecycleState === "draining" && this.pendingTasks.size === 0) + ) { + abortController.abort(); + return abortController; } - const abortController = options.abortController ?? new AbortController(); + this.initializeIfNeeded(); + + const previousLatest = this.tasks.get(taskId); + const staleTimeoutMs = options.staleTimeoutMs === undefined || options.staleTimeoutMs <= 0 ? DEFAULT_STALE_TASK_TIMEOUT_MS : options.staleTimeoutMs; const now = Date.now(); + let resolveTask!: () => void; + let rejectTask!: (reason?: unknown) => void; + const promise = new Promise((resolve, reject) => { + resolveTask = resolve; + rejectTask = reject; + }); + const taskInfo: TaskInfo = { + taskId, promise, abortController, createdAt: now, @@ -123,6 +139,22 @@ class AsyncTaskManagerClass { }; this.tasks.set(taskId, taskInfo); + this.pendingTasks.add(taskInfo); + + if (previousLatest) { + logger.warn("[AsyncTaskManager] Task already exists, cancelling old task", { + taskId, + taskType, + }); + if (!previousLatest.abortController.signal.aborted) { + previousLatest.abortController.abort(); + } + logger.info("[AsyncTaskManager] Task cancelled", { + taskId, + taskType: previousLatest.taskType, + age: Date.now() - previousLatest.createdAt, + }); + } // 任务完成后自动清理 promise @@ -159,9 +191,19 @@ class AsyncTaskManagerClass { logger.debug("[AsyncTaskManager] Task registered", { taskId, taskType, - activeTasks: this.tasks.size, + activeTasks: this.pendingTasks.size, }); + if (abortController.signal.aborted) { + resolveTask(); + } else { + try { + Promise.resolve(factory(abortController.signal)).then(resolveTask, rejectTask); + } catch (error) { + rejectTask(error); + } + } + return abortController; } @@ -208,18 +250,19 @@ class AsyncTaskManagerClass { * @param taskId 任务唯一标识 */ private cleanup(taskId: string, expectedTask: TaskInfo): boolean { - if (this.tasks.get(taskId) !== expectedTask) { + if (!this.pendingTasks.delete(expectedTask)) { return false; } - const deleted = this.tasks.delete(taskId); - if (deleted) { - logger.debug("[AsyncTaskManager] Task cleaned up", { - taskId, - remainingTasks: this.tasks.size, - }); + if (this.tasks.get(taskId) === expectedTask) { + this.tasks.delete(taskId); } - return deleted; + + logger.debug("[AsyncTaskManager] Task cleaned up", { + taskId, + remainingTasks: this.pendingTasks.size, + }); + return true; } /** @@ -228,7 +271,7 @@ class AsyncTaskManagerClass { * 遍历所有活跃任务,对于空闲时间超过任务级 staleTimeoutMs 的任务: * 1. 记录警告日志 * 2. 触发 AbortController 取消任务 - * 3. 从任务 Map 中移除 + * 3. 保持 pending 跟踪,直到真实 Promise settled * * 注意:这是清理"空闲超时"的任务。活跃流应在收到上游 chunk 时 * 调用 touch() 更新 lastActivityAt,避免被误判为挂死任务。 @@ -236,23 +279,23 @@ class AsyncTaskManagerClass { private cleanupCompletedTasks(): void { const now = Date.now(); - for (const [taskId, taskInfo] of this.tasks.entries()) { + for (const taskInfo of this.pendingTasks) { + const { taskId } = taskInfo; const age = now - taskInfo.createdAt; const idleAge = now - taskInfo.lastActivityAt; const staleTimeoutMs = taskInfo.staleTimeoutMs || DEFAULT_STALE_TASK_TIMEOUT_MS; - // 如果任务超过阈值没有任何进展,记录警告、取消并从 Map 断开强引用。 - if (idleAge > staleTimeoutMs) { - logger.warn("[AsyncTaskManager] Task timeout, cancelling and detaching", { + // stale cleanup 只负责发出一次取消;settlement 才拥有移除 pending 跟踪的权限。 + if (idleAge > staleTimeoutMs && !taskInfo.abortController.signal.aborted) { + logger.warn("[AsyncTaskManager] Task timeout, cancelling", { taskId, taskType: taskInfo.taskType, age, idleAge, staleTimeoutMs, }); - this.cancel(taskId); - this.cleanup(taskId, taskInfo); + taskInfo.abortController.abort(); } } } @@ -261,13 +304,15 @@ class AsyncTaskManagerClass { * 清理所有任务(进程退出时调用) */ cleanupAll(): void { + this.lifecycleState = "closed"; logger.info("[AsyncTaskManager] Cleaning up all tasks", { - count: this.tasks.size, + count: this.pendingTasks.size, }); - for (const [taskId, taskInfo] of Array.from(this.tasks.entries())) { - this.cancel(taskId); - this.cleanup(taskId, taskInfo); + for (const taskInfo of Array.from(this.pendingTasks)) { + if (!taskInfo.abortController.signal.aborted) { + taskInfo.abortController.abort(); + } } if (this.cleanupInterval) { @@ -276,11 +321,66 @@ class AsyncTaskManagerClass { } } + /** + * 取消并等待 shutdown 时仍在飞的全部任务 settled。 + * + * task 的 finally 可能在等待期间注册尾部任务,因此循环到 pending 集合为空;并发 shutdown + * 调用共享同一个 Promise,避免重复取消或提前返回。 + */ + shutdownAll(): Promise { + if (this.shutdownPromise) { + return this.shutdownPromise; + } + + let resolveShutdown!: () => void; + let rejectShutdown!: (reason?: unknown) => void; + const shutdownPromise = new Promise((resolve, reject) => { + resolveShutdown = resolve; + rejectShutdown = reject; + }); + this.shutdownPromise = shutdownPromise; + this.lifecycleState = "draining"; + + // 先发布共享 Promise,再同步开始 abort;这样既保留既有同步取消语义, + // 同步 abort listener 重入时也会复用同一次 shutdown。 + void (async () => { + if (this.cleanupInterval) { + clearInterval(this.cleanupInterval); + this.cleanupInterval = null; + } + + while (true) { + if (this.pendingTasks.size === 0) { + this.lifecycleState = "closed"; + return; + } + + const activeTasks = Array.from(this.pendingTasks); + logger.info("[AsyncTaskManager] Cancelling and joining active tasks", { + count: activeTasks.length, + }); + + for (const taskInfo of activeTasks) { + if (!taskInfo.abortController.signal.aborted) { + taskInfo.abortController.abort(); + } + } + + await Promise.allSettled(activeTasks.map((taskInfo) => taskInfo.promise)); + } + })().then(resolveShutdown, (error) => { + this.lifecycleState = "closed"; + rejectShutdown(error); + }); + + return shutdownPromise; + } + /** * 获取当前活跃任务数 */ getActiveTaskCount(): number { - return this.tasks.size; + return this.pendingTasks.size; } /** @@ -288,8 +388,8 @@ class AsyncTaskManagerClass { */ getActiveTasks(): Array<{ taskId: string; taskType: string; age: number }> { const now = Date.now(); - return Array.from(this.tasks.entries()).map(([taskId, taskInfo]) => ({ - taskId, + return Array.from(this.pendingTasks).map((taskInfo) => ({ + taskId: taskInfo.taskId, taskType: taskInfo.taskType, age: now - taskInfo.createdAt, })); @@ -303,6 +403,6 @@ export const AsyncTaskManager = // 供 shutdown 编排器调用:在 cleanup 阶段(server.close 完成后)才取消残留任务, // 避免 drain 期间打断流式响应。 -export function shutdownAllAsyncTasks(): void { - AsyncTaskManager.cleanupAll(); +export function shutdownAllAsyncTasks(): Promise { + return AsyncTaskManager.shutdownAll(); } diff --git a/src/lib/config/env.schema.ts b/src/lib/config/env.schema.ts index d6998bfa1..80d56a986 100644 --- a/src/lib/config/env.schema.ts +++ b/src/lib/config/env.schema.ts @@ -39,7 +39,7 @@ export const EnvSchema = z.object({ }, z.string().url("数据库URL格式无效")), // PostgreSQL 连接池配置(postgres.js) // - 多副本部署(k8s)需要结合数据库 max_connections 分摊配置 - // - 这些值为“每个应用进程”的连接池上限 + // - DB_POOL_MAX 是每个应用进程内 data/control/writer 三类 pool 的连接总预算 DB_POOL_MAX: optionalNumber( z.number().int().min(1, "DB_POOL_MAX 不能小于 1").max(200, "DB_POOL_MAX 不能大于 200") ), @@ -57,6 +57,22 @@ export const EnvSchema = z.object({ .min(1, "DB_POOL_CONNECT_TIMEOUT 不能小于 1") .max(120, "DB_POOL_CONNECT_TIMEOUT 不能大于 120") ), + // 活动语句超时(毫秒),必须早于流式结算的 120 秒应用层 deadline + DB_STATEMENT_TIMEOUT_MS: optionalNumber( + z + .number() + .int() + .min(1000, "DB_STATEMENT_TIMEOUT_MS 不能小于 1000") + .max(119000, "DB_STATEMENT_TIMEOUT_MS 不能大于 119000") + ).default(90_000), + // 等待数据库锁的最长时间(毫秒) + DB_LOCK_TIMEOUT_MS: optionalNumber( + z + .number() + .int() + .min(100, "DB_LOCK_TIMEOUT_MS 不能小于 100") + .max(60000, "DB_LOCK_TIMEOUT_MS 不能大于 60000") + ).default(5_000), // message_request 写入模式 // - sync:同步写入(兼容旧行为,但高并发下会增加请求尾部阻塞) // - async:异步批量写入(默认,降低 DB 写放大与连接占用) @@ -108,6 +124,13 @@ export const EnvSchema = z.object({ PORT: z.coerce.number().default(23000), REDIS_URL: z.string().optional(), REDIS_TLS_REJECT_UNAUTHORIZED: z.string().default("true").transform(booleanTransform), + REDIS_COMMAND_TIMEOUT_MS: optionalNumber( + z + .number() + .int() + .min(100, "REDIS_COMMAND_TIMEOUT_MS 不能小于 100") + .max(120000, "REDIS_COMMAND_TIMEOUT_MS 不能大于 120000") + ).default(10_000), ENABLE_RATE_LIMIT: z.string().default("true").transform(booleanTransform), ENABLE_SECURE_COOKIES: z.string().default("true").transform(booleanTransform), ENABLE_LEGACY_ACTIONS_API: z.string().default("true").transform(booleanTransform), diff --git a/src/lib/langfuse/emit-proxy-trace.ts b/src/lib/langfuse/emit-proxy-trace.ts index 1d143c117..7a0b0bc22 100644 --- a/src/lib/langfuse/emit-proxy-trace.ts +++ b/src/lib/langfuse/emit-proxy-trace.ts @@ -106,9 +106,19 @@ export function emitProxyLangfuseTrace( ): void { if (!process.env.LANGFUSE_PUBLIC_KEY || !process.env.LANGFUSE_SECRET_KEY) return; - // 必须在异步 import 之前截断,避免动态加载/SDK 发送期间闭包继续强引用完整大响应。 - const responseText = truncateResponseTextForLangfuse(data.responseText); - const sessionSnapshot = buildLangfuseSessionSnapshot(session); + let responseText: string; + let sessionSnapshot: ProxySession; + try { + // 必须在异步 import 之前截断,避免动态加载/SDK 发送期间闭包继续强引用完整大响应。 + responseText = truncateResponseTextForLangfuse(data.responseText); + sessionSnapshot = buildLangfuseSessionSnapshot(session); + } catch (err) { + logger.warn("[Langfuse] Proxy trace snapshot failed", { + error: err instanceof Error ? err.message : String(err), + }); + return; + } + const { responseHeaders, durationMs, diff --git a/src/lib/langfuse/trace-proxy-request.ts b/src/lib/langfuse/trace-proxy-request.ts index abbe359ea..cf5d3e2d4 100644 --- a/src/lib/langfuse/trace-proxy-request.ts +++ b/src/lib/langfuse/trace-proxy-request.ts @@ -1,5 +1,7 @@ +import { sanitizeHeaders } from "@/app/v1/_lib/proxy/errors"; import type { UsageMetrics } from "@/app/v1/_lib/proxy/response-handler"; import type { ProxySession } from "@/app/v1/_lib/proxy/session"; +import { redactHeaders } from "@/lib/api/v1/_shared/redaction"; import { isLangfuseEnabled } from "@/lib/langfuse/index"; import { logger } from "@/lib/logger"; import type { CostBreakdown } from "@/lib/utils/cost-calculation"; @@ -38,23 +40,32 @@ function getStatusCategory(statusCode: number): string { return `${Math.floor(statusCode / 100)}xx`; } -/** - * Convert Headers to a plain record. - * - * Security note: session.headers are the CLIENT's original request headers - * (user -> CCH), which may include the user's own CCH auth key. These are - * safe to log -- the user already knows their own credentials. - * - * The upstream PROVIDER API key (outboundKey) is injected by ProxyForwarder - * into a separate Headers object and is NEVER present in session.headers or - * ctx.responseHeaders, so no redaction is needed here. - */ -function headersToRecord(headers: Headers): Record { - const result: Record = {}; +function sanitizeProviderChainValue(value: unknown, key?: string): unknown { + if (key === "headers" && typeof value === "string") { + return sanitizeHeaders(value); + } + if (Array.isArray(value)) { + return value.map((item) => sanitizeProviderChainValue(item)); + } + if (value && typeof value === "object") { + return Object.fromEntries( + Object.entries(value).map(([entryKey, entryValue]) => [ + entryKey, + sanitizeProviderChainValue(entryValue, entryKey), + ]) + ); + } + return value; +} + +function redactLangfuseHeaders(headers: Headers): Record { + const externalHeaders = new Headers(); headers.forEach((value, key) => { - result[key] = value; + if (!key.toLowerCase().startsWith("x-cch-")) { + externalHeaders.append(key, value); + } }); - return result; + return redactHeaders(externalHeaders); } const SUCCESS_REASONS = new Set([ @@ -242,13 +253,15 @@ export async function traceProxyRequest(ctx: TraceContext): Promise { requestSequence: String(session.getRequestSequence()), }; - // Build generation metadata - all request detail fields, raw headers (no redaction) + const requestHeaders = redactLangfuseHeaders(session.headers); + const responseHeaders = redactLangfuseHeaders(ctx.responseHeaders); + const generationMetadata: Record = { // Provider providerId: provider?.id, providerName: provider?.name, providerType: provider?.providerType, - providerChain: session.getProviderChain(), + providerChain: sanitizeProviderChainValue(session.getProviderChain()), // Model model: session.getCurrentModel(), originalModel: session.getOriginalModel(), @@ -277,9 +290,8 @@ export async function traceProxyRequest(ctx: TraceContext): Promise { requestSummary: buildRequestBodySummary(session), // SSE sseEventCount: ctx.sseEventCount, - // Headers (raw, no redaction) - requestHeaders: headersToRecord(session.headers), - responseHeaders: headersToRecord(ctx.responseHeaders), + requestHeaders, + responseHeaders, }; // Build usage details for Langfuse generation @@ -344,7 +356,8 @@ export async function traceProxyRequest(ctx: TraceContext): Promise { } // 2. Provider attempt events (one per failed/hedge chain item) - for (const item of session.getProviderChain()) { + for (const rawItem of session.getProviderChain()) { + const item = sanitizeProviderChainValue(rawItem) as typeof rawItem; // Hedge trigger: informational event (not a success or failure) if (item.reason === "hedge_triggered") { const hedgeObs = rootSpan.startObservation( diff --git a/src/lib/ledger-backfill/service.ts b/src/lib/ledger-backfill/service.ts index b71fba5ba..c708a499a 100644 --- a/src/lib/ledger-backfill/service.ts +++ b/src/lib/ledger-backfill/service.ts @@ -11,13 +11,17 @@ export interface BackfillUsageLedgerSummary { alreadyExisted: number; } -export async function backfillUsageLedger(): Promise { +export async function backfillUsageLedger( + signal?: AbortSignal +): Promise { const startTime = Date.now(); const LOCK_KEY = 20260101; + signal?.throwIfAborted(); // Use pg_try_advisory_xact_lock (transaction-scoped) so lock/unlock always happen // on the same connection — safe with connection pools. return await db.transaction(async (tx) => { + signal?.throwIfAborted(); const lockResult = await tx.execute(sql` SELECT pg_try_advisory_xact_lock(${LOCK_KEY}) AS acquired `); @@ -39,6 +43,7 @@ export async function backfillUsageLedger(): Promise let lastId = 0; while (true) { + signal?.throwIfAborted(); const batchResult = await tx.execute(sql` WITH batch AS ( SELECT @@ -189,6 +194,7 @@ export async function backfillUsageLedger(): Promise const inserted = Number(batchRow?.inserted ?? 0); const updated = Number(batchRow?.updated ?? 0); const maxId = Number(batchRow?.max_id ?? 0); + signal?.throwIfAborted(); if (processed === 0) { break; diff --git a/src/lib/lifecycle/shutdown.ts b/src/lib/lifecycle/shutdown.ts index 00521d875..42aa4f800 100644 --- a/src/lib/lifecycle/shutdown.ts +++ b/src/lib/lifecycle/shutdown.ts @@ -46,15 +46,42 @@ async function withTimeout(p: Promise, ms: number, label: string): Promise } } +async function awaitWithWarning(p: Promise, ms: number, label: string): Promise { + const timer = setTimeout(() => { + logger.warn(`[Shutdown] ${label} still pending`, { ms }); + }, ms); + try { + return await p; + } finally { + clearTimeout(timer); + } +} + +async function awaitQuiescenceBestEffort( + promise: Promise, + warningMs: number, + label: string +): Promise { + try { + await awaitWithWarning(promise, warningMs, label); + } catch (error) { + logger.warn(`[Shutdown] ${label} failed`, { + error: error instanceof Error ? error.message : String(error), + }); + } +} + const DEFAULT_STEP_TIMEOUT_MS = 3000; const DEFAULT_TOTAL_TIMEOUT_MS = 10000; export interface RunCleanupOptions { + // Cleanup 的慢操作告警阈值;最终强制退出由 server.js hard watchdog 负责。 totalTimeoutMs?: number; perStepTimeoutMs?: number; } -// 串行执行每一步的资源回收。每步超时不阻塞后续步骤;整体超时是兜底保护。 +// 串行执行资源回收。非关键步骤超时后继续;async task、writer 与 DB pool 是不可 detach 的 +// critical barrier,超时只告警,失败则向 server.js 传播并触发非零退出。 export async function runApplicationCleanup( signal: string, opts: RunCleanupOptions = {} @@ -64,10 +91,12 @@ export async function runApplicationCleanup( const startedAt = Date.now(); logger.info("[Shutdown] application cleanup starting", { signal, totalMs, stepMs }); + let writerQuiescencePending = false; + const deferredErrors: Error[] = []; const work = (async () => { // 1. 停止本地周期任务(不需要做 IO,几乎是同步) - await withTimeout( + await awaitQuiescenceBestEffort( (async () => { const { stopCacheCleanup } = await import("@/lib/cache/session-cache"); stopCacheCleanup(); @@ -77,19 +106,19 @@ export async function runApplicationCleanup( ); // 2. 端点探测调度器 - await withTimeout( + await awaitQuiescenceBestEffort( (async () => { const { stopEndpointProbeScheduler } = await import( "@/lib/provider-endpoints/probe-scheduler" ); - stopEndpointProbeScheduler(); + await stopEndpointProbeScheduler(); })(), stepMs, "stopEndpointProbeScheduler" ); // 3. 公共状态重建调度器 - await withTimeout( + await awaitQuiescenceBestEffort( (async () => { const { stopPublicStatusRebuildScheduler } = await import("@/lib/public-status/scheduler"); await stopPublicStatusRebuildScheduler(); @@ -99,40 +128,95 @@ export async function runApplicationCleanup( ); // 4. 端点探测日志清理 - await withTimeout( + await awaitQuiescenceBestEffort( (async () => { const { stopEndpointProbeLogCleanup } = await import( "@/lib/provider-endpoints/probe-log-cleanup" ); - stopEndpointProbeLogCleanup(); + await stopEndpointProbeLogCleanup(); })(), stepMs, "stopEndpointProbeLogCleanup" ); - // 5. 取消仍在飞的后台异步任务。 + // 5. Bull queues own Redis connections and may still ACK jobs or emit DB work. + // Join them before closing either backing resource. + try { + await awaitWithWarning( + (async () => { + const stopQueues = ( + globalThis as typeof globalThis & { + __CCH_STOP_BACKGROUND_QUEUES__?: () => Promise; + } + ).__CCH_STOP_BACKGROUND_QUEUES__; + if (stopQueues) await stopQueues(); + })(), + stepMs, + "stopBackgroundQueues" + ); + } catch (error) { + const queueError = error instanceof Error ? error : new Error(String(error)); + deferredErrors.push(queueError); + logger.error("[Shutdown] background queues failed to stop; continuing critical cleanup", { + error: queueError.message, + }); + } + + // 6. 取消仍在飞的后台异步任务。 // 必须排在 message-buffer flush 之前——任务被 abort 时仍会写出尾部日志/用量记录, // flush 才能把这些尾部更新真正落库。 - await withTimeout( - (async () => { - const { shutdownAllAsyncTasks } = await import("@/lib/async-task-manager"); - shutdownAllAsyncTasks(); - })(), - stepMs, - "shutdownAllAsyncTasks" - ); + const asyncTasksWarningTimer = setTimeout(() => { + logger.warn("[Shutdown] shutdownAllAsyncTasks still pending", { ms: stepMs }); + }, stepMs); + try { + const { shutdownAllAsyncTasks } = await import("@/lib/async-task-manager"); + await shutdownAllAsyncTasks(); + } catch (error) { + logger.error("[Shutdown] async tasks failed to settle", { + error: error instanceof Error ? error.message : String(error), + }); + throw error; + } finally { + clearTimeout(asyncTasksWarningTimer); + } - // 6. 刷写 message_request 异步写缓冲 - await withTimeout( - (async () => { - const { stopMessageRequestWriteBuffer } = await import("@/repository/message-write-buffer"); - await stopMessageRequestWriteBuffer(); - })(), - stepMs, - "stopMessageRequestWriteBuffer" - ); + // 7. 刷写 message_request 异步写缓冲。这里不能用可脱离的单步 timeout: + // closeDbPools 必须等 writer 真正 settled,否则会关闭仍在执行终态 SQL 的连接。 + writerQuiescencePending = true; + const writerWarningTimer = setTimeout(() => { + logger.warn("[Shutdown] stopMessageRequestWriteBuffer still pending", { ms: stepMs }); + }, stepMs); + try { + const { stopMessageRequestWriteBuffer } = await import("@/repository/message-write-buffer"); + await stopMessageRequestWriteBuffer(); + } catch (error) { + logger.error("[Shutdown] message writer failed to quiesce; database pools remain open", { + error: error instanceof Error ? error.message : String(error), + }); + throw error; + } finally { + clearTimeout(writerWarningTimer); + writerQuiescencePending = false; + } - // 7. Langfuse 自带超时(LANGFUSE_SHUTDOWN_TIMEOUT_MS),这里再加一层兜底 + // 8. writer flush 完成后再关闭数据库 pool。pool close 也是 critical barrier, + // 单步 deadline 只能告警,不能让底层 client.end() 脱离 shutdown 生命周期。 + const dbWarningTimer = setTimeout(() => { + logger.warn("[Shutdown] closeDbPools still pending", { ms: stepMs }); + }, stepMs); + try { + const { closeDbPools } = await import("@/drizzle/db"); + await closeDbPools(); + } catch (error) { + logger.error("[Shutdown] database pools failed to close", { + error: error instanceof Error ? error.message : String(error), + }); + throw error; + } finally { + clearTimeout(dbWarningTimer); + } + + // 9. Langfuse 自带超时(LANGFUSE_SHUTDOWN_TIMEOUT_MS),这里再加一层兜底 await withTimeout( (async () => { const { shutdownLangfuse } = await import("@/lib/langfuse"); @@ -142,7 +226,7 @@ export async function runApplicationCleanup( "shutdownLangfuse" ); - // 8. Redis 连接最后关:上面的步骤可能仍在写日志/缓存 + // 10. Redis 连接最后关:上面的步骤可能仍在写日志/缓存 await withTimeout( (async () => { const { closeRedis } = await import("@/lib/redis"); @@ -152,7 +236,7 @@ export async function runApplicationCleanup( "closeRedis" ); - // 9. API Key Vacuum Filter 订阅清理 —— 同步函数,不需要 timeout + // 11. API Key Vacuum Filter 订阅清理 —— 同步函数,不需要 timeout try { const g = globalThis as unknown as { __CCH_API_KEY_VF_SYNC_CLEANUP__?: (() => void) | null; @@ -164,7 +248,7 @@ export async function runApplicationCleanup( }); } - // 10. 云价格定时同步 + // 12. 云价格定时同步 try { const g = globalThis as unknown as { __CCH_CLOUD_PRICE_SYNC_INTERVAL_ID__?: ReturnType; @@ -178,17 +262,30 @@ export async function runApplicationCleanup( error: error instanceof Error ? error.message : String(error), }); } + + if (deferredErrors.length > 0) { + throw new AggregateError(deferredErrors, "Application cleanup completed with errors"); + } })(); - const total = new Promise((resolve) => { - const t = setTimeout(() => { - logger.warn("[Shutdown] application cleanup total timeout reached", { totalMs }); - resolve(); - }, totalMs); - work.finally(() => clearTimeout(t)); - }); + const totalWarningTimer = setTimeout(() => { + logger.warn( + "[Shutdown] application cleanup total timeout reached; continuing critical cleanup", + { totalMs } + ); + if (writerQuiescencePending) { + logger.error( + "[Shutdown] cleanup deadline reached with message writer still active; continuing to wait", + { totalMs } + ); + } + }, totalMs); - await Promise.race([work, total]); + try { + await work; + } finally { + clearTimeout(totalWarningTimer); + } logger.info("[Shutdown] application cleanup complete", { signal, diff --git a/src/lib/price-sync/cloud-price-updater.ts b/src/lib/price-sync/cloud-price-updater.ts index 98d5c5733..0258fae8f 100644 --- a/src/lib/price-sync/cloud-price-updater.ts +++ b/src/lib/price-sync/cloud-price-updater.ts @@ -202,7 +202,7 @@ export function requestCloudPriceTableSync(options: { AsyncTaskManager.register( taskId, - (async () => { + async () => { try { const result = await syncCloudPriceTableToDatabase(); if (!result.ok) { @@ -223,7 +223,7 @@ export function requestCloudPriceTableSync(options: { } finally { g.__CCH_CLOUD_PRICE_SYNC_LAST_AT__ = Date.now(); } - })(), + }, "cloud_price_table_sync" ); } catch (error) { diff --git a/src/lib/provider-endpoints/probe-log-cleanup.ts b/src/lib/provider-endpoints/probe-log-cleanup.ts index 9e6bd0108..11f93668f 100644 --- a/src/lib/provider-endpoints/probe-log-cleanup.ts +++ b/src/lib/provider-endpoints/probe-log-cleanup.ts @@ -30,12 +30,15 @@ const cleanupState = globalThis as unknown as { __CCH_ENDPOINT_PROBE_LOG_CLEANUP_INTERVAL_ID__?: ReturnType; __CCH_ENDPOINT_PROBE_LOG_CLEANUP_LOCK__?: LeaderLock; __CCH_ENDPOINT_PROBE_LOG_CLEANUP_RUNNING__?: boolean; + __CCH_ENDPOINT_PROBE_LOG_CLEANUP_CURRENT_PROMISE__?: Promise; + __CCH_ENDPOINT_PROBE_LOG_CLEANUP_STOP_REQUESTED__?: boolean; }; async function runCleanupOnce(): Promise { if (cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_RUNNING__) { return; } + if (cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_STOP_REQUESTED__) return; cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_RUNNING__ = true; @@ -69,7 +72,7 @@ async function runCleanupOnce(): Promise { let totalDeleted = 0; while (true) { - if (leadershipLost) { + if (leadershipLost || cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_STOP_REQUESTED__) { return; } @@ -78,6 +81,10 @@ async function runCleanupOnce(): Promise { batchSize: CLEANUP_BATCH_SIZE, }); + if (cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_STOP_REQUESTED__) { + return; + } + if (deleted <= 0) { break; } @@ -110,6 +117,16 @@ async function runCleanupOnce(): Promise { } } +function launchCleanup(): void { + if (cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_CURRENT_PROMISE__) return; + const current = runCleanupOnce().finally(() => { + if (cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_CURRENT_PROMISE__ === current) { + cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_CURRENT_PROMISE__ = undefined; + } + }); + cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_CURRENT_PROMISE__ = current; +} + export function startEndpointProbeLogCleanup(): void { if (process.env.CI === "true") { return; @@ -120,15 +137,17 @@ export function startEndpointProbeLogCleanup(): void { } cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_STARTED__ = true; + cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_STOP_REQUESTED__ = false; - void runCleanupOnce(); + launchCleanup(); cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_INTERVAL_ID__ = setInterval(() => { - void runCleanupOnce(); + launchCleanup(); }, CLEANUP_INTERVAL_MS); } -export function stopEndpointProbeLogCleanup(): void { +export async function stopEndpointProbeLogCleanup(): Promise { + cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_STOP_REQUESTED__ = true; const intervalId = cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_INTERVAL_ID__; if (intervalId) { clearInterval(intervalId); @@ -136,11 +155,12 @@ export function stopEndpointProbeLogCleanup(): void { cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_INTERVAL_ID__ = undefined; cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_STARTED__ = false; - cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_RUNNING__ = false; + + await cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_CURRENT_PROMISE__; const lock = cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_LOCK__; cleanupState.__CCH_ENDPOINT_PROBE_LOG_CLEANUP_LOCK__ = undefined; if (lock) { - void releaseLeaderLock(lock); + await releaseLeaderLock(lock); } } diff --git a/src/lib/provider-endpoints/probe-scheduler.ts b/src/lib/provider-endpoints/probe-scheduler.ts index d5c8d9d3b..1b48572b1 100644 --- a/src/lib/provider-endpoints/probe-scheduler.ts +++ b/src/lib/provider-endpoints/probe-scheduler.ts @@ -51,6 +51,7 @@ const schedulerState = globalThis as unknown as { __CCH_ENDPOINT_PROBE_SCHEDULER_STOP_REQUESTED__?: boolean; __CCH_ENDPOINT_PROBE_SCHEDULER_NEXT_DUE_AT_MS__?: number; __CCH_ENDPOINT_PROBE_SCHEDULER_NEXT_DB_POLL_AT_MS__?: number; + __CCH_ENDPOINT_PROBE_SCHEDULER_CURRENT_PROMISE__?: Promise; }; function sleep(ms: number): Promise { @@ -324,6 +325,16 @@ async function runProbeCycle(): Promise { } } +function launchProbeCycle(): void { + if (schedulerState.__CCH_ENDPOINT_PROBE_SCHEDULER_CURRENT_PROMISE__) return; + const current = runProbeCycle().finally(() => { + if (schedulerState.__CCH_ENDPOINT_PROBE_SCHEDULER_CURRENT_PROMISE__ === current) { + schedulerState.__CCH_ENDPOINT_PROBE_SCHEDULER_CURRENT_PROMISE__ = undefined; + } + }); + schedulerState.__CCH_ENDPOINT_PROBE_SCHEDULER_CURRENT_PROMISE__ = current; +} + export function startEndpointProbeScheduler(): void { if (schedulerState.__CCH_ENDPOINT_PROBE_SCHEDULER_STARTED__) { return; @@ -333,10 +344,10 @@ export function startEndpointProbeScheduler(): void { schedulerState.__CCH_ENDPOINT_PROBE_SCHEDULER_STARTED__ = true; clearNextWorkHints(); - void runProbeCycle(); + launchProbeCycle(); schedulerState.__CCH_ENDPOINT_PROBE_SCHEDULER_INTERVAL_ID__ = setInterval(() => { - void runProbeCycle(); + launchProbeCycle(); }, TICK_INTERVAL_MS); logger.info("[EndpointProbeScheduler] Started", { @@ -352,7 +363,7 @@ export function startEndpointProbeScheduler(): void { }); } -export function stopEndpointProbeScheduler(): void { +export async function stopEndpointProbeScheduler(): Promise { schedulerState.__CCH_ENDPOINT_PROBE_SCHEDULER_STOP_REQUESTED__ = true; clearNextWorkHints(); @@ -364,10 +375,12 @@ export function stopEndpointProbeScheduler(): void { schedulerState.__CCH_ENDPOINT_PROBE_SCHEDULER_INTERVAL_ID__ = undefined; schedulerState.__CCH_ENDPOINT_PROBE_SCHEDULER_STARTED__ = false; + await schedulerState.__CCH_ENDPOINT_PROBE_SCHEDULER_CURRENT_PROMISE__; + const lock = schedulerState.__CCH_ENDPOINT_PROBE_SCHEDULER_LOCK__; schedulerState.__CCH_ENDPOINT_PROBE_SCHEDULER_LOCK__ = undefined; if (lock) { - void releaseLeaderLock(lock); + await releaseLeaderLock(lock); } } diff --git a/src/lib/provider-testing/test-service.test.ts b/src/lib/provider-testing/test-service.test.ts index 49fa29ef6..18d093435 100644 --- a/src/lib/provider-testing/test-service.test.ts +++ b/src/lib/provider-testing/test-service.test.ts @@ -169,33 +169,33 @@ describe("executeProviderTest", () => { expectRequestUrl("https://relay.example.com/openai/v1/responses"); }); - test.each([ - "https://api.gptclubapi.xyz/openai", - "https://api.gptclubapi.xyz/openai/", - ])("codex bare /openai base preserves absolute versioned request url: %s", async (providerUrl) => { - mockJsonResponse({ - id: "resp_test", - model: "gpt-5.5", - output: [ - { - type: "message", - role: "assistant", - content: [{ type: "output_text", text: "pong" }], - }, - ], - }); + test.each(["https://api.gptclubapi.xyz/openai", "https://api.gptclubapi.xyz/openai/"])( + "codex bare /openai base preserves absolute versioned request url: %s", + async (providerUrl) => { + mockJsonResponse({ + id: "resp_test", + model: "gpt-5.5", + output: [ + { + type: "message", + role: "assistant", + content: [{ type: "output_text", text: "pong" }], + }, + ], + }); - const result = await executeProviderTest({ - providerUrl, - apiKey: "sk-test-codex", - providerType: "codex", - model: "gpt-5.5", - }); + const result = await executeProviderTest({ + providerUrl, + apiKey: "sk-test-codex", + providerType: "codex", + model: "gpt-5.5", + }); - expect(result.success).toBe(true); - expect(result.requestUrl).toBe("https://api.gptclubapi.xyz/openai/v1/responses"); - expect(fetchMock.mock.calls[0]?.[0]).toBe("https://api.gptclubapi.xyz/openai/v1/responses"); - }); + expect(result.success).toBe(true); + expect(result.requestUrl).toBe("https://api.gptclubapi.xyz/openai/v1/responses"); + expect(fetchMock.mock.calls[0]?.[0]).toBe("https://api.gptclubapi.xyz/openai/v1/responses"); + } + ); test("openai-compatible 版本根路径应只追加 endpoint,不重复拼接 /v1", async () => { mockJsonResponse({ diff --git a/src/lib/public-status/scheduler.ts b/src/lib/public-status/scheduler.ts index 1edd2e378..6709dc268 100644 --- a/src/lib/public-status/scheduler.ts +++ b/src/lib/public-status/scheduler.ts @@ -23,6 +23,7 @@ const schedulerState = globalThis as unknown as { __CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_RUNNING__?: boolean; __CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_LOCK__?: LeaderLock; __CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_STOP_REQUESTED__?: boolean; + __CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_CURRENT_PROMISE__?: Promise; }; function parseRebuildHintKey( @@ -168,7 +169,7 @@ async function runCycle(): Promise { const targets = await collectTargets(); for (const target of targets) { - if (leadershipLost) { + if (leadershipLost || schedulerState.__CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_STOP_REQUESTED__) { break; } const result = await rebuildPublicStatusProjection({ @@ -176,6 +177,9 @@ async function runCycle(): Promise { rangeHours: target.rangeHours, redis, }); + if (schedulerState.__CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_STOP_REQUESTED__) { + break; + } if (result.status === "updated" && target.hintKey) { await redis.del(target.hintKey); } @@ -190,6 +194,16 @@ async function runCycle(): Promise { } } +function launchCycle(): void { + if (schedulerState.__CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_CURRENT_PROMISE__) return; + const current = runCycle().finally(() => { + if (schedulerState.__CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_CURRENT_PROMISE__ === current) { + schedulerState.__CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_CURRENT_PROMISE__ = undefined; + } + }); + schedulerState.__CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_CURRENT_PROMISE__ = current; +} + export function startPublicStatusRebuildScheduler(): void { if (schedulerState.__CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_STARTED__) { return; @@ -198,7 +212,7 @@ export function startPublicStatusRebuildScheduler(): void { schedulerState.__CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_STOP_REQUESTED__ = false; schedulerState.__CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_STARTED__ = true; schedulerState.__CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_INTERVAL_ID__ = setInterval(() => { - void runCycle(); + launchCycle(); }, TICK_INTERVAL_MS); const timer = schedulerState.__CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_INTERVAL_ID__ as @@ -206,7 +220,7 @@ export function startPublicStatusRebuildScheduler(): void { | undefined; timer?.unref?.(); - void runCycle(); + launchCycle(); } export async function stopPublicStatusRebuildScheduler(): Promise { @@ -217,6 +231,8 @@ export async function stopPublicStatusRebuildScheduler(): Promise { clearInterval(intervalId); } + await schedulerState.__CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_CURRENT_PROMISE__; + const currentLock = schedulerState.__CCH_PUBLIC_STATUS_REBUILD_SCHEDULER_LOCK__; if (currentLock) { await releaseLeaderLock(currentLock); diff --git a/src/lib/rate-limit/lease-service.ts b/src/lib/rate-limit/lease-service.ts index 6406afca0..99ba9a29f 100644 --- a/src/lib/rate-limit/lease-service.ts +++ b/src/lib/rate-limit/lease-service.ts @@ -66,10 +66,55 @@ export interface DecrementLeaseBudgetResult { failOpen?: boolean; } +export interface LeaseSettlementEntity { + id: number; + resetModes?: Partial>; +} + +export interface SettleLeaseBudgetsParams { + requestId: string | number; + cost: number; + entities: { + key: LeaseSettlementEntity; + user: LeaseSettlementEntity; + provider: LeaseSettlementEntity; + }; +} + +export type LeaseBudgetSettlementStatus = "decremented" | "missing" | "insufficient"; + +export interface LeaseBudgetSettlement { + entityType: LeaseEntityTypeType; + entityId: number; + window: LeaseWindowType; + status: LeaseBudgetSettlementStatus; + newRemaining: number; +} + +export interface SettleLeaseBudgetsResult { + requestId: string; + status: "settled" | "duplicate" | "fail_open"; + settlements: LeaseBudgetSettlement[]; + failOpen?: boolean; +} + +interface LeaseSettlementTarget { + entityType: LeaseEntityTypeType; + entityId: number; + window: LeaseWindowType; + resetMode?: DailyResetMode; +} + /** * Lease Service - manages budget leases for rate limiting */ export class LeaseService { + private static readonly SETTLEMENT_MARKER_TTL_SECONDS = 5 * 60; + + private static readonly SETTLEMENT_ENTITY_TYPES = ["key", "user", "provider"] as const; + + private static readonly SETTLEMENT_WINDOWS = ["5h", "daily", "weekly", "monthly"] as const; + private static get redis() { return getRedisClient(); } @@ -387,6 +432,225 @@ export class LeaseService { return {newRemaining, 1} `; + /** + * Atomically settle the fixed 4 windows x 3 entity lease set. + * + * KEYS[1] is a bounded idempotency marker. KEYS[2..13] are the lease keys in + * key/user/provider then 5h/daily/weekly/monthly order. ARGV[1] is the actual + * cost and ARGV[2] is the marker TTL in seconds. + * + * The marker survives the Redis client's bounded reconnect retry cycle while + * expiring after five minutes so marker cardinality remains bounded. + */ + private static readonly SETTLE_LEASE_BUDGETS_LUA_SCRIPT = ` + local markerKey = KEYS[1] + local previousSettlement = redis.call("GET", markerKey) + if previousSettlement then + return {1, previousSettlement} + end + + local cost = tonumber(ARGV[1]) + local markerTtlSeconds = tonumber(ARGV[2]) + local settlements = {} + local pendingWrites = {} + + for keyIndex = 2, #KEYS do + local leaseKey = KEYS[keyIndex] + local leaseReply = redis.pcall("GET", leaseKey) + local leaseReadFailed = type(leaseReply) == "table" and leaseReply.err + + if leaseReadFailed or not leaseReply then + settlements[#settlements + 1] = {0, -1} + else + local decoded, lease = pcall(cjson.decode, leaseReply) + local remaining = nil + if decoded and type(lease) == "table" then + remaining = tonumber(lease.remainingBudget) + end + local ttl = redis.call("TTL", leaseKey) + + if not remaining or ttl <= 0 then + settlements[#settlements + 1] = {0, -1} + elseif remaining < cost then + -- Consume the cached slice when the request is larger than the + -- remaining lease. Keeping a positive balance here lets every + -- request in the refresh window repeat the same overshoot. + lease.remainingBudget = 0 + local encodedLeaseOk, encodedLease = pcall(cjson.encode, lease) + if not encodedLeaseOk then + settlements[#settlements + 1] = {0, -1} + else + pendingWrites[#pendingWrites + 1] = {leaseKey, ttl, encodedLease} + settlements[#settlements + 1] = {-1, 0} + end + else + local newRemaining = remaining - cost + lease.remainingBudget = newRemaining + local encodedLeaseOk, encodedLease = pcall(cjson.encode, lease) + if not encodedLeaseOk then + settlements[#settlements + 1] = {0, -1} + else + pendingWrites[#pendingWrites + 1] = {leaseKey, ttl, encodedLease} + settlements[#settlements + 1] = {1, newRemaining} + end + end + end + end + + local encodedOk, encoded = pcall(cjson.encode, settlements) + if not encodedOk then + return redis.error_reply("failed to encode lease settlement results") + end + + for writeIndex = 1, #pendingWrites do + local pendingWrite = pendingWrites[writeIndex] + redis.call("SETEX", pendingWrite[1], pendingWrite[2], pendingWrite[3]) + end + + redis.call("SETEX", markerKey, markerTtlSeconds, encoded) + return {0, encoded} + `; + + private static buildSettlementTargets(params: SettleLeaseBudgetsParams): LeaseSettlementTarget[] { + const targets: LeaseSettlementTarget[] = []; + + for (const entityType of LeaseService.SETTLEMENT_ENTITY_TYPES) { + const entity = params.entities[entityType]; + + for (const window of LeaseService.SETTLEMENT_WINDOWS) { + targets.push({ + entityType, + entityId: entity.id, + window, + resetMode: + window === "5h" || window === "daily" ? entity.resetModes?.[window] : undefined, + }); + } + } + + return targets; + } + + private static parseSettlementResults( + rawSettlements: unknown, + targets: LeaseSettlementTarget[] + ): LeaseBudgetSettlement[] { + if (typeof rawSettlements !== "string") { + throw new Error("Invalid lease settlement payload"); + } + + const parsed = JSON.parse(rawSettlements) as unknown; + if (!Array.isArray(parsed) || parsed.length !== targets.length) { + throw new Error("Invalid lease settlement result count"); + } + + return targets.map((target, index) => { + const rawResult = parsed[index]; + if (!Array.isArray(rawResult) || rawResult.length !== 2) { + throw new Error("Invalid lease settlement item"); + } + + const statusCode = Number(rawResult[0]); + const newRemaining = Number(rawResult[1]); + if (!Number.isFinite(newRemaining)) { + throw new Error("Invalid lease settlement remaining budget"); + } + + let status: LeaseBudgetSettlementStatus; + if (statusCode === 1) { + status = "decremented"; + } else if (statusCode === 0) { + status = "missing"; + } else if (statusCode === -1) { + status = "insufficient"; + } else { + throw new Error("Invalid lease settlement status"); + } + + return { + entityType: target.entityType, + entityId: target.entityId, + window: target.window, + status, + newRemaining, + }; + }); + } + + /** + * Settle one request's actual cost against all twelve lease budgets. + * + * A request marker and all lease mutations run in one bounded Lua invocation. + * If ioredis resends a command after losing the first reply, the marker returns + * the original result without applying the cost again. + */ + static async settleLeaseBudgets( + params: SettleLeaseBudgetsParams + ): Promise { + const requestId = String(params.requestId).trim(); + + try { + const redis = LeaseService.redis; + if (redis?.status !== "ready") { + logger.warn("[LeaseService] Redis not ready, fail-open for batch settlement", { + requestId, + cost: params.cost, + }); + return { requestId, status: "fail_open", settlements: [], failOpen: true }; + } + + if (!requestId || !Number.isFinite(params.cost) || params.cost <= 0) { + logger.warn("[LeaseService] Invalid batch settlement input, fail-open", { + requestId, + cost: params.cost, + }); + return { requestId, status: "fail_open", settlements: [], failOpen: true }; + } + + const targets = LeaseService.buildSettlementTargets(params); + const markerKey = `lease:settlement:${requestId}`; + const leaseKeys = targets.map((target) => + buildLeaseKey(target.entityType, target.entityId, target.window, target.resetMode) + ); + + const rawResult = (await redis.eval( + LeaseService.SETTLE_LEASE_BUDGETS_LUA_SCRIPT, + 1 + leaseKeys.length, + markerKey, + ...leaseKeys, + params.cost.toString(), + LeaseService.SETTLEMENT_MARKER_TTL_SECONDS.toString() + )) as unknown; + + if (!Array.isArray(rawResult) || rawResult.length !== 2) { + throw new Error("Invalid lease settlement response"); + } + + const duplicateFlag = Number(rawResult[0]); + if (duplicateFlag !== 0 && duplicateFlag !== 1) { + throw new Error("Invalid lease settlement duplicate flag"); + } + + const settlements = LeaseService.parseSettlementResults(rawResult[1], targets); + const status = duplicateFlag === 1 ? "duplicate" : "settled"; + + logger.debug("[LeaseService] Batch lease settlement completed", { + requestId, + status, + cost: params.cost, + }); + + return { requestId, status, settlements }; + } catch (error) { + logger.error("[LeaseService] settleLeaseBudgets failed, fail-open", { + requestId, + cost: params.cost, + error, + }); + return { requestId, status: "fail_open", settlements: [], failOpen: true }; + } + } + /** * Decrement lease budget atomically using Lua script * diff --git a/src/lib/rate-limit/service.ts b/src/lib/rate-limit/service.ts index 1e7c7f304..e1927c308 100644 --- a/src/lib/rate-limit/service.ts +++ b/src/lib/rate-limit/service.ts @@ -65,6 +65,7 @@ * ============================================================================ */ +import type { ChainableCommander } from "ioredis"; import { logger } from "@/lib/logger"; import { getRedisClient } from "@/lib/redis"; import { @@ -78,8 +79,7 @@ import { GET_COST_5H_ROLLING_WINDOW, GET_COST_DAILY_ROLLING_WINDOW, RELEASE_PROVIDER_SESSION, - TRACK_COST_5H_ROLLING_WINDOW, - TRACK_COST_DAILY_ROLLING_WINDOW, + TRACK_COST_ROLLING_WINDOW, } from "@/lib/redis/lua-scripts"; import { SessionTracker } from "@/lib/session-tracker"; import { ERROR_CODES } from "@/lib/utils/error-messages"; @@ -91,7 +91,12 @@ import { } from "@/repository/statistics"; import { clipStartByResetAt, resolveUser5hCostResetAt } from "./cost-reset-utils"; import type { LeaseWindowType } from "./lease"; -import { type DecrementLeaseBudgetResult, LeaseService } from "./lease-service"; +import { + type DecrementLeaseBudgetResult, + LeaseService, + type SettleLeaseBudgetsParams, + type SettleLeaseBudgetsResult, +} from "./lease-service"; import { type DailyResetMode, getResetAtFromTtlSeconds, @@ -187,15 +192,13 @@ export class RateLimitService { return state.resetAt; } - private static async trackFixedCostWindow( + private static queueFixedCostWindow( + pipeline: ChainableCommander, key: string, cost: number, ttlSeconds: number - ): Promise { - const redis = RateLimitService.redis; - if (!redis || redis.status !== "ready") return; - - await redis.eval( + ): void { + pipeline.eval( RateLimitService.TRACK_FIXED_COST_WINDOW_LUA, 1, key, @@ -204,6 +207,48 @@ export class RateLimitService { ); } + private static queueRollingCostWindow( + pipeline: ChainableCommander, + key: string, + cost: number, + now: number, + windowMs: number, + requestId: string, + ttlSeconds: number + ): void { + pipeline.eval( + TRACK_COST_ROLLING_WINDOW, + 1, + key, + cost.toString(), + now.toString(), + windowMs.toString(), + requestId, + ttlSeconds.toString() + ); + } + + private static logCostPipelineErrors( + results: Array<[Error | null, unknown]> | null, + operation: "trackCost" | "trackUserDailyCost" + ): void { + if (!results) { + logger.error("[RateLimit] Cost pipeline returned null", { operation }); + return; + } + + for (let commandIndex = 0; commandIndex < results.length; commandIndex += 1) { + const error = results[commandIndex]?.[0]; + if (!error) continue; + + logger.error("[RateLimit] Cost pipeline command failed", { + operation, + commandIndex, + error: error.message, + }); + } + } + private static async warmRollingCostZset( key: string, entries: Array<{ id: number; createdAt: Date; costUsd: number }>, @@ -927,56 +972,71 @@ export class RateLimitService { providerResetTime?: string; providerResetMode?: DailyResetMode; user5hResetMode?: DailyResetMode; - requestId?: number; + userResetTime?: string; + userResetMode?: DailyResetMode; + requestId?: string | number; createdAtMs?: number; } ): Promise { - if (!RateLimitService.redis || cost <= 0) return; + const redis = RateLimitService.redis; + if (!redis || redis.status !== "ready" || cost <= 0) return; try { const keyDailyReset = RateLimitService.resolveDailyReset(options?.keyResetTime); const providerDailyReset = RateLimitService.resolveDailyReset(options?.providerResetTime); + const userDailyReset = RateLimitService.resolveDailyReset(options?.userResetTime); const key5hMode = options?.key5hResetMode ?? "rolling"; const keyDailyMode = options?.keyResetMode ?? "fixed"; const provider5hMode = options?.provider5hResetMode ?? "rolling"; const providerDailyMode = options?.providerResetMode ?? "fixed"; const user5hMode = options?.user5hResetMode ?? "rolling"; + const userDailyMode = options?.userResetMode ?? "fixed"; const now = options?.createdAtMs ?? Date.now(); const requestId = options?.requestId != null ? String(options.requestId) : ""; const window5h = 5 * 60 * 60 * 1000; // 5 hours in ms const window24h = 24 * 60 * 60 * 1000; // 24 hours in ms - // 计算动态 TTL(daily/周/月) - const ttlDailyKey = await getTTLForPeriodWithMode( - "daily", - keyDailyReset.normalized, - keyDailyMode - ); - const ttlDailyProvider = - keyDailyReset.normalized === providerDailyReset.normalized && - keyDailyMode === providerDailyMode - ? ttlDailyKey - : await getTTLForPeriodWithMode( - "daily", - providerDailyReset.normalized, - providerDailyMode - ); - const ttlWeekly = await getTTLForPeriod("weekly"); - const ttlMonthly = await getTTLForPeriod("monthly"); + const dailyTtlPromises = new Map>(); + const getFixedDailyTtl = (normalizedResetTime: string): Promise => { + const cached = dailyTtlPromises.get(normalizedResetTime); + if (cached) return cached; + + const pending = getTTLForPeriodWithMode("daily", normalizedResetTime, "fixed"); + dailyTtlPromises.set(normalizedResetTime, pending); + return pending; + }; + + const [ttlDailyKey, ttlDailyProvider, ttlDailyUser, ttlWeekly, ttlMonthly] = + await Promise.all([ + keyDailyMode === "fixed" + ? getFixedDailyTtl(keyDailyReset.normalized) + : Promise.resolve(0), + providerDailyMode === "fixed" + ? getFixedDailyTtl(providerDailyReset.normalized) + : Promise.resolve(0), + options?.userId != null && userDailyMode === "fixed" + ? getFixedDailyTtl(userDailyReset.normalized) + : Promise.resolve(0), + getTTLForPeriod("weekly"), + getTTLForPeriod("monthly"), + ]); + + const pipeline = redis.pipeline(); // 1. 5h 窗口:rolling 使用 ZSET,fixed 仅在首个成功记账时创建 TTL 窗口 if (key5hMode === "rolling") { - await RateLimitService.redis.eval( - TRACK_COST_5H_ROLLING_WINDOW, - 1, // KEYS count - RateLimitService.get5hCostKey("key", keyId, "rolling"), // KEYS[1] - cost.toString(), // ARGV[1]: cost - now.toString(), // ARGV[2]: now - window5h.toString(), // ARGV[3]: window - requestId // ARGV[4]: request_id (optional) + RateLimitService.queueRollingCostWindow( + pipeline, + RateLimitService.get5hCostKey("key", keyId, "rolling"), + cost, + now, + window5h, + requestId, + 21600 ); } else { - await RateLimitService.trackFixedCostWindow( + RateLimitService.queueFixedCostWindow( + pipeline, RateLimitService.get5hCostKey("key", keyId, "fixed"), cost, 5 * 3600 @@ -984,17 +1044,18 @@ export class RateLimitService { } if (provider5hMode === "rolling") { - await RateLimitService.redis.eval( - TRACK_COST_5H_ROLLING_WINDOW, - 1, + RateLimitService.queueRollingCostWindow( + pipeline, RateLimitService.get5hCostKey("provider", providerId, "rolling"), - cost.toString(), - now.toString(), - window5h.toString(), - requestId + cost, + now, + window5h, + requestId, + 21600 ); } else { - await RateLimitService.trackFixedCostWindow( + RateLimitService.queueFixedCostWindow( + pipeline, RateLimitService.get5hCostKey("provider", providerId, "fixed"), cost, 5 * 3600 @@ -1003,17 +1064,18 @@ export class RateLimitService { if (options?.userId != null) { if (user5hMode === "rolling") { - await RateLimitService.redis.eval( - TRACK_COST_5H_ROLLING_WINDOW, - 1, + RateLimitService.queueRollingCostWindow( + pipeline, RateLimitService.get5hCostKey("user", options.userId, "rolling"), - cost.toString(), - now.toString(), - window5h.toString(), - requestId + cost, + now, + window5h, + requestId, + 21600 ); } else { - await RateLimitService.trackFixedCostWindow( + RateLimitService.queueFixedCostWindow( + pipeline, RateLimitService.get5hCostKey("user", options.userId, "fixed"), cost, 5 * 3600 @@ -1023,59 +1085,70 @@ export class RateLimitService { // 2. daily 滚动窗口:使用 Lua 脚本(ZSET) if (keyDailyMode === "rolling") { - await RateLimitService.redis.eval( - TRACK_COST_DAILY_ROLLING_WINDOW, - 1, + RateLimitService.queueRollingCostWindow( + pipeline, `key:${keyId}:cost_daily_rolling`, - cost.toString(), - now.toString(), - window24h.toString(), - requestId + cost, + now, + window24h, + requestId, + 90000 ); + } else { + const keyDailyKey = `key:${keyId}:cost_daily_${keyDailyReset.suffix}`; + pipeline.incrbyfloat(keyDailyKey, cost); + pipeline.expire(keyDailyKey, ttlDailyKey); } if (providerDailyMode === "rolling") { - await RateLimitService.redis.eval( - TRACK_COST_DAILY_ROLLING_WINDOW, - 1, + RateLimitService.queueRollingCostWindow( + pipeline, `provider:${providerId}:cost_daily_rolling`, - cost.toString(), - now.toString(), - window24h.toString(), - requestId + cost, + now, + window24h, + requestId, + 90000 ); + } else { + const providerDailyKey = `provider:${providerId}:cost_daily_${providerDailyReset.suffix}`; + pipeline.incrbyfloat(providerDailyKey, cost); + pipeline.expire(providerDailyKey, ttlDailyProvider); } - // 3. daily fixed/周/月固定窗口:使用 STRING + 动态 TTL - const pipeline = RateLimitService.redis.pipeline(); - - // Key 的 daily fixed/周/月消费 - if (keyDailyMode === "fixed") { - const keyDailyKey = `key:${keyId}:cost_daily_${keyDailyReset.suffix}`; - pipeline.incrbyfloat(keyDailyKey, cost); - pipeline.expire(keyDailyKey, ttlDailyKey); + if (options?.userId != null) { + if (userDailyMode === "rolling") { + RateLimitService.queueRollingCostWindow( + pipeline, + `user:${options.userId}:cost_daily_rolling`, + cost, + now, + window24h, + requestId, + 90000 + ); + } else { + const userDailyKey = `user:${options.userId}:cost_daily_${userDailyReset.suffix}`; + pipeline.incrbyfloat(userDailyKey, cost); + pipeline.expire(userDailyKey, ttlDailyUser); + } } + // 3. Key 与 Provider 的周/月固定窗口。User 长周期限额由 lease + PostgreSQL 权威用量负责。 pipeline.incrbyfloat(`key:${keyId}:cost_weekly`, cost); pipeline.expire(`key:${keyId}:cost_weekly`, ttlWeekly); pipeline.incrbyfloat(`key:${keyId}:cost_monthly`, cost); pipeline.expire(`key:${keyId}:cost_monthly`, ttlMonthly); - // Provider 的 daily fixed/周/月消费 - if (providerDailyMode === "fixed") { - const providerDailyKey = `provider:${providerId}:cost_daily_${providerDailyReset.suffix}`; - pipeline.incrbyfloat(providerDailyKey, cost); - pipeline.expire(providerDailyKey, ttlDailyProvider); - } - pipeline.incrbyfloat(`provider:${providerId}:cost_weekly`, cost); pipeline.expire(`provider:${providerId}:cost_weekly`, ttlWeekly); pipeline.incrbyfloat(`provider:${providerId}:cost_monthly`, cost); pipeline.expire(`provider:${providerId}:cost_monthly`, ttlMonthly); - await pipeline.exec(); + const results = await pipeline.exec(); + RateLimitService.logCostPipelineErrors(results, "trackCost"); logger.debug(`[RateLimit] Tracked cost: key=${keyId}, provider=${providerId}, cost=${cost}`); } catch (error) { @@ -1501,42 +1574,47 @@ export class RateLimitService { cost: number, resetTime?: string, resetMode?: DailyResetMode, - options?: { requestId?: number; createdAtMs?: number } + options?: { requestId?: string | number; createdAtMs?: number } ): Promise { - if (!RateLimitService.redis || cost <= 0) return; + const redis = RateLimitService.redis; + if (!redis || redis.status !== "ready" || cost <= 0) return; const mode = resetMode ?? "fixed"; const normalizedResetTime = normalizeResetTime(resetTime); try { + const pipeline = redis.pipeline(); + if (mode === "rolling") { - // Rolling 模式:使用 ZSET + Lua 脚本 const key = `user:${userId}:cost_daily_rolling`; const now = options?.createdAtMs ?? Date.now(); const window24h = 24 * 60 * 60 * 1000; const requestId = options?.requestId != null ? String(options.requestId) : ""; - await RateLimitService.redis.eval( - TRACK_COST_DAILY_ROLLING_WINDOW, - 1, + RateLimitService.queueRollingCostWindow( + pipeline, key, - cost.toString(), - now.toString(), - window24h.toString(), - requestId + cost, + now, + window24h, + requestId, + 90000 ); logger.debug(`[RateLimit] Tracked user daily cost (rolling): user=${userId}, cost=${cost}`); } else { - // Fixed 模式:使用 STRING 类型 const suffix = normalizedResetTime.replace(":", ""); const key = `user:${userId}:cost_daily_${suffix}`; const ttl = await getTTLForPeriodWithMode("daily", normalizedResetTime, "fixed"); - await RateLimitService.redis.pipeline().incrbyfloat(key, cost).expire(key, ttl).exec(); + pipeline.incrbyfloat(key, cost); + pipeline.expire(key, ttl); logger.debug(`[RateLimit] Tracked user daily cost (fixed): user=${userId}, cost=${cost}`); } + + const results = await pipeline.exec(); + RateLimitService.logCostPipelineErrors(results, "trackUserDailyCost"); } catch (error) { logger.error(`[RateLimit] Failed to track user daily cost:`, error); } @@ -1828,4 +1906,13 @@ export class RateLimitService { resetMode: options?.resetMode, }); } + + /** + * Settle one request's actual cost against the fixed twelve lease budgets. + */ + static async settleLeaseBudgets( + params: SettleLeaseBudgetsParams + ): Promise { + return LeaseService.settleLeaseBudgets(params); + } } diff --git a/src/lib/redis/client.ts b/src/lib/redis/client.ts index af7b1506f..21bdc38dc 100644 --- a/src/lib/redis/client.ts +++ b/src/lib/redis/client.ts @@ -1,4 +1,5 @@ import Redis, { type RedisOptions } from "ioredis"; +import { getEnvConfig } from "@/lib/config/env.schema"; import { logger } from "@/lib/logger"; let redisClient: Redis | null = null; @@ -46,6 +47,7 @@ export function buildRedisOptionsForUrl(redisUrl: string): { isTLS: boolean; options: RedisOptions; } { + const env = getEnvConfig(); const isTLS = (() => { try { const parsed = new URL(redisUrl); @@ -59,6 +61,14 @@ export function buildRedisOptionsForUrl(redisUrl: string): { const baseOptions: RedisOptions = { enableOfflineQueue: false, // 快速失败 maxRetriesPerRequest: 3, + commandTimeout: env.REDIS_COMMAND_TIMEOUT_MS, + // commandTimeout only rejects the caller Promise; it does not remove a sent + // command from ioredis' RESP-order queue. Destroy a no-progress socket shortly + // afterwards so the shared queue cannot grow without bound under a TCP blackhole. + socketTimeout: env.REDIS_COMMAND_TIMEOUT_MS + 5_000, + // Timed-out writes have already been treated as fail-open by the application. + // Replaying them after reconnect would mutate Redis after the request completed. + autoResendUnfulfilledCommands: false, retryStrategy(times: number) { if (times > 5) { logger.error("[Redis] Max retries reached, giving up"); diff --git a/src/lib/redis/lua-scripts.ts b/src/lib/redis/lua-scripts.ts index 7513e4119..c0da78e89 100644 --- a/src/lib/redis/lua-scripts.ts +++ b/src/lib/redis/lua-scripts.ts @@ -265,30 +265,33 @@ return results `; /** - * 追踪 5小时滚动窗口消费(使用 ZSET) + * 追踪滚动窗口消费(写路径专用,使用 ZSET) * - * 功能: - * 1. 清理 5 小时前的消费记录 - * 2. 添加当前消费记录(带时间戳) - * 3. 计算当前窗口内的总消费 - * 4. 设置兜底 TTL(6 小时) + * 写路径只负责清理、追加和恢复 TTL。精确总额由 GET 脚本在真正需要 + * 限额判断时计算,避免每次成功请求都扫描整个窗口。 * - * KEYS[1]: key:${id}:cost_5h_rolling 或 provider:${id}:cost_5h_rolling + * KEYS[1]: {entity}:${id}:cost_{window}_rolling * ARGV[1]: cost(本次消费金额) * ARGV[2]: now(当前时间戳,毫秒) - * ARGV[3]: window(窗口时长,毫秒,默认 18000000 = 5小时) - * ARGV[4]: request_id(可选,用于 member 去重) + * ARGV[3]: window(窗口时长,毫秒) + * ARGV[4]: request_id(可选,用于相同时间轴上的 member 去重) + * ARGV[5]: ttl_seconds(兜底 TTL,秒) * - * 返回值:string - 当前窗口内的总消费 + * 返回值:integer - 1 表示写入完成 */ -export const TRACK_COST_5H_ROLLING_WINDOW = ` +export const TRACK_COST_ROLLING_WINDOW = ` local key = KEYS[1] local cost = tonumber(ARGV[1]) local now_ms = tonumber(ARGV[2]) -local window_ms = tonumber(ARGV[3]) -- 5 hours = 18000000 ms +local window_ms = tonumber(ARGV[3]) local request_id = ARGV[4] +local ttl_seconds = tonumber(ARGV[5]) --- 1. 清理过期记录(5 小时前的数据) +if not cost or not now_ms or not window_ms or not ttl_seconds then + return redis.error_reply('invalid rolling cost arguments') +end + +-- 1. 清理窗口外的消费记录 redis.call('ZREMRANGEBYSCORE', key, '-inf', now_ms - window_ms) -- 2. 添加当前消费记录(member = timestamp:cost 或 timestamp:requestId:cost,便于调试和追踪) @@ -300,21 +303,10 @@ else end redis.call('ZADD', key, now_ms, member) --- 3. 计算窗口内总消费 -local records = redis.call('ZRANGE', key, 0, -1) -local total = 0 -for _, record in ipairs(records) do - -- 解析 member 格式:"timestamp:cost" 或 "timestamp:id:cost" - local cost_str = string.match(record, '.*:(.+)') - if cost_str then - total = total + tonumber(cost_str) - end -end - --- 4. 设置兜底 TTL(6 小时,防止数据永久堆积) -redis.call('EXPIRE', key, 21600) +-- 3. 恢复兜底 TTL,允许写路径修复缺失 TTL 的合法或脏 ZSET +redis.call('EXPIRE', key, ttl_seconds) -return tostring(total) +return 1 `; /** @@ -351,59 +343,6 @@ end return tostring(total) `; -/** - * 追踪 24小时滚动窗口消费(使用 ZSET) - * - * 功能: - * 1. 清理 24 小时前的消费记录 - * 2. 添加当前消费记录(带时间戳) - * 3. 计算当前窗口内的总消费 - * 4. 设置兜底 TTL(25 小时) - * - * KEYS[1]: key:${id}:cost_daily_rolling 或 provider:${id}:cost_daily_rolling - * ARGV[1]: cost(本次消费金额) - * ARGV[2]: now(当前时间戳,毫秒) - * ARGV[3]: window(窗口时长,毫秒,默认 86400000 = 24小时) - * ARGV[4]: request_id(可选,用于 member 去重) - * - * 返回值:string - 当前窗口内的总消费 - */ -export const TRACK_COST_DAILY_ROLLING_WINDOW = ` -local key = KEYS[1] -local cost = tonumber(ARGV[1]) -local now_ms = tonumber(ARGV[2]) -local window_ms = tonumber(ARGV[3]) -- 24 hours = 86400000 ms -local request_id = ARGV[4] - --- 1. 清理过期记录(24 小时前的数据) -redis.call('ZREMRANGEBYSCORE', key, '-inf', now_ms - window_ms) - --- 2. 添加当前消费记录(member = timestamp:cost 或 timestamp:requestId:cost,便于调试和追踪) -local member -if request_id and request_id ~= '' then - member = now_ms .. ':' .. request_id .. ':' .. cost -else - member = now_ms .. ':' .. cost -end -redis.call('ZADD', key, now_ms, member) - --- 3. 计算窗口内总消费 -local records = redis.call('ZRANGE', key, 0, -1) -local total = 0 -for _, record in ipairs(records) do - -- 解析 member 格式:"timestamp:cost" 或 "timestamp:id:cost" - local cost_str = string.match(record, '.*:(.+)') - if cost_str then - total = total + tonumber(cost_str) - end -end - --- 4. 设置兜底 TTL(25 小时,防止数据永久堆积) -redis.call('EXPIRE', key, 90000) - -return tostring(total) -`; - /** * 查询 24小时滚动窗口当前消费 * diff --git a/src/lib/session-manager.ts b/src/lib/session-manager.ts index eaec22ae8..d44272252 100644 --- a/src/lib/session-manager.ts +++ b/src/lib/session-manager.ts @@ -3,6 +3,7 @@ import "server-only"; import crypto from "node:crypto"; import { extractCodexSessionId } from "@/app/v1/_lib/codex/session-extractor"; import { sanitizeHeaders, sanitizeUrl } from "@/app/v1/_lib/proxy/errors"; +import { RESERVED_INTERNAL_HEADERS } from "@/app/v1/_lib/responses-ws/internal-secret"; import { parseClaudeMetadataUserId } from "@/lib/claude-code/metadata-user-id"; import { getEnvConfig } from "@/lib/config/env.schema"; import { logger } from "@/lib/logger"; @@ -32,6 +33,15 @@ import { } from "./redis/active-session-keys"; import { SessionTracker } from "./session-tracker"; +const RESERVED_INTERNAL_HEADER_SET = new Set( + RESERVED_INTERNAL_HEADERS.map((header) => header.toLowerCase()) +); + +function isReservedInternalHeader(name: string): boolean { + const lowerName = name.toLowerCase(); + return lowerName.startsWith("x-cch-") || RESERVED_INTERNAL_HEADER_SET.has(lowerName); +} + /** * 将已脱敏的 header 文本解析为可序列化对象(用于写入 Session 元信息)。 */ @@ -73,7 +83,7 @@ function parseHeaderRecord(value: string): Record | null { const record: Record = {}; for (const [key, raw] of Object.entries(parsed as Record)) { - if (typeof raw === "string") { + if (typeof raw === "string" && !isReservedInternalHeader(key)) { record[key] = raw; } } @@ -128,7 +138,9 @@ function normalizeSnapshotHeaders( } const normalized = Object.fromEntries( - Object.entries(headers).filter(([, value]) => typeof value === "string") + Object.entries(headers).filter( + ([key, value]) => typeof value === "string" && !isReservedInternalHeader(key) + ) ); return Object.keys(normalized).length > 0 ? normalized : null; } @@ -685,15 +697,40 @@ export class SessionManager { /** * 清除 session 绑定的 provider(用于跨模型 session 绑定过时时) */ - static async clearSessionProvider(sessionId: string): Promise { + static async clearSessionProvider( + sessionId: string, + expectedProviderId?: number | null + ): Promise { const redis = getRedisClient(); - if (!redis || redis.status !== "ready") return; + if (!redis || redis.status !== "ready") return false; try { - await redis.del(`session:${sessionId}:provider`); - logger.trace("SessionManager: Cleared session provider binding", { sessionId }); + const key = `session:${sessionId}:provider`; + const deleted = + expectedProviderId == null + ? await redis.del(key) + : Number( + await redis.eval( + ` + if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("DEL", KEYS[1]) + end + return 0 + `, + 1, + key, + expectedProviderId.toString() + ) + ); + logger.trace("SessionManager: Cleared session provider binding", { + sessionId, + expectedProviderId: expectedProviderId ?? null, + deleted: deleted > 0, + }); + return deleted > 0; } catch (error) { logger.error("SessionManager: Failed to clear session provider", { error, sessionId }); + return false; } } diff --git a/src/lib/utils/upstream-error-detection.test.ts b/src/lib/utils/upstream-error-detection.test.ts index 957ef374b..1b35ad6eb 100644 --- a/src/lib/utils/upstream-error-detection.test.ts +++ b/src/lib/utils/upstream-error-detection.test.ts @@ -74,13 +74,13 @@ describe("detectUpstreamErrorFromSseOrJsonText", () => { expect(res.isError).toBe(true); }); - test.each([ - '{"error":true}', - '{"error":42}', - ])("纯 JSON:error 为非字符串类型也应视为错误(%s)", (body) => { - const res = detectUpstreamErrorFromSseOrJsonText(body); - expect(res.isError).toBe(true); - }); + test.each(['{"error":true}', '{"error":42}'])( + "纯 JSON:error 为非字符串类型也应视为错误(%s)", + (body) => { + const res = detectUpstreamErrorFromSseOrJsonText(body); + expect(res.isError).toBe(true); + } + ); test("JSON 数组输入不视为错误(目前不做解析)", () => { const res = detectUpstreamErrorFromSseOrJsonText('[{"error":"something"}]'); diff --git a/src/repository/message-write-buffer.ts b/src/repository/message-write-buffer.ts index cab71d102..e58f871e8 100644 --- a/src/repository/message-write-buffer.ts +++ b/src/repository/message-write-buffer.ts @@ -2,7 +2,8 @@ import "server-only"; import type { SQL } from "drizzle-orm"; import { sql } from "drizzle-orm"; -import { db } from "@/drizzle/db"; +import { findSafeDatabaseError } from "@/drizzle/admitted-client"; +import { getMessageWriterDb } from "@/drizzle/db"; import { getEnvConfig } from "@/lib/config/env.schema"; import { logger } from "@/lib/logger"; import type { StoredCostBreakdown } from "@/types/cost-breakdown"; @@ -43,12 +44,180 @@ export type MessageRequestUpdateRecord = { patch: MessageRequestUpdatePatch; }; +export type DurableMessageRequestUpdateOptions = { + timeoutMs?: number; + onCommitted?: (patch: Readonly) => void | Promise; +}; + +type DurableAcknowledgement = { + id: number; + promise: Promise; + resolve: () => void; + reject: (error: Error) => void; + state: "pending" | "in-flight"; + settled: boolean; + timeoutId: NodeJS.Timeout | null; + commitNotified: boolean; + onCommittedCallbacks: Set>; +}; + +type PendingMessageRequestUpdate = { + patch: MessageRequestUpdatePatch; + durableAcknowledgement?: DurableAcknowledgement; +}; + +type MessageRequestUpdateBatchRecord = MessageRequestUpdateRecord & { + durableAcknowledgement?: DurableAcknowledgement; +}; + type WriterConfig = { flushIntervalMs: number; batchSize: number; maxPending: number; }; +const DEFAULT_DURABLE_ACK_TIMEOUT_MS = 120_000; +const OVERFLOW_LOG_AGGREGATION_MS = 1_000; + +type EvictablePendingEntry = { + id: number; + priority: number; + order: number; +}; + +class EvictablePendingIndex { + private readonly heap: EvictablePendingEntry[] = []; + private readonly positions = new Map(); + private nextOrder = 0; + + upsert(id: number, priority: number): void { + const position = this.positions.get(id); + if (position === undefined) { + const entry = { id, priority, order: this.nextOrder++ }; + this.heap.push(entry); + this.positions.set(id, this.heap.length - 1); + this.bubbleUp(this.heap.length - 1); + return; + } + + const entry = this.heap[position]; + if (!entry || entry.priority === priority) { + return; + } + + const previousPriority = entry.priority; + entry.priority = priority; + if (priority < previousPriority) { + this.bubbleUp(position); + } else { + this.bubbleDown(position); + } + } + + remove(id: number): EvictablePendingEntry | undefined { + const position = this.positions.get(id); + if (position === undefined) { + return undefined; + } + return this.removeAt(position); + } + + popLowestPriority(): EvictablePendingEntry | undefined { + if (this.heap.length === 0) { + return undefined; + } + return this.removeAt(0); + } + + clear(): void { + this.heap.length = 0; + this.positions.clear(); + } + + private removeAt(position: number): EvictablePendingEntry | undefined { + const removed = this.heap[position]; + if (!removed) { + return undefined; + } + + const last = this.heap.pop(); + this.positions.delete(removed.id); + if (last && position < this.heap.length) { + this.heap[position] = last; + this.positions.set(last.id, position); + const parentPosition = Math.floor((position - 1) / 2); + if (position > 0 && this.isLowerPriority(last, this.heap[parentPosition])) { + this.bubbleUp(position); + } else { + this.bubbleDown(position); + } + } + return removed; + } + + private bubbleUp(startPosition: number): void { + let position = startPosition; + while (position > 0) { + const parentPosition = Math.floor((position - 1) / 2); + const entry = this.heap[position]; + const parent = this.heap[parentPosition]; + if (!entry || !parent || !this.isLowerPriority(entry, parent)) { + break; + } + this.swap(position, parentPosition); + position = parentPosition; + } + } + + private bubbleDown(startPosition: number): void { + let position = startPosition; + while (true) { + const leftPosition = position * 2 + 1; + const rightPosition = leftPosition + 1; + let lowestPosition = position; + + if ( + this.heap[leftPosition] && + this.heap[lowestPosition] && + this.isLowerPriority(this.heap[leftPosition], this.heap[lowestPosition]) + ) { + lowestPosition = leftPosition; + } + if ( + this.heap[rightPosition] && + this.heap[lowestPosition] && + this.isLowerPriority(this.heap[rightPosition], this.heap[lowestPosition]) + ) { + lowestPosition = rightPosition; + } + if (lowestPosition === position) { + return; + } + this.swap(position, lowestPosition); + position = lowestPosition; + } + } + + private swap(firstPosition: number, secondPosition: number): void { + const first = this.heap[firstPosition]; + const second = this.heap[secondPosition]; + if (!first || !second) { + return; + } + this.heap[firstPosition] = second; + this.heap[secondPosition] = first; + this.positions.set(first.id, secondPosition); + this.positions.set(second.id, firstPosition); + } + + private isLowerPriority(first: EvictablePendingEntry, second: EvictablePendingEntry): boolean { + return ( + first.priority < second.priority || + (first.priority === second.priority && first.order < second.order) + ); + } +} + const COLUMN_MAP: Record = { durationMs: "duration_ms", costUsd: "cost_usd", @@ -83,10 +252,22 @@ function loadWriterConfig(): WriterConfig { }; } -function takeBatch(map: Map, batchSize: number) { - const items: MessageRequestUpdateRecord[] = []; - for (const [id, patch] of map) { - items.push({ id, patch }); +function takeBatch( + map: Map, + evictableIndex: EvictablePendingIndex, + batchSize: number +): MessageRequestUpdateBatchRecord[] { + const items: MessageRequestUpdateBatchRecord[] = []; + for (const [id, pending] of map) { + if (pending.durableAcknowledgement && !pending.durableAcknowledgement.settled) { + pending.durableAcknowledgement.state = "in-flight"; + } + items.push({ + id, + patch: pending.patch, + durableAcknowledgement: pending.durableAcknowledgement, + }); + evictableIndex.remove(id); map.delete(id); if (items.length >= batchSize) { break; @@ -95,7 +276,10 @@ function takeBatch(map: Map, batchSize: numbe return items; } -export function buildBatchUpdateSql(updates: MessageRequestUpdateRecord[]): SQL | null { +export function buildBatchUpdateSql( + updates: MessageRequestUpdateRecord[], + options: { returnUpdatedIds?: boolean; durableIds?: readonly number[] } = {} +): SQL | null { if (updates.length === 0) { return null; } @@ -152,12 +336,26 @@ export function buildBatchUpdateSql(updates: MessageRequestUpdateRecord[]): SQL ids.map((id) => sql`${id}`), sql`, ` ); - - return sql` + const updateIds = new Set(ids); + const durableIds = Array.from(new Set(options.durableIds ?? [])).filter((id) => + updateIds.has(id) + ); + const durableFence = + durableIds.length === 0 + ? sql`` + : durableIds.length === ids.length + ? sql` AND ${sql.identifier("status_code")} IS NULL` + : sql` AND (id NOT IN (${sql.join( + durableIds.map((id) => sql`${id}`), + sql`, ` + )}) OR ${sql.identifier("status_code")} IS NULL)`; + + const query = sql` UPDATE message_request SET ${sql.join(setClauses, sql`, `)} - WHERE id IN (${idList}) AND deleted_at IS NULL + WHERE id IN (${idList}) AND deleted_at IS NULL${durableFence} `; + return options.returnUpdatedIds ? sql`${query} RETURNING id` : query; } /** @@ -195,10 +393,19 @@ function getPatchRetentionPriority(patch: MessageRequestUpdatePatch): number { class MessageRequestWriteBuffer { private readonly config: WriterConfig; - private readonly pending = new Map(); + private readonly pending = new Map(); + private readonly evictableIndex = new EvictablePendingIndex(); + private readonly durableAcknowledgements = new Map(); private flushTimer: NodeJS.Timeout | null = null; + private overflowLogTimer: NodeJS.Timeout | null = null; + private overflowDroppedCount = 0; + private overflowDroppedWithDurationMs = 0; + private overflowDroppedWithStatusCode = 0; + private overflowLowestPriority = Number.POSITIVE_INFINITY; + private overflowLastDroppedId: number | undefined; private flushAgainAfterCurrent = false; private flushInFlight: Promise | null = null; + private readonly commitCallbacksInFlight = new Set>(); private stopping = false; constructor(config: WriterConfig) { @@ -206,51 +413,262 @@ class MessageRequestWriteBuffer { } enqueue(id: number, patch: MessageRequestUpdatePatch): void { - const existing = this.pending.get(id) ?? {}; + const existing = this.pending.get(id); // existing is older, patch is newer -> for replacement fields newer wins. - this.pending.set(id, mergePatch(existing, patch)); - - // 队列上限保护:DB 异常时避免无限增长导致 OOM - if (this.pending.size > this.config.maxPending) { - // 优先保留更接近终态的 patch: - // statusCode > durationMs > metadata-only - // 这样 Gemini passthrough 等 statusCode-only 终态更新不会比 duration-only 更容易被丢弃。 - let droppedId: number | undefined; - let droppedPatch: MessageRequestUpdatePatch | undefined; - let lowestPriority = Number.POSITIVE_INFINITY; - - for (const [candidateId, candidatePatch] of this.pending) { - const priority = getPatchRetentionPriority(candidatePatch); - if (priority < lowestPriority) { - lowestPriority = priority; - droppedId = candidateId; - droppedPatch = candidatePatch; - } + this.setPending(id, mergePatch(existing?.patch ?? {}, patch), existing?.durableAcknowledgement); + + this.enforcePendingLimit(); + this.scheduleFlushIfNeeded(); + } + + enqueueDurably( + id: number, + patch: MessageRequestUpdatePatch, + options: DurableMessageRequestUpdateOptions = {} + ): Promise { + if (this.stopping) { + return Promise.reject(new Error("message_request writer is stopping")); + } + + const activeAcknowledgement = this.durableAcknowledgements.get(id); + if (activeAcknowledgement && !activeAcknowledgement.settled) { + // The first durable claimant owns both the terminal patch and its commit + // callback. Later contenders may observe its SQL acknowledgement, but + // must not merge a contradictory terminal outcome or publish side effects. + return activeAcknowledgement.promise.then(() => false); + } + + if (this.durableAcknowledgements.size >= this.config.maxPending) { + return Promise.reject(new Error("durable message_request queue is full")); + } + + const acknowledgement = this.createDurableAcknowledgement(id, options); + const existing = this.pending.get(id); + this.setPending(id, mergePatch(existing?.patch ?? {}, patch), acknowledgement); + + if (!this.enforcePendingLimit()) { + this.deletePending(id); + this.rejectDurableAcknowledgement( + acknowledgement, + new Error("durable message_request queue is full") + ); + return acknowledgement.promise.then(() => true); + } + + this.scheduleFlushIfNeeded(); + return acknowledgement.promise.then(() => true); + } + + private createDurableAcknowledgement( + id: number, + options: DurableMessageRequestUpdateOptions + ): DurableAcknowledgement { + let resolvePromise!: () => void; + let rejectPromise!: (error: Error) => void; + const promise = new Promise((resolve, reject) => { + resolvePromise = resolve; + rejectPromise = reject; + }); + const acknowledgement: DurableAcknowledgement = { + id, + promise, + resolve: resolvePromise, + reject: rejectPromise, + state: "pending", + settled: false, + timeoutId: null, + commitNotified: false, + onCommittedCallbacks: new Set(options.onCommitted ? [options.onCommitted] : []), + }; + + const timeoutMs = options.timeoutMs ?? DEFAULT_DURABLE_ACK_TIMEOUT_MS; + const effectiveTimeoutMs = + Number.isFinite(timeoutMs) && timeoutMs > 0 ? timeoutMs : DEFAULT_DURABLE_ACK_TIMEOUT_MS; + acknowledgement.timeoutId = setTimeout(() => { + const pending = this.pending.get(id); + if (pending?.durableAcknowledgement === acknowledgement) { + this.deletePending(id); } + this.rejectDurableAcknowledgement( + acknowledgement, + new Error("durable message_request acknowledgement timed out") + ); + }, effectiveTimeoutMs); + acknowledgement.timeoutId.unref?.(); + + this.durableAcknowledgements.set(id, acknowledgement); + return acknowledgement; + } - if (droppedId === undefined) { - const first = this.pending.entries().next().value as - | [number, MessageRequestUpdatePatch] - | undefined; - if (first) { - droppedId = first[0]; - droppedPatch = first[1]; + private notifyDurableCommit( + acknowledgement: DurableAcknowledgement | undefined, + patch: Readonly + ): void { + if (!acknowledgement || acknowledgement.commitNotified) return; + acknowledgement.commitNotified = true; + + for (const callback of acknowledgement.onCommittedCallbacks) { + try { + const result = callback(patch); + if (result && typeof result.then === "function") { + let callbackPromise: Promise; + callbackPromise = Promise.resolve(result) + .catch((error: unknown) => { + logger.error("[MessageRequestWriteBuffer] Durable commit callback failed", { + error: error instanceof Error ? error.message : String(error), + messageRequestId: acknowledgement.id, + }); + }) + .finally(() => { + this.commitCallbacksInFlight.delete(callbackPromise); + }); + this.commitCallbacksInFlight.add(callbackPromise); } + } catch (error) { + logger.error("[MessageRequestWriteBuffer] Durable commit callback failed", { + error: error instanceof Error ? error.message : String(error), + messageRequestId: acknowledgement.id, + }); } + } + acknowledgement.onCommittedCallbacks.clear(); + } - if (droppedId !== undefined) { - this.pending.delete(droppedId); - logger.warn("[MessageRequestWriteBuffer] Pending queue overflow, dropping update", { - maxPending: this.config.maxPending, - droppedId, - droppedPriority: lowestPriority, - droppedHasDurationMs: droppedPatch?.durationMs !== undefined, - droppedHasStatusCode: droppedPatch?.statusCode !== undefined, - currentPending: this.pending.size, - }); + private resolveDurableAcknowledgement(acknowledgement?: DurableAcknowledgement): void { + if (!acknowledgement || acknowledgement.settled) return; + acknowledgement.settled = true; + if (acknowledgement.timeoutId) { + clearTimeout(acknowledgement.timeoutId); + acknowledgement.timeoutId = null; + } + if (this.durableAcknowledgements.get(acknowledgement.id) === acknowledgement) { + this.durableAcknowledgements.delete(acknowledgement.id); + } + acknowledgement.resolve(); + } + + private rejectDurableAcknowledgement( + acknowledgement: DurableAcknowledgement | undefined, + error: Error + ): void { + if (!acknowledgement || acknowledgement.settled) return; + acknowledgement.settled = true; + if (acknowledgement.timeoutId) { + clearTimeout(acknowledgement.timeoutId); + acknowledgement.timeoutId = null; + } + if (this.durableAcknowledgements.get(acknowledgement.id) === acknowledgement) { + this.durableAcknowledgements.delete(acknowledgement.id); + } + acknowledgement.reject(error); + } + + private rejectAllDurableAcknowledgements(error: Error): void { + for (const acknowledgement of this.durableAcknowledgements.values()) { + this.rejectDurableAcknowledgement(acknowledgement, error); + } + } + + private setPending( + id: number, + patch: MessageRequestUpdatePatch, + durableAcknowledgement?: DurableAcknowledgement + ): void { + const activeDurableAcknowledgement = + durableAcknowledgement && !durableAcknowledgement.settled + ? durableAcknowledgement + : undefined; + this.pending.set(id, { + patch, + durableAcknowledgement: activeDurableAcknowledgement, + }); + if (activeDurableAcknowledgement) { + this.evictableIndex.remove(id); + } else { + this.evictableIndex.upsert(id, getPatchRetentionPriority(patch)); + } + } + + private deletePending(id: number): PendingMessageRequestUpdate | undefined { + const pending = this.pending.get(id); + if (!pending) { + return undefined; + } + this.pending.delete(id); + this.evictableIndex.remove(id); + return pending; + } + + private enforcePendingLimit(): boolean { + while (this.pending.size > this.config.maxPending) { + const droppedEntry = this.evictableIndex.popLowestPriority(); + if (!droppedEntry) { + return false; + } + const dropped = this.pending.get(droppedEntry.id); + if (!dropped || (dropped.durableAcknowledgement && !dropped.durableAcknowledgement.settled)) { + continue; } + + this.pending.delete(droppedEntry.id); + this.recordOverflowDrop(droppedEntry, dropped.patch); } + return true; + } + + private recordOverflowDrop( + droppedEntry: EvictablePendingEntry, + droppedPatch: MessageRequestUpdatePatch + ): void { + this.overflowDroppedCount++; + this.overflowLastDroppedId = droppedEntry.id; + this.overflowLowestPriority = Math.min(this.overflowLowestPriority, droppedEntry.priority); + if (droppedPatch.durationMs !== undefined) { + this.overflowDroppedWithDurationMs++; + } + if (droppedPatch.statusCode !== undefined) { + this.overflowDroppedWithStatusCode++; + } + if (this.overflowLogTimer) { + return; + } + this.overflowLogTimer = setTimeout(() => { + this.overflowLogTimer = null; + this.flushOverflowLog(); + }, OVERFLOW_LOG_AGGREGATION_MS); + this.overflowLogTimer.unref?.(); + } + + private flushOverflowLog(): void { + if (this.overflowDroppedCount === 0) { + return; + } + logger.warn("[MessageRequestWriteBuffer] Pending queue overflow, dropping updates", { + maxPending: this.config.maxPending, + droppedCount: this.overflowDroppedCount, + lowestDroppedPriority: this.overflowLowestPriority, + droppedWithDurationMs: this.overflowDroppedWithDurationMs, + droppedWithStatusCode: this.overflowDroppedWithStatusCode, + lastDroppedId: this.overflowLastDroppedId, + currentPending: this.pending.size, + }); + this.overflowDroppedCount = 0; + this.overflowDroppedWithDurationMs = 0; + this.overflowDroppedWithStatusCode = 0; + this.overflowLowestPriority = Number.POSITIVE_INFINITY; + this.overflowLastDroppedId = undefined; + } + + private clearOverflowLogTimer(): void { + if (!this.overflowLogTimer) { + return; + } + clearTimeout(this.overflowLogTimer); + this.overflowLogTimer = null; + } + + private scheduleFlushIfNeeded(): void { // flush 过程中有新任务:标记需要再跑一轮(避免刚好 flush 完成时遗漏) if (this.flushInFlight) { this.flushAgainAfterCurrent = true; @@ -300,24 +718,81 @@ class MessageRequestWriteBuffer { this.flushAgainAfterCurrent = false; while (this.pending.size > 0) { - const batch = takeBatch(this.pending, this.config.batchSize); - const query = buildBatchUpdateSql(batch); + const batch = takeBatch(this.pending, this.evictableIndex, this.config.batchSize); + const requiresUpdatedIds = batch.some( + (item) => item.durableAcknowledgement && !item.durableAcknowledgement.settled + ); + const durableIds = batch.flatMap((item) => + item.durableAcknowledgement ? [item.id] : [] + ); + const query = buildBatchUpdateSql(batch, { + returnUpdatedIds: requiresUpdatedIds, + durableIds, + }); if (!query) { + for (const item of batch) { + this.rejectDurableAcknowledgement( + item.durableAcknowledgement, + new Error("durable message_request update contains no writable fields") + ); + } continue; } try { - await db.execute(query); + const result = await getMessageWriterDb().execute(query); + const updatedIds = new Set( + requiresUpdatedIds + ? Array.from(result, (row) => Number((row as { id?: unknown }).id)) + : [] + ); + for (const item of batch) { + const acknowledgement = item.durableAcknowledgement; + if (!acknowledgement) { + continue; + } + if (updatedIds.has(item.id)) { + this.notifyDurableCommit(acknowledgement, item.patch); + if (!acknowledgement.settled) { + this.resolveDurableAcknowledgement(acknowledgement); + } + } else if (!acknowledgement.settled) { + this.rejectDurableAcknowledgement( + acknowledgement, + new Error(`durable message_request update did not persist id ${item.id}`) + ); + } + } } catch (error) { // 失败重试:将 batch 放回队列 // 合并策略:保留“更新更晚”的字段(existing 优先),避免覆盖新数据。 for (const item of batch) { - const existing = this.pending.get(item.id) ?? {}; - this.pending.set(item.id, mergePatch(item.patch, existing)); + if (item.durableAcknowledgement?.settled) { + continue; + } + const existing = this.pending.get(item.id); + const durableAcknowledgement = + item.durableAcknowledgement && !item.durableAcknowledgement.settled + ? item.durableAcknowledgement + : existing?.durableAcknowledgement; + if (durableAcknowledgement) { + durableAcknowledgement.state = "pending"; + } + this.setPending( + item.id, + mergePatch(item.patch, existing?.patch ?? {}), + durableAcknowledgement + ); } + this.enforcePendingLimit(); + const databaseError = findSafeDatabaseError(error); logger.error("[MessageRequestWriteBuffer] Flush failed, will retry later", { - error: error instanceof Error ? error.message : String(error), + error: + databaseError?.message ?? (error instanceof Error ? error.message : String(error)), + databaseCode: databaseError?.code, + admissionPool: databaseError?.pool, + admissionMaxOutstanding: databaseError?.maxOutstanding, pending: this.pending.size, batchSize: batch.length, }); @@ -346,17 +821,39 @@ class MessageRequestWriteBuffer { if (this.pending.size > 0) { await this.flush(); } + if (this.pending.size > 0) { + const error = new Error("message_request writer shutdown persistence failed"); + this.rejectAllDurableAcknowledgements(error); + this.clearOverflowLogTimer(); + this.flushOverflowLog(); + this.pending.clear(); + this.evictableIndex.clear(); + throw error; + } + if (this.durableAcknowledgements.size > 0) { + this.rejectAllDurableAcknowledgements( + new Error("message_request writer stopped before durable commit") + ); + } + while (this.commitCallbacksInFlight.size > 0) { + await Promise.allSettled([...this.commitCallbacksInFlight]); + } + this.clearOverflowLogTimer(); + this.flushOverflowLog(); + this.pending.clear(); + this.evictableIndex.clear(); } } let _buffer: MessageRequestWriteBuffer | null = null; let _bufferState: "running" | "stopping" | "stopped" = "running"; +let _stopPromise: Promise | null = null; function getBuffer(): MessageRequestWriteBuffer | null { + if (_bufferState !== "running") { + return null; + } if (!_buffer) { - if (_bufferState !== "running") { - return null; - } _buffer = new MessageRequestWriteBuffer(loadWriterConfig()); } return _buffer; @@ -374,6 +871,23 @@ export function enqueueMessageRequestUpdate(id: number, patch: MessageRequestUpd buffer.enqueue(id, patch); } +export function enqueueMessageRequestUpdateDurably( + id: number, + patch: MessageRequestUpdatePatch, + options?: DurableMessageRequestUpdateOptions +): Promise { + if (getEnvConfig().MESSAGE_REQUEST_WRITE_MODE !== "async") { + return Promise.reject( + new Error("durable message_request buffer API requires async write mode") + ); + } + const buffer = getBuffer(); + if (!buffer) { + return Promise.reject(new Error("message_request writer is not running")); + } + return buffer.enqueueDurably(id, patch, options); +} + export async function flushMessageRequestWriteBuffer(): Promise { if (!_buffer) { return; @@ -381,18 +895,30 @@ export async function flushMessageRequestWriteBuffer(): Promise { await _buffer.flush(); } -export async function stopMessageRequestWriteBuffer(): Promise { - if (_bufferState === "stopped") { - return; +export function stopMessageRequestWriteBuffer(): Promise { + if (_stopPromise) { + return _stopPromise; } _bufferState = "stopping"; - - if (!_buffer) { + const buffer = _buffer; + + let resolveStop!: () => void; + let rejectStop!: (reason?: unknown) => void; + const stopPromise = new Promise((resolve, reject) => { + resolveStop = resolve; + rejectStop = reject; + }); + _stopPromise = stopPromise; + + void (async () => { + if (buffer) { + await buffer.stop(); + if (_buffer === buffer) { + _buffer = null; + } + } _bufferState = "stopped"; - return; - } + })().then(resolveStop, rejectStop); - await _buffer.stop(); - _buffer = null; - _bufferState = "stopped"; + return stopPromise; } diff --git a/src/repository/message.ts b/src/repository/message.ts index 7a3c53cd6..4d72a6944 100644 --- a/src/repository/message.ts +++ b/src/repository/message.ts @@ -1,7 +1,7 @@ "use server"; import { and, asc, desc, eq, gt, inArray, isNull, lt, sql } from "drizzle-orm"; -import { db } from "@/drizzle/db"; +import { db, getMessageWriterDb } from "@/drizzle/db"; import { keys as keysTable, messageRequest, providers, usageLedger, users } from "@/drizzle/schema"; import { getEnvConfig } from "@/lib/config/env.schema"; import { isLedgerOnlyMode } from "@/lib/ledger-fallback"; @@ -17,7 +17,12 @@ import type { SpecialSetting } from "@/types/special-settings"; import { LEDGER_BILLING_CONDITION } from "./_shared/ledger-conditions"; import { EXCLUDE_WARMUP_CONDITION } from "./_shared/message-request-conditions"; import { toMessageRequest } from "./_shared/transformers"; -import { enqueueMessageRequestUpdate } from "./message-write-buffer"; +import { + type DurableMessageRequestUpdateOptions, + enqueueMessageRequestUpdate, + enqueueMessageRequestUpdateDurably, + type MessageRequestUpdatePatch, +} from "./message-write-buffer"; type PublicStatusRequestSeed = { createdAt: Date; @@ -27,6 +32,7 @@ type PublicStatusRequestSeed = { }; type PublicStatusFinalDetails = { + durationMs?: number; statusCode?: number; outputTokens?: number; ttfbMs?: number | null; @@ -137,7 +143,7 @@ async function readPublicStatusRequestSeedFallback( function queuePublicStatusRollupForFinalDetails( id: number, details: PublicStatusFinalDetails -): void { +): Promise | undefined { if (!isPublicStatusFinalDetails(details) || !markPublicStatusRequestInFlight(id)) { return; } @@ -146,7 +152,7 @@ function queuePublicStatusRollupForFinalDetails( return; } - void (async () => { + return (async () => { try { const seed = peekPublicStatusRequestSeed(id) ?? (await readPublicStatusRequestSeedFallback(id)); @@ -201,6 +207,18 @@ function queuePublicStatusRollupForFinalDetails( })(); } +function publishCommittedMessageRequestDetails( + id: number, + details: PublicStatusFinalDetails +): Promise | undefined { + if (details.durationMs !== undefined) { + updatePublicStatusRequestSeed(id, { durationMs: details.durationMs }); + } + if (details.providerChain !== undefined && details.statusCode !== undefined) { + return queuePublicStatusRollupForFinalDetails(id, details); + } +} + /** * 创建消息请求记录 */ @@ -458,48 +476,55 @@ export async function addMessageRequestHedgeLoserCost( throw lastError; } +export type MessageRequestDetailsUpdate = { + durationMs?: number; + statusCode?: number; + inputTokens?: number; + outputTokens?: number; + ttfbMs?: number | null; + cacheCreationInputTokens?: number; + cacheReadInputTokens?: number; + cacheCreation5mInputTokens?: number; + cacheCreation1hInputTokens?: number; + cacheTtlApplied?: string | null; + providerChain?: CreateMessageRequestData["provider_chain"]; + errorMessage?: string; + errorStack?: string; // 完整堆栈信息 + errorCause?: string; // 嵌套错误原因(JSON 格式) + model?: string; // ⭐ 新增:支持更新重定向后的模型名称 + actualResponseModel?: string | null; // 上游响应实际返回的模型名(audit 用途,不影响计费) + providerId?: number; // ⭐ 新增:支持更新最终供应商ID(重试切换后) + context1mApplied?: boolean; // 是否应用了1M上下文窗口 + swapCacheTtlApplied?: boolean; // Swap Cache TTL Billing active at request time + specialSettings?: CreateMessageRequestData["special_settings"]; // 特殊设置(审计/展示) +}; + /** * 更新消息请求的扩展信息(status code, tokens, provider chain, error) */ export async function updateMessageRequestDetails( id: number, - details: { - statusCode?: number; - inputTokens?: number; - outputTokens?: number; - ttfbMs?: number | null; - cacheCreationInputTokens?: number; - cacheReadInputTokens?: number; - cacheCreation5mInputTokens?: number; - cacheCreation1hInputTokens?: number; - cacheTtlApplied?: string | null; - providerChain?: CreateMessageRequestData["provider_chain"]; - errorMessage?: string; - errorStack?: string; // 完整堆栈信息 - errorCause?: string; // 嵌套错误原因(JSON 格式) - model?: string; // ⭐ 新增:支持更新重定向后的模型名称 - actualResponseModel?: string | null; // 上游响应实际返回的模型名(audit 用途,不影响计费) - providerId?: number; // ⭐ 新增:支持更新最终供应商ID(重试切换后) - context1mApplied?: boolean; // 是否应用了1M上下文窗口 - swapCacheTtlApplied?: boolean; // Swap Cache TTL Billing active at request time - specialSettings?: CreateMessageRequestData["special_settings"]; // 特殊设置(审计/展示) - } -): Promise { - const shouldQueuePublicStatusRollup = - details.providerChain !== undefined && details.statusCode !== undefined; - - if (getEnvConfig().MESSAGE_REQUEST_WRITE_MODE === "async") { - enqueueMessageRequestUpdate(id, details); - if (shouldQueuePublicStatusRollup) { - queuePublicStatusRollupForFinalDetails(id, details); + details: MessageRequestDetailsUpdate, + options: { onlyIfUnfinalized?: boolean; awaitCommitObservers?: boolean } = {} +): Promise { + if (getEnvConfig().MESSAGE_REQUEST_WRITE_MODE === "async" && !options.onlyIfUnfinalized) { + // 终态 patch 必须观察 SQL commit 后再发布 public-status rollup。 + // 非终态 metadata 仍保持轻量 enqueue,但不能伪称已提交。 + if (details.statusCode !== undefined) { + await updateMessageRequestDetailsDurably(id, details); + return true; } - return; + enqueueMessageRequestUpdate(id, details); + return true; } const updateData: Record = { updatedAt: new Date(), }; + if (details.durationMs !== undefined) { + updateData.durationMs = details.durationMs; + } if (details.statusCode !== undefined) { updateData.statusCode = details.statusCode; } @@ -558,10 +583,111 @@ export async function updateMessageRequestDetails( updateData.specialSettings = details.specialSettings; } - await db.update(messageRequest).set(updateData).where(eq(messageRequest.id, id)); - if (shouldQueuePublicStatusRollup) { - queuePublicStatusRollupForFinalDetails(id, details); + if (options.onlyIfUnfinalized) { + const terminalDb = + getEnvConfig().MESSAGE_REQUEST_WRITE_MODE === "async" ? getMessageWriterDb() : db; + const updated = await terminalDb + .update(messageRequest) + .set(updateData) + .where(and(eq(messageRequest.id, id), isNull(messageRequest.statusCode))) + .returning({ id: messageRequest.id }); + if (updated.length === 0) { + return false; + } + } else { + await db.update(messageRequest).set(updateData).where(eq(messageRequest.id, id)); + } + const rollupPromise = publishCommittedMessageRequestDetails(id, details); + if (options.awaitCommitObservers === false) { + void rollupPromise; + } else { + await rollupPromise; + } + return true; +} + +export async function updateMessageRequestDetailsIfUnfinalized( + id: number, + details: MessageRequestDetailsUpdate, + options?: Pick +): Promise { + const committed = await updateMessageRequestDetails(id, details, { + onlyIfUnfinalized: true, + awaitCommitObservers: false, + }); + if (committed && options?.onCommitted) { + try { + const callbackResult = options.onCommitted(details); + if (callbackResult && typeof callbackResult.then === "function") { + void Promise.resolve(callbackResult).catch((error) => { + logger.warn("[MessageRequest] Conditional commit callback failed", { + messageRequestId: id, + error: error instanceof Error ? error.message : String(error), + }); + }); + } + } catch (error) { + logger.warn("[MessageRequest] Conditional commit callback failed", { + messageRequestId: id, + error: error instanceof Error ? error.message : String(error), + }); + } + } + return committed; +} + +/** + * Persist terminal request details with an acknowledgement that the backing SQL batch committed. + * Ordinary async metadata updates continue to use updateMessageRequestDetails(). + */ +export async function updateMessageRequestDetailsDurably( + id: number, + details: MessageRequestDetailsUpdate, + options?: DurableMessageRequestUpdateOptions +): Promise { + if (getEnvConfig().MESSAGE_REQUEST_WRITE_MODE !== "async") { + const committed = await updateMessageRequestDetails(id, details, { + onlyIfUnfinalized: true, + awaitCommitObservers: false, + }); + if (committed) { + try { + const callbackResult = options?.onCommitted?.(details); + if (callbackResult && typeof callbackResult.then === "function") { + void Promise.resolve(callbackResult).catch((error) => { + logger.warn("[MessageRequest] onCommitted callback failed", { + messageRequestId: id, + error: error instanceof Error ? error.message : String(error), + }); + }); + } + } catch (error) { + logger.warn("[MessageRequest] onCommitted callback failed", { + messageRequestId: id, + error: error instanceof Error ? error.message : String(error), + }); + } + } + return committed; } + + let commitPublished = false; + const publishCommit = (committedPatch: Readonly) => { + if (commitPublished) return; + commitPublished = true; + const rollupPromise = publishCommittedMessageRequestDetails(id, committedPatch); + const callbackResult = options?.onCommitted?.(committedPatch); + if (rollupPromise && callbackResult) { + return Promise.all([rollupPromise, callbackResult]).then(() => undefined); + } + return callbackResult ?? rollupPromise; + }; + + const committed = await enqueueMessageRequestUpdateDurably(id, details, { + ...options, + onCommitted: publishCommit, + }); + return committed; } /** diff --git a/tests/configs/integration.config.ts b/tests/configs/integration.config.ts index 064a31fab..ca2152353 100644 --- a/tests/configs/integration.config.ts +++ b/tests/configs/integration.config.ts @@ -8,6 +8,12 @@ export default createTestRunnerConfig({ testFiles: [ "tests/integration/usage-ledger.test.ts", "tests/integration/my-usage-imported-ledger.test.ts", + "tests/integration/rolling-cost-redis.test.ts", + "tests/integration/lease-settlement-redis.test.ts", + "tests/integration/db-pool-isolation-postgres.test.ts", + "tests/integration/db-pool-slow-close-postgres.test.ts", + "tests/integration/message-write-buffer-recovery-postgres.test.ts", + "tests/integration/proxy-hedge-lifecycle.test.ts", ], api: { host: process.env.VITEST_API_HOST || "127.0.0.1", diff --git a/tests/integration/billing-model-source.test.ts b/tests/integration/billing-model-source.test.ts index 1caab13b7..a4d223e13 100644 --- a/tests/integration/billing-model-source.test.ts +++ b/tests/integration/billing-model-source.test.ts @@ -3,14 +3,37 @@ import type { ModelPrice, ModelPriceData } from "@/types/model-price"; import type { SystemSettings } from "@/types/system-config"; const asyncTasks: Promise[] = []; +const asyncTaskControllers = new Map, AbortController>(); +let asyncTaskAdmissionOpen = true; const cloudPriceSyncRequests: Array<{ reason: string }> = []; vi.mock("@/lib/async-task-manager", () => ({ AsyncTaskManager: { - register: (_taskId: string, promise: Promise) => { + register: ( + _taskId: string, + factory: (signal: AbortSignal) => Promise, + options?: string | { abortController?: AbortController } + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + if (!asyncTaskAdmissionOpen) { + controller.abort(); + return controller; + } + + let promise: Promise; + try { + promise = Promise.resolve(factory(controller.signal)); + } catch (error) { + promise = Promise.reject(error); + } asyncTasks.push(promise); - return new AbortController(); + asyncTaskControllers.set(promise, controller); + return controller; }, + touch: () => true, cleanup: () => {}, cancel: () => {}, }, @@ -43,6 +66,7 @@ vi.mock("@/repository/system-config", () => ({ vi.mock("@/repository/message", () => ({ updateMessageRequestCostWithBreakdown: vi.fn(), updateMessageRequestDetails: vi.fn(), + updateMessageRequestDetailsDurably: vi.fn(), updateMessageRequestDuration: vi.fn(), })); @@ -60,6 +84,7 @@ vi.mock("@/lib/rate-limit", () => ({ RateLimitService: { trackCost: vi.fn(), trackUserDailyCost: vi.fn(), + settleLeaseBudgets: vi.fn(), }, })); @@ -79,6 +104,7 @@ vi.mock("@/lib/proxy-status-tracker", () => ({ import { finalizeRequestStats, ProxyResponseHandler } from "@/app/v1/_lib/proxy/response-handler"; import { ProxySession } from "@/app/v1/_lib/proxy/session"; +import { AsyncTaskManager } from "@/lib/async-task-manager"; import { getCachedSystemSettings, invalidateSystemSettingsCache } from "@/lib/config"; import { SessionManager } from "@/lib/session-manager"; import { RateLimitService } from "@/lib/rate-limit"; @@ -86,6 +112,7 @@ import { SessionTracker } from "@/lib/session-tracker"; import { updateMessageRequestCostWithBreakdown, updateMessageRequestDetails, + updateMessageRequestDetailsDurably, updateMessageRequestDuration, } from "@/repository/message"; import { findLatestPriceByModel } from "@/repository/model-price"; @@ -93,6 +120,13 @@ import { getSystemSettings } from "@/repository/system-config"; beforeEach(() => { vi.clearAllMocks(); + asyncTaskAdmissionOpen = false; + for (const controller of asyncTaskControllers.values()) { + controller.abort(); + } + asyncTasks.splice(0, asyncTasks.length); + asyncTaskControllers.clear(); + asyncTaskAdmissionOpen = true; cloudPriceSyncRequests.splice(0, cloudPriceSyncRequests.length); invalidateSystemSettingsCache(); }); @@ -312,10 +346,195 @@ function createStreamResponse(usage: { input_tokens: number; output_tokens: numb } async function drainAsyncTasks(): Promise { - const tasks = asyncTasks.splice(0, asyncTasks.length); - await Promise.all(tasks); + const errors: unknown[] = []; + const maxDrainRounds = 100; + let round = 0; + + while (asyncTasks.length > 0) { + if (round >= maxDrainRounds) { + asyncTaskAdmissionOpen = false; + const overflowTasks = asyncTasks.splice(0, asyncTasks.length); + for (const task of overflowTasks) { + asyncTaskControllers.get(task)?.abort(); + } + const overflowResults = await Promise.allSettled(overflowTasks); + for (let index = 0; index < overflowResults.length; index += 1) { + asyncTaskControllers.delete(overflowTasks[index]); + const result = overflowResults[index]; + if (result.status === "rejected") { + errors.push(result.reason); + } + } + errors.push(new Error(`Async task drain exceeded ${maxDrainRounds} rounds`)); + break; + } + round += 1; + + const tasks = asyncTasks.splice(0, asyncTasks.length); + const results = await Promise.allSettled(tasks); + for (let index = 0; index < results.length; index += 1) { + asyncTaskControllers.delete(tasks[index]); + const result = results[index]; + if (result.status === "rejected") { + errors.push(result.reason); + } + } + } + + if (errors.length > 0) { + throw new AggregateError(errors, "Async task drain failed"); + } } +describe("drainAsyncTasks", () => { + it("waits for a tail task registered while draining the primary task", async () => { + let markTailStarted: () => void = () => {}; + let releaseTail: () => void = () => {}; + const tailStarted = new Promise((resolve) => { + markTailStarted = resolve; + }); + const tailCompleted = vi.fn(); + + AsyncTaskManager.register("primary", async () => { + await Promise.resolve(); + AsyncTaskManager.register("tail", async () => { + markTailStarted(); + await new Promise((resolve) => { + releaseTail = resolve; + }); + tailCompleted(); + }); + }); + + const drainPromise = drainAsyncTasks(); + await tailStarted; + + try { + const outcome = await Promise.race([ + drainPromise.then(() => "drained" as const), + new Promise<"pending">((resolve) => { + setTimeout(() => resolve("pending"), 0); + }), + ]); + + expect(outcome).toBe("pending"); + } finally { + releaseTail(); + } + + await drainPromise; + expect(tailCompleted).toHaveBeenCalledTimes(1); + }); + + it("waits for sibling tail work before reporting tail rejections", async () => { + const tailError = new Error("tail task failed"); + let markPendingTailStarted: () => void = () => {}; + let releasePendingTail: () => void = () => {}; + const pendingTailStarted = new Promise((resolve) => { + markPendingTailStarted = resolve; + }); + + AsyncTaskManager.register("primary", async () => { + await Promise.resolve(); + AsyncTaskManager.register("rejecting-tail", async () => { + throw tailError; + }); + void asyncTasks.at(-1)?.catch(() => {}); + AsyncTaskManager.register("pending-tail", async () => { + markPendingTailStarted(); + await new Promise((resolve) => { + releasePendingTail = resolve; + }); + }); + }); + + const drainPromise = drainAsyncTasks(); + await pendingTailStarted; + + try { + const earlyOutcome = await Promise.race([ + drainPromise.then( + () => "resolved" as const, + () => "rejected" as const + ), + new Promise<"pending">((resolve) => { + setTimeout(() => resolve("pending"), 0); + }), + ]); + + expect(earlyOutcome).toBe("pending"); + } finally { + releasePendingTail(); + } + + const rejection = await drainPromise.then( + () => undefined, + (error: unknown) => error + ); + expect(rejection).toBeInstanceOf(AggregateError); + expect((rejection as AggregateError).errors).toEqual([tailError]); + }); + + it("closes admission and observes overflow work when the drain guard trips", async () => { + const overflowError = new Error("overflow task aborted"); + const blockedTailStarted = vi.fn(); + let overflowController: AbortController | undefined; + + const registerGeneration = (generation: number): void => { + const controller = AsyncTaskManager.register(`generation-${generation}`, async (signal) => { + await new Promise((resolve) => { + setTimeout(resolve, 0); + }); + + if (generation <= 100) { + registerGeneration(generation + 1); + return; + } + + await new Promise((_resolve, reject) => { + const rejectOnAbort = () => { + AsyncTaskManager.register("blocked-overflow-tail", async () => { + blockedTailStarted(); + }); + reject(overflowError); + }; + + if (signal.aborted) { + rejectOnAbort(); + return; + } + signal.addEventListener("abort", rejectOnAbort, { once: true }); + }); + }); + + if (generation === 101) { + overflowController = controller; + } + }; + + registerGeneration(1); + const rejection = await drainAsyncTasks().then( + () => undefined, + (error: unknown) => error + ); + + try { + expect(rejection).toBeInstanceOf(AggregateError); + expect((rejection as AggregateError).errors).toContain(overflowError); + expect((rejection as AggregateError).errors).toEqual( + expect.arrayContaining([ + expect.objectContaining({ message: "Async task drain exceeded 100 rounds" }), + ]) + ); + expect(blockedTailStarted).not.toHaveBeenCalled(); + expect(asyncTasks).toHaveLength(0); + } finally { + overflowController?.abort(); + await Promise.allSettled(asyncTasks.splice(0, asyncTasks.length)); + } + }); +}); + function captureRateLimitCosts(): number[] { const rateLimitCosts: number[] = []; vi.mocked(RateLimitService.trackCost).mockImplementation( @@ -367,7 +586,7 @@ async function runScenario({ return null; }); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); @@ -494,7 +713,7 @@ describe("Billing model source - Redis session cost vs DB cost", () => { it("nested pricing: gpt-5.5 alias model should bill from pricing.openai when provider is chatgpt", async () => { vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("redirected")); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); @@ -556,7 +775,7 @@ describe("Billing model source - Redis session cost vs DB cost", () => { it("codex fast: requested mode ignores actual priority when request tier is default", async () => { vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("redirected")); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); @@ -625,7 +844,7 @@ describe("Billing model source - Redis session cost vs DB cost", () => { it("codex fast: falls back to requested priority pricing when response omits service_tier", async () => { vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("redirected")); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); @@ -681,7 +900,7 @@ describe("Billing model source - Redis session cost vs DB cost", () => { it("codex fast: uses long-context priority pricing when request is priority and response omits service_tier", async () => { vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("redirected")); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); @@ -751,7 +970,7 @@ describe("Billing model source - Redis session cost vs DB cost", () => { it("codex fast: requested mode keeps priority pricing even when actual tier is downgraded", async () => { vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("redirected")); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); @@ -810,7 +1029,7 @@ describe("Billing model source - Redis session cost vs DB cost", () => { it("codex fast: actual mode uses priority pricing when response reports service_tier=priority", async () => { vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("redirected", "actual")); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); @@ -869,7 +1088,7 @@ describe("Billing model source - Redis session cost vs DB cost", () => { it("codex fast: actual mode does not use priority pricing when response explicitly reports non-priority tier", async () => { vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("redirected", "actual")); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); @@ -928,7 +1147,7 @@ describe("Billing model source - Redis session cost vs DB cost", () => { it("codex fast: actual mode falls back to requested priority pricing when response omits service_tier", async () => { vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("redirected", "actual")); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); @@ -987,7 +1206,7 @@ describe("Billing model source - Redis session cost vs DB cost", () => { await getCachedSystemSettings(); vi.mocked(getSystemSettings).mockRejectedValueOnce(new Error("db down")); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); @@ -1075,7 +1294,7 @@ describe("模型重定向后的图片按次计费", () => { return null; }); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); @@ -1185,7 +1404,7 @@ describe("模型重定向后的图片按次计费", () => { return makePriceRecord(modelName, { input_cost_per_request: 0 }, "manual"); }); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); @@ -1229,7 +1448,7 @@ describe("模型重定向后的图片按次计费", () => { throw new Error("pricing db unavailable"); }); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); @@ -1277,7 +1496,7 @@ describe("模型重定向后的图片按次计费", () => { return makePriceRecord(modelName, { input_cost_per_request: 0.01 }, "manual"); }); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); @@ -1327,7 +1546,7 @@ describe("模型重定向后的图片按次计费", () => { }); vi.mocked(updateMessageRequestCostWithBreakdown).mockResolvedValue(undefined); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackCost).mockResolvedValue(undefined); let sessionUsagePayload: Record | undefined; @@ -1390,7 +1609,7 @@ describe("价格表缺失/查询失败:不计费放行", () => { }); } - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); diff --git a/tests/integration/db-pool-isolation-postgres.test.ts b/tests/integration/db-pool-isolation-postgres.test.ts new file mode 100644 index 000000000..197046577 --- /dev/null +++ b/tests/integration/db-pool-isolation-postgres.test.ts @@ -0,0 +1,135 @@ +import { performance } from "node:perf_hooks"; +import { sql } from "drizzle-orm"; +import { afterAll, beforeAll, describe, expect, test, vi } from "vitest"; + +const HAS_DB = Boolean(process.env.DSN || process.env.DATABASE_URL); +const run = describe.skipIf(!HAS_DB); + +run("PostgreSQL pool isolation integration", () => { + let dbModule: typeof import("@/drizzle/db"); + let previousDsn: string | undefined; + let previousPoolMax: string | undefined; + + async function rows(query: ReturnType): Promise { + return Array.from(await dbModule.getDb().execute(query)) as T[]; + } + + async function waitForActiveDataConnections(expected: number): Promise { + const deadline = performance.now() + 2_000; + while (performance.now() < deadline) { + const [row] = await rows<{ active: number }>(sql` + SELECT COUNT(*)::int AS active + FROM pg_stat_activity + WHERE datname = current_database() + AND application_name = 'claude-code-hub:data' + AND state = 'active' + `); + if (Number(row?.active ?? 0) >= expected) return; + await new Promise((resolve) => setTimeout(resolve, 10)); + } + throw new Error(`Timed out waiting for ${expected} active data connections`); + } + + async function startDataSleep(seconds: number): Promise { + return dbModule.runWithDataDbScope(async () => { + await dbModule.getDb().execute(sql`SELECT pg_sleep(${seconds})`); + }); + } + + beforeAll(async () => { + previousDsn = process.env.DSN; + previousPoolMax = process.env.DB_POOL_MAX; + if (!process.env.DSN && process.env.DATABASE_URL) { + process.env.DSN = process.env.DATABASE_URL; + } + process.env.DB_POOL_MAX = "6"; + vi.resetModules(); + dbModule = await import("@/drizzle/db"); + }); + + afterAll(async () => { + try { + await dbModule.closeDbPools(); + } finally { + if (previousPoolMax === undefined) { + delete process.env.DB_POOL_MAX; + } else { + process.env.DB_POOL_MAX = previousPoolMax; + } + if (previousDsn === undefined) { + delete process.env.DSN; + } else { + process.env.DSN = previousDsn; + } + } + }); + + test("creates observable data, control, and writer lanes", async () => { + await dbModule.getDb().execute(sql`SELECT 1`); + await dbModule.runWithDataDbScope(() => dbModule.getDb().execute(sql`SELECT 1`)); + await dbModule.getMessageWriterDb().execute(sql`SELECT 1`); + + const activity = await rows<{ applicationName: string; connections: number }>(sql` + SELECT application_name AS "applicationName", COUNT(*)::int AS connections + FROM pg_stat_activity + WHERE datname = current_database() + AND application_name IN ( + 'claude-code-hub:data', + 'claude-code-hub:control', + 'claude-code-hub:writer' + ) + GROUP BY application_name + `); + + expect(new Set(activity.map(({ applicationName }) => applicationName))).toEqual( + new Set(["claude-code-hub:data", "claude-code-hub:control", "claude-code-hub:writer"]) + ); + }); + + test("keeps control and writer queries responsive while every data connection is busy", async () => { + const dataSleeps = Array.from({ length: 4 }, () => startDataSleep(0.5)); + await waitForActiveDataConnections(4); + + const startedAt = performance.now(); + await Promise.all([ + dbModule.getDb().execute(sql`SELECT 1`), + dbModule.getMessageWriterDb().execute(sql`SELECT 1`), + ]); + const isolatedLatencyMs = performance.now() - startedAt; + + expect(isolatedLatencyMs).toBeLessThan(300); + await Promise.all(dataSleeps); + }); + + test("rejects the 33rd outstanding data query before it joins the postgres.js queue", async () => { + const pending = Array.from({ length: 32 }, () => startDataSleep(0.2)); + + const rejectedAt = performance.now(); + let admissionError: unknown; + try { + await startDataSleep(0.2); + } catch (error) { + admissionError = error; + } + const rejectionLatencyMs = performance.now() - rejectedAt; + + expect(admissionError).toMatchObject({ + cause: { + name: "DbPoolAdmissionError", + code: "DB_POOL_ADMISSION_EXCEEDED", + pool: "data", + maxOutstanding: 32, + }, + }); + expect(rejectionLatencyMs).toBeLessThan(50); + + const controlAndWriterStartedAt = performance.now(); + await Promise.all([ + dbModule.getDb().execute(sql`SELECT 1`), + dbModule.getMessageWriterDb().execute(sql`SELECT 1`), + ]); + expect(performance.now() - controlAndWriterStartedAt).toBeLessThan(300); + + await Promise.all(pending); + }); +}); diff --git a/tests/integration/db-pool-slow-close-postgres.test.ts b/tests/integration/db-pool-slow-close-postgres.test.ts new file mode 100644 index 000000000..5d62cc117 --- /dev/null +++ b/tests/integration/db-pool-slow-close-postgres.test.ts @@ -0,0 +1,282 @@ +import { performance } from "node:perf_hooks"; +import { sql } from "drizzle-orm"; +import postgres from "postgres"; +import { describe, expect, test, vi } from "vitest"; + +const HAS_DB = Boolean(process.env.DSN || process.env.DATABASE_URL); +const run = describe.skipIf(!HAS_DB); + +const CLEANUP_BOUNDARY_MODULES = [ + "@/lib/cache/session-cache", + "@/lib/provider-endpoints/probe-scheduler", + "@/lib/public-status/scheduler", + "@/lib/provider-endpoints/probe-log-cleanup", + "@/lib/async-task-manager", + "@/repository/message-write-buffer", + "@/lib/langfuse", + "@/lib/redis", +] as const; + +interface ActivityRow { + state: string; + waitEventType: string | null; + waitEvent: string | null; +} + +interface ConnectionRow { + applicationName: string; + connections: number; +} + +function delay(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +function observeSettlement(promise: Promise): () => boolean { + let settled = false; + void promise.then( + () => { + settled = true; + }, + () => { + settled = true; + } + ); + return () => settled; +} + +async function waitFor( + probe: () => T | undefined | Promise, + description: string, + timeoutMs = 3_000 +): Promise { + const deadline = performance.now() + timeoutMs; + while (performance.now() < deadline) { + const value = await probe(); + if (value !== undefined) return value; + await delay(10); + } + throw new Error(`Timed out waiting for ${description}`); +} + +run.sequential("PostgreSQL slow pool close integration", () => { + test("keeps cleanup attached to an advisory-lock-blocked data query until every pool closes", async () => { + const originalEnv = { + DSN: process.env.DSN, + DB_POOL_MAX: process.env.DB_POOL_MAX, + DB_LOCK_TIMEOUT_MS: process.env.DB_LOCK_TIMEOUT_MS, + DB_STATEMENT_TIMEOUT_MS: process.env.DB_STATEMENT_TIMEOUT_MS, + }; + const dsn = process.env.DSN ?? process.env.DATABASE_URL; + if (!dsn) throw new Error("DSN or DATABASE_URL is required"); + + process.env.DSN = dsn; + process.env.DB_POOL_MAX = "3"; + process.env.DB_LOCK_TIMEOUT_MS = "4000"; + process.env.DB_STATEMENT_TIMEOUT_MS = "10000"; + + // tests/setup.ts 可能已通过旧 module instance 建立 control pool。先用 public close + // 清空该 harness 连接,再 reset 出本测试独占的三 lane,避免失去旧 pool 的 close handle。 + const harnessDbModule = await import("@/drizzle/db"); + await harnessDbModule.closeDbPools(); + vi.resetModules(); + + const holder = postgres(dsn, { + max: 1, + connect_timeout: 5, + connection: { application_name: "cch-pool-close-test:holder" }, + }); + const observer = postgres(dsn, { + max: 1, + connect_timeout: 5, + connection: { application_name: "cch-pool-close-test:observer" }, + }); + + const lockNamespace = 0x434348; + const lockId = Math.floor(Math.random() * 1_000_000_000); + let lockHeld = false; + let dbModule: typeof import("@/drizzle/db") | undefined; + let activeQuery: Promise | undefined; + let closePromise: Promise | undefined; + let cleanup: Promise | undefined; + + try { + const [{ databaseName }] = await observer<{ databaseName: string }[]>` + SELECT current_database() AS "databaseName" + `; + expect(databaseName).toMatch(/test/i); + + await holder`SELECT pg_advisory_lock(${lockNamespace}, ${lockId})`; + lockHeld = true; + + vi.doMock("@/lib/cache/session-cache", () => ({ stopCacheCleanup: vi.fn() })); + vi.doMock("@/lib/provider-endpoints/probe-scheduler", () => ({ + stopEndpointProbeScheduler: vi.fn(), + })); + vi.doMock("@/lib/public-status/scheduler", () => ({ + stopPublicStatusRebuildScheduler: vi.fn(async () => {}), + })); + vi.doMock("@/lib/provider-endpoints/probe-log-cleanup", () => ({ + stopEndpointProbeLogCleanup: vi.fn(), + })); + vi.doMock("@/lib/async-task-manager", () => ({ + shutdownAllAsyncTasks: vi.fn(async () => {}), + })); + vi.doMock("@/repository/message-write-buffer", () => ({ + stopMessageRequestWriteBuffer: vi.fn(async () => {}), + })); + vi.doMock("@/lib/langfuse", () => ({ shutdownLangfuse: vi.fn(async () => {}) })); + vi.doMock("@/lib/redis", () => ({ closeRedis: vi.fn(async () => {}) })); + + dbModule = await import("@/drizzle/db"); + const { runApplicationCleanup } = await import("@/lib/lifecycle/shutdown"); + + await Promise.all([ + dbModule.getDb().execute(sql`SELECT 1`), + dbModule.runWithDataDbScope(() => dbModule?.getDb().execute(sql`SELECT 1`)), + dbModule.getMessageWriterDb().execute(sql`SELECT 1`), + ]); + + const warmedConnections = await observer` + SELECT application_name AS "applicationName", COUNT(*)::int AS connections + FROM pg_stat_activity + WHERE datname = current_database() + AND application_name IN ( + 'claude-code-hub:data', + 'claude-code-hub:control', + 'claude-code-hub:writer' + ) + GROUP BY application_name + `; + expect( + new Set( + warmedConnections + .filter(({ connections }) => connections > 0) + .map(({ applicationName }) => applicationName) + ) + ).toEqual( + new Set(["claude-code-hub:data", "claude-code-hub:control", "claude-code-hub:writer"]) + ); + + activeQuery = dbModule + .runWithDataDbScope(() => + dbModule?.getDb().execute(sql`SELECT pg_advisory_xact_lock(${lockNamespace}, ${lockId})`) + ) + .then(() => undefined); + const isActiveQuerySettled = observeSettlement(activeQuery); + + const blockedActivity = await waitFor(async () => { + const [row] = await observer` + SELECT + state, + wait_event_type AS "waitEventType", + wait_event AS "waitEvent" + FROM pg_stat_activity + WHERE datname = current_database() + AND application_name = 'claude-code-hub:data' + AND query LIKE '%pg_advisory_xact_lock%' + ORDER BY query_start DESC + LIMIT 1 + `; + if ( + row?.state === "active" && + row.waitEventType === "Lock" && + row.waitEvent === "advisory" + ) { + return row; + } + return undefined; + }, "data query to wait on an advisory lock"); + expect(blockedActivity).toEqual({ + state: "active", + waitEventType: "Lock", + waitEvent: "advisory", + }); + + const perStepTimeoutMs = 50; + const totalTimeoutMs = 100; + cleanup = runApplicationCleanup("integration-test", { + perStepTimeoutMs, + totalTimeoutMs, + }); + const isCleanupSettled = observeSettlement(cleanup); + + await waitFor(() => { + try { + dbModule?.getDb(); + return undefined; + } catch (error) { + if (error instanceof Error && error.message === "Database pools are closing") { + return true; + } + throw error; + } + }, "database pools to enter closing state"); + + const firstClose = dbModule.closeDbPools(); + const secondClose = dbModule.closeDbPools(); + expect(firstClose).toBe(secondClose); + closePromise = firstClose; + const isCloseSettled = observeSettlement(closePromise); + + expect(() => dbModule?.getDb()).toThrow("Database pools are closing"); + expect(() => dbModule?.runWithDataDbScope(() => dbModule?.getDb())).toThrow( + "Database pools are closing" + ); + expect(() => dbModule?.getMessageWriterDb()).toThrow("Database pools are closing"); + + await delay(totalTimeoutMs + perStepTimeoutMs); + + expect(isActiveQuerySettled()).toBe(false); + expect(isCloseSettled()).toBe(false); + expect(isCleanupSettled()).toBe(false); + + const [{ unlocked }] = await holder<{ unlocked: boolean }[]>` + SELECT pg_advisory_unlock(${lockNamespace}, ${lockId}) AS unlocked + `; + expect(unlocked).toBe(true); + lockHeld = false; + + await Promise.all([activeQuery, closePromise, cleanup]); + + await waitFor(async () => { + const [{ connections }] = await observer<{ connections: number }[]>` + SELECT COUNT(*)::int AS connections + FROM pg_stat_activity + WHERE datname = current_database() + AND application_name IN ( + 'claude-code-hub:data', + 'claude-code-hub:control', + 'claude-code-hub:writer' + ) + `; + return connections === 0 ? true : undefined; + }, "all production database connections to close"); + } finally { + if (lockHeld) { + await holder`SELECT pg_advisory_unlock(${lockNamespace}, ${lockId})`; + } + + const pending: Promise[] = []; + if (activeQuery) pending.push(activeQuery); + if (dbModule) pending.push(dbModule.closeDbPools()); + if (closePromise) pending.push(closePromise); + if (cleanup) pending.push(cleanup); + await Promise.allSettled(pending); + await Promise.allSettled([holder.end({ timeout: 1 }), observer.end({ timeout: 1 })]); + + for (const moduleId of CLEANUP_BOUNDARY_MODULES) { + vi.doUnmock(moduleId); + } + vi.resetModules(); + + for (const [key, value] of Object.entries(originalEnv)) { + if (value === undefined) { + delete process.env[key]; + } else { + process.env[key] = value; + } + } + } + }); +}); diff --git a/tests/integration/lease-settlement-redis.test.ts b/tests/integration/lease-settlement-redis.test.ts new file mode 100644 index 000000000..ecf4a299a --- /dev/null +++ b/tests/integration/lease-settlement-redis.test.ts @@ -0,0 +1,243 @@ +import { afterAll, afterEach, beforeAll, describe, expect, test } from "vitest"; +import { + buildLeaseKey, + type BudgetLease, + type LeaseEntityTypeType, + type LeaseWindowType, + serializeLease, +} from "@/lib/rate-limit/lease"; +import { LeaseService, type SettleLeaseBudgetsParams } from "@/lib/rate-limit/lease-service"; +import type { DailyResetMode } from "@/lib/rate-limit/time-utils"; +import { closeRedis, getRedisClient } from "@/lib/redis/client"; + +const HAS_REDIS = Boolean(process.env.REDIS_URL); +const run = describe.skipIf(!HAS_REDIS); +const TEST_PREFIX = `it-lease-settlement-${Date.now()}-${Math.random().toString(16).slice(2)}`; + +interface SettlementTarget { + entityType: LeaseEntityTypeType; + entityId: number; + window: LeaseWindowType; + resetMode: DailyResetMode; + key: string; +} + +function buildParams(tag: string): SettleLeaseBudgetsParams { + const baseId = + 800_000_000 + (Date.now() % 10_000_000) + Math.floor(Math.random() * 10_000) * 10 + tag.length; + + return { + requestId: `${TEST_PREFIX}:${tag}`, + cost: 1.25, + entities: { + key: { + id: baseId, + resetModes: { "5h": "rolling", daily: "fixed" }, + }, + user: { + id: baseId + 1, + resetModes: { "5h": "fixed", daily: "rolling" }, + }, + provider: { + id: baseId + 2, + resetModes: { "5h": "rolling", daily: "fixed" }, + }, + }, + }; +} + +function buildTargets(params: SettleLeaseBudgetsParams): SettlementTarget[] { + const targets: SettlementTarget[] = []; + const entityTypes = ["key", "user", "provider"] as const; + const windows = ["5h", "daily", "weekly", "monthly"] as const; + + for (const entityType of entityTypes) { + const entity = params.entities[entityType]; + for (const window of windows) { + const resetMode = + window === "5h" || window === "daily" + ? (entity.resetModes?.[window] ?? (window === "5h" ? "rolling" : "fixed")) + : "fixed"; + targets.push({ + entityType, + entityId: entity.id, + window, + resetMode, + key: buildLeaseKey(entityType, entity.id, window, resetMode), + }); + } + } + + return targets; +} + +run("lease settlement Lua integration", () => { + const touchedKeys = new Set(); + let redis: NonNullable>; + let previousEnableRateLimit: string | undefined; + + async function waitForRedisReady() { + const client = getRedisClient({ allowWhenRateLimitDisabled: true }); + if (!client) { + throw new Error("Redis client unavailable for integration test"); + } + + if (client.status !== "ready") { + await new Promise((resolve, reject) => { + const timeout = setTimeout(() => reject(new Error("Redis ready timeout")), 5_000); + client.once("ready", () => { + clearTimeout(timeout); + resolve(); + }); + }); + } + + if (client.status !== "ready") { + throw new Error(`Redis not ready: ${client.status}`); + } + return client; + } + + function rememberParams(params: SettleLeaseBudgetsParams): SettlementTarget[] { + const targets = buildTargets(params); + for (const target of targets) touchedKeys.add(target.key); + touchedKeys.add(`lease:settlement:${String(params.requestId)}`); + return targets; + } + + function makeLease(target: SettlementTarget, remainingBudget: number): BudgetLease { + return { + entityType: target.entityType, + entityId: target.entityId, + window: target.window, + resetMode: target.resetMode, + resetTime: "00:00", + snapshotAtMs: Date.now(), + currentUsage: 10, + limitAmount: 200, + remainingBudget, + ttlSeconds: 120, + }; + } + + async function seedValidLeases(targets: SettlementTarget[], remainingBudget = 100) { + const pipeline = redis.pipeline(); + for (let index = 0; index < targets.length; index += 1) { + pipeline.set( + targets[index].key, + serializeLease(makeLease(targets[index], remainingBudget + index)), + "EX", + 120 + ); + } + const results = await pipeline.exec(); + expect(results?.every(([error]) => error === null)).toBe(true); + } + + beforeAll(async () => { + previousEnableRateLimit = process.env.ENABLE_RATE_LIMIT; + process.env.ENABLE_RATE_LIMIT = "true"; + redis = await waitForRedisReady(); + await expect(redis.ping()).resolves.toBe("PONG"); + }); + + afterEach(async () => { + if (touchedKeys.size > 0) { + await redis.del(...touchedKeys); + touchedKeys.clear(); + } + }); + + afterAll(async () => { + if (previousEnableRateLimit === undefined) { + delete process.env.ENABLE_RATE_LIMIT; + } else { + process.env.ENABLE_RATE_LIMIT = previousEnableRateLimit; + } + await closeRedis(); + }); + + test("settles all twelve leases once and replays the marker without a second decrement", async () => { + const params = buildParams("replay"); + const targets = rememberParams(params); + await seedValidLeases(targets); + + const first = await LeaseService.settleLeaseBudgets(params); + expect(first.status).toBe("settled"); + expect(first.settlements).toHaveLength(12); + expect(first.settlements.every(({ status }) => status === "decremented")).toBe(true); + + for (let index = 0; index < targets.length; index += 1) { + const raw = await redis.get(targets[index].key); + expect(raw).not.toBeNull(); + const lease = JSON.parse(raw!) as BudgetLease; + expect(lease.remainingBudget).toBeCloseTo(100 + index - params.cost, 10); + expect(await redis.ttl(targets[index].key)).toBeGreaterThan(0); + } + + const markerKey = `lease:settlement:${String(params.requestId)}`; + const marker = await redis.get(markerKey); + expect(marker).not.toBeNull(); + expect(JSON.parse(marker!)).toHaveLength(12); + const markerTtl = await redis.ttl(markerKey); + expect(markerTtl).toBeGreaterThan(0); + expect(markerTtl).toBeLessThanOrEqual(5 * 60); + + const duplicate = await LeaseService.settleLeaseBudgets(params); + expect(duplicate.status).toBe("duplicate"); + expect(duplicate.settlements).toEqual(first.settlements); + + for (let index = 0; index < targets.length; index += 1) { + const lease = JSON.parse((await redis.get(targets[index].key))!) as BudgetLease; + expect(lease.remainingBudget).toBeCloseTo(100 + index - params.cost, 10); + } + }); + + test("malformed JSON, missing TTL, insufficient budget, and WRONGTYPE are isolated", async () => { + const params = buildParams("faults"); + const targets = rememberParams(params); + await seedValidLeases(targets); + + const malformed = targets[0]; + const noTtl = targets[1]; + const insufficient = targets[2]; + const wrongType = targets[3]; + + await redis.set(malformed.key, "{malformed", "EX", 120); + await redis.set(noTtl.key, serializeLease(makeLease(noTtl, 50))); + await redis.set(insufficient.key, serializeLease(makeLease(insufficient, 0.5)), "EX", 120); + await redis.del(wrongType.key); + await redis.lpush(wrongType.key, "not-a-string-lease"); + await redis.expire(wrongType.key, 120); + + const result = await LeaseService.settleLeaseBudgets(params); + expect(result.status).toBe("settled"); + expect(result.settlements).toHaveLength(12); + + const byTarget = new Map( + result.settlements.map((settlement) => [ + `${settlement.entityType}:${settlement.window}`, + settlement, + ]) + ); + expect(byTarget.get("key:5h")?.status).toBe("missing"); + expect(byTarget.get("key:daily")?.status).toBe("missing"); + expect(byTarget.get("key:weekly")?.status).toBe("insufficient"); + expect(byTarget.get("key:monthly")?.status).toBe("missing"); + expect(result.settlements.filter(({ status }) => status === "decremented")).toHaveLength(8); + + await expect(redis.get(malformed.key)).resolves.toBe("{malformed"); + expect(JSON.parse((await redis.get(noTtl.key))!) as BudgetLease).toMatchObject({ + remainingBudget: 50, + }); + await expect(redis.ttl(noTtl.key)).resolves.toBe(-1); + expect(JSON.parse((await redis.get(insufficient.key))!) as BudgetLease).toMatchObject({ + remainingBudget: 0, + }); + await expect(redis.type(wrongType.key)).resolves.toBe("list"); + + const markerKey = `lease:settlement:${String(params.requestId)}`; + expect(JSON.parse((await redis.get(markerKey))!)).toHaveLength(12); + expect(await redis.ttl(markerKey)).toBeGreaterThan(0); + }); +}); diff --git a/tests/integration/message-write-buffer-recovery-postgres.test.ts b/tests/integration/message-write-buffer-recovery-postgres.test.ts new file mode 100644 index 000000000..2686fd379 --- /dev/null +++ b/tests/integration/message-write-buffer-recovery-postgres.test.ts @@ -0,0 +1,291 @@ +import { randomUUID } from "node:crypto"; +import { performance } from "node:perf_hooks"; +import { like } from "drizzle-orm"; +import postgres from "postgres"; +import { afterAll, beforeAll, describe, expect, test, vi } from "vitest"; +import { messageRequest, usageLedger } from "@/drizzle/schema"; + +const ENV_KEYS = [ + "DSN", + "DB_POOL_MAX", + "DB_LOCK_TIMEOUT_MS", + "DB_STATEMENT_TIMEOUT_MS", + "MESSAGE_REQUEST_WRITE_MODE", + "MESSAGE_REQUEST_ASYNC_FLUSH_INTERVAL_MS", + "MESSAGE_REQUEST_ASYNC_BATCH_SIZE", + "MESSAGE_REQUEST_ASYNC_MAX_PENDING", +] as const; +const originalEnv = new Map(ENV_KEYS.map((key) => [key, process.env[key]] as const)); +const dsn = process.env.DSN ?? process.env.DATABASE_URL; + +if (dsn) { + process.env.DSN = dsn; + process.env.DB_POOL_MAX = "4"; + process.env.DB_LOCK_TIMEOUT_MS = "100"; + process.env.DB_STATEMENT_TIMEOUT_MS = "5000"; + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + process.env.MESSAGE_REQUEST_ASYNC_FLUSH_INTERVAL_MS = "50"; + process.env.MESSAGE_REQUEST_ASYNC_BATCH_SIZE = "200"; + process.env.MESSAGE_REQUEST_ASYNC_MAX_PENDING = "5000"; +} +vi.resetModules(); + +const run = describe.skipIf(!dsn); +const KEY_PREFIX = `it-message-buffer-recovery-${randomUUID()}`; + +run.sequential("message write buffer PostgreSQL recovery", () => { + let dbModule: typeof import("@/drizzle/db"); + let messageRepository: typeof import("@/repository/message"); + let writeBuffer: typeof import("@/repository/message-write-buffer"); + let lockClient: ReturnType; + + async function createRequest(tag: string, costUsd = "0.250000000000000"): Promise { + const request = await messageRepository.createMessageRequest({ + provider_id: 910_000_001, + user_id: 920_000_001, + key: `${KEY_PREFIX}-${tag}`, + model: "integration-model", + original_model: "integration-model", + endpoint: "/v1/messages", + cost_usd: costUsd, + }); + return request.id; + } + + async function cleanupTestRows(): Promise { + const keyPattern = `${KEY_PREFIX}%`; + await dbModule.getDb().delete(messageRequest).where(like(messageRequest.key, keyPattern)); + await dbModule.getDb().delete(usageLedger).where(like(usageLedger.key, keyPattern)); + } + + async function expectRequest( + id: number, + expected: Readonly> + ): Promise { + await expect(messageRepository.findMessageRequestById(id)).resolves.toMatchObject(expected); + } + + function restoreEnvironment(): void { + for (const [key, value] of originalEnv) { + if (value === undefined) delete process.env[key]; + else process.env[key] = value; + } + } + + beforeAll(async () => { + if (!dsn) throw new TypeError("DSN or DATABASE_URL is required"); + + const harnessDb = await import("@/drizzle/db"); + await harnessDb.closeDbPools(); + vi.resetModules(); + + [dbModule, messageRepository, writeBuffer] = await Promise.all([ + import("@/drizzle/db"), + import("@/repository/message"), + import("@/repository/message-write-buffer"), + ]); + lockClient = postgres(dsn, { + max: 1, + connect_timeout: 5, + connection: { application_name: "cch-message-buffer-recovery:lock" }, + }); + + const [{ databaseName }] = await lockClient<{ databaseName: string }[]>` + SELECT current_database() AS "databaseName" + `; + expect(databaseName).toMatch(/test/i); + await cleanupTestRows(); + }); + + afterAll(async () => { + const failures: unknown[] = []; + for (const cleanup of [ + () => writeBuffer.stopMessageRequestWriteBuffer(), + cleanupTestRows, + () => lockClient.end({ timeout: 5 }), + () => dbModule.closeDbPools(), + ]) { + try { + await cleanup(); + } catch (error) { + failures.push(error instanceof Error ? error : new Error(String(error))); + } + } + restoreEnvironment(); + if (failures.length > 0) throw new AggregateError(failures, "integration cleanup failed"); + }); + + test("settles a mixed durable batch only after PostgreSQL commits", async () => { + const ordinaryId = await createRequest("mixed-ordinary"); + const durableId = await createRequest("mixed-durable"); + + const operations = await lockClient.begin(async (transaction) => { + await transaction`SELECT id FROM message_request WHERE id = ${durableId} FOR UPDATE`; + writeBuffer.enqueueMessageRequestUpdate(ordinaryId, { durationMs: 450 }); + const durable = writeBuffer.enqueueMessageRequestUpdateDurably(durableId, { + durationMs: 900, + statusCode: 200, + costUsd: "0.375000000000000", + }); + const flush = writeBuffer.flushMessageRequestWriteBuffer(); + + let settled = false; + void durable.then( + () => { + settled = true; + }, + () => { + settled = true; + } + ); + await new Promise((resolve) => setImmediate(resolve)); + expect(settled).toBe(false); + await expectRequest(ordinaryId, { durationMs: null }); + return { durable, flush }; + }); + + await Promise.all([operations.flush, operations.durable]); + await expectRequest(ordinaryId, { durationMs: 450 }); + await expectRequest(durableId, { + durationMs: 900, + statusCode: 200, + costUsd: "0.375000000000000", + }); + await messageRepository.updateMessageRequestDetailsIfUnfinalized(durableId, { + durationMs: 1_800, + statusCode: 503, + }); + await expectRequest(durableId, { durationMs: 900, statusCode: 200 }); + }); + + test("keeps fallback ownership when a timed-out primary commits late", async () => { + const requestId = await createRequest("fallback-first"); + let committedReceipts = 0; + + const { flush } = await lockClient.begin(async (transaction) => { + await transaction`SELECT id FROM message_request WHERE id = ${requestId} FOR UPDATE`; + const primary = writeBuffer.enqueueMessageRequestUpdateDurably( + requestId, + { + durationMs: 300, + statusCode: 200, + costUsd: "0.100000000000000", + errorMessage: "retired-primary", + }, + { + timeoutMs: 20, + onCommitted: () => { + committedReceipts++; + }, + } + ); + const flush = writeBuffer.flushMessageRequestWriteBuffer(); + + await expect(primary).rejects.toThrow("durable message_request acknowledgement timed out"); + const fallbackReceipts = await transaction<{ id: number }[]>` + UPDATE message_request + SET duration_ms = 2_400, + status_code = 504, + cost_usd = 0.625, + error_message = 'fallback-owner', + updated_at = NOW() + WHERE id = ${requestId} AND status_code IS NULL + RETURNING id + `; + expect(fallbackReceipts).toEqual([{ id: requestId }]); + return { flush }; + }); + + await flush; + expect(committedReceipts).toBe(0); + await expectRequest(requestId, { + durationMs: 2_400, + statusCode: 504, + costUsd: "0.625000000000000", + errorMessage: "fallback-owner", + }); + }); + + test("reinserted generation excludes a retired pending patch", async () => { + const requestId = await createRequest("generation-reinsert"); + const retired = writeBuffer.enqueueMessageRequestUpdateDurably( + requestId, + { statusCode: 500, costUsd: "0.100000000000000", errorMessage: "retired-pending" }, + { timeoutMs: 20 } + ); + + await expect(retired).rejects.toThrow("durable message_request acknowledgement timed out"); + const current = writeBuffer.enqueueMessageRequestUpdateDurably(requestId, { + durationMs: 700, + statusCode: 201, + costUsd: "0.450000000000000", + errorMessage: "current-generation", + }); + await Promise.all([writeBuffer.flushMessageRequestWriteBuffer(), current]); + + await expectRequest(requestId, { + durationMs: 700, + statusCode: 201, + costUsd: "0.450000000000000", + errorMessage: "current-generation", + }); + }); + + test("bounds a saturated 5,000-entry queue while retaining terminal priority", async () => { + const lockedId = await createRequest("saturation-lock"); + const terminalId = await createRequest("saturation-terminal"); + const evictedId = await createRequest("saturation-evicted"); + const retainedId = await createRequest("saturation-retained"); + + const { flush } = await lockClient.begin(async (transaction) => { + await transaction`SELECT id FROM message_request WHERE id = ${lockedId} FOR UPDATE`; + writeBuffer.enqueueMessageRequestUpdate(lockedId, { durationMs: 1 }); + for (let index = 0; index < 199; index++) { + writeBuffer.enqueueMessageRequestUpdate(-1_000_000 - index, { durationMs: index }); + } + + writeBuffer.enqueueMessageRequestUpdate(terminalId, { statusCode: 202 }); + writeBuffer.enqueueMessageRequestUpdate(evictedId, { model: "evicted-model" }); + for (let index = 0; index < 4_998; index++) { + writeBuffer.enqueueMessageRequestUpdate(-2_000_000 - index, { model: `filler-${index}` }); + } + writeBuffer.enqueueMessageRequestUpdate(retainedId, { model: "retained-model" }); + return { flush: writeBuffer.flushMessageRequestWriteBuffer() }; + }); + + await flush; + await expectRequest(lockedId, { durationMs: 1 }); + await expectRequest(terminalId, { statusCode: 202 }); + await expectRequest(evictedId, { model: "integration-model" }); + await expectRequest(retainedId, { model: "retained-model" }); + }); + + test("retries at bounded cadence and retains authoritative loser costs", async () => { + const requestId = await createRequest("cost-retry", "0"); + await messageRepository.addMessageRequestHedgeLoserCost(requestId, "0.02", { + providerId: 31, + providerName: "first-loser", + attemptNumber: 1, + costUsd: "0.02", + }); + + const startedAt = performance.now(); + await lockClient.begin(async (transaction) => { + await transaction`SELECT id FROM message_request WHERE id = ${requestId} FOR UPDATE`; + await expect( + messageRepository.updateMessageRequestWinnerCost(requestId, "0.10") + ).rejects.toMatchObject({ cause: { code: "55P03" } }); + }); + expect(performance.now() - startedAt).toBeGreaterThanOrEqual(350); + + await messageRepository.updateMessageRequestWinnerCost(requestId, "0.10"); + await messageRepository.addMessageRequestHedgeLoserCost(requestId, "0.03", { + providerId: 32, + providerName: "second-loser", + attemptNumber: 2, + costUsd: "0.03", + }); + await messageRepository.updateMessageRequestWinnerCost(requestId, "0.10"); + await expectRequest(requestId, { costUsd: "0.150000000000000" }); + }); +}); diff --git a/tests/integration/proxy-hedge-lifecycle.test.ts b/tests/integration/proxy-hedge-lifecycle.test.ts new file mode 100644 index 000000000..ec3e9a533 --- /dev/null +++ b/tests/integration/proxy-hedge-lifecycle.test.ts @@ -0,0 +1,627 @@ +// Real loopback transport exercises production lifecycle code; persistence/control-plane seams are mocked. +import { createServer, type ServerResponse } from "node:http"; +import type { Socket } from "node:net"; +import { Context } from "hono"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { ProxyForwarder } from "@/app/v1/_lib/proxy/forwarder"; +import { ProxyResponseHandler } from "@/app/v1/_lib/proxy/response-handler"; +import { type MessageContext, ProxySession } from "@/app/v1/_lib/proxy/session"; +import { DbPoolAdmissionError } from "@/drizzle/admitted-client"; +import { getGlobalAgentPool, resetGlobalAgentPool } from "@/lib/proxy-agent"; +import type { Key } from "@/types/key"; +import type { Provider } from "@/types/provider"; +import type { User } from "@/types/user"; + +const state = vi.hoisted(() => { + return { + addLoserCost: vi.fn(), + billHedgeLosers: false, + durableTerminal: vi.fn(async () => {}), + http2Error: ((): Error | null => null)(), + loserBilled: Promise.withResolvers(), + pickAlternative: vi.fn(), + providers: Array.from([]), + recordFailure: vi.fn(async () => {}), + settleLeaseBudgets: vi.fn(async () => {}), + tasks: Array.from>([]), + trackCost: vi.fn(async () => {}), + updateMessageRequestCostWithBreakdown: vi.fn(async () => {}), + updateMessageRequestDetailsIfUnfinalized: vi.fn(async () => {}), + updateWinnerCost: vi.fn(async () => {}), + }; +}); + +vi.mock("@/lib/logger", () => ({ + logger: Object.fromEntries( + ["debug", "error", "fatal", "info", "trace", "warn"].map((level) => [level, vi.fn()]) + ), +})); +vi.mock("@/lib/config", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + getCachedSystemSettings: async () => ({ + billHedgeLosers: state.billHedgeLosers, + enableBillingHeaderRectifier: false, + enableClaudeMetadataUserIdInjection: false, + enableThinkingBudgetRectifier: false, + enableThinkingSignatureRectifier: false, + }), + isHttp2Enabled: async () => { + if (state.http2Error) throw state.http2Error; + return true; + }, + }; +}); +vi.mock("@/lib/config/system-settings-cache", () => ({ + getCachedSystemSettings: async () => ({ billNonSuccessfulRequests: false }), +})); +vi.mock("@/app/v1/_lib/proxy/provider-selector", () => ({ + ProxyProviderResolver: { pickRandomProviderWithExclusion: state.pickAlternative }, +})); +vi.mock("@/lib/provider-endpoints/endpoint-selector", () => ({ + getEndpointFilterStats: vi.fn(async () => null), + getPreferredProviderEndpoints: vi.fn(async () => []), +})); +vi.mock("@/lib/circuit-breaker", () => ({ + getCircuitState: vi.fn(() => "closed"), + getProviderHealthInfo: vi.fn(async () => ({ + config: { failureThreshold: 3 }, + health: { failureCount: 0 }, + })), + recordFailure: state.recordFailure, + recordSuccess: vi.fn(), +})); +vi.mock("@/lib/endpoint-circuit-breaker", () => ({ + recordEndpointFailure: vi.fn(), + recordEndpointSuccess: vi.fn(), + resetEndpointCircuit: vi.fn(), +})); +vi.mock("@/lib/vendor-type-circuit-breaker", () => ({ + isVendorTypeCircuitOpen: vi.fn(async () => false), + recordVendorTypeAllEndpointsTimeout: vi.fn(), +})); +vi.mock("@/lib/rate-limit/service", () => ({ + RateLimitService: { + checkAndTrackProviderSession: vi.fn(async () => ({ allowed: true })), + releaseProviderSession: vi.fn(), + }, +})); +vi.mock("@/lib/rate-limit", () => ({ + RateLimitService: { + settleLeaseBudgets: state.settleLeaseBudgets, + trackCost: state.trackCost, + trackUserDailyCost: vi.fn(async () => {}), + }, +})); +vi.mock("@/lib/request-filter-engine", () => ({ + requestFilterEngine: { applyFinal: vi.fn(async () => {}) }, +})); +vi.mock("@/app/v1/_lib/proxy/response-fixer", () => ({ + ResponseFixer: { process: async (_session: ProxySession, response: Response) => response }, +})); +vi.mock("@/lib/async-task-manager", () => ({ + AsyncTaskManager: { + cancel: vi.fn(), + cleanup: vi.fn(), + register: ( + _taskId: string, + factory: (signal: AbortSignal) => Promise, + options?: string | { readonly abortController?: AbortController } + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + const task = Promise.resolve().then(() => factory(controller.signal)); + void task.catch(() => undefined); + state.tasks.push(task); + return controller; + }, + touch: vi.fn(() => true), + }, +})); +vi.mock("@/repository/message", () => ({ + addMessageRequestHedgeLoserCost: state.addLoserCost, + updateMessageRequestCostWithBreakdown: state.updateMessageRequestCostWithBreakdown, + updateMessageRequestDetails: vi.fn(async () => {}), + updateMessageRequestDetailsDurably: state.durableTerminal, + updateMessageRequestDetailsIfUnfinalized: state.updateMessageRequestDetailsIfUnfinalized, + updateMessageRequestDuration: vi.fn(async () => {}), + updateMessageRequestWinnerCost: state.updateWinnerCost, +})); +vi.mock("@/repository/model-price", () => ({ + findLatestPriceByModel: vi.fn(async (modelName: string) => ({ + createdAt: new Date(0), + id: 1, + modelName, + priceData: { input_cost_per_token: 0.001, output_cost_per_token: 0.002 }, + source: "litellm", + updatedAt: new Date(0), + })), +})); +vi.mock("@/repository/system-config", () => ({ + getSystemSettings: vi.fn(async () => ({ + billingModelSource: "redirected", + codexPriorityBillingSource: "requested", + })), +})); +vi.mock("@/lib/price-sync/cloud-price-updater", () => ({ requestCloudPriceTableSync: vi.fn() })); +vi.mock("@/lib/langfuse/emit-proxy-trace", () => ({ emitProxyLangfuseTrace: vi.fn() })); +vi.mock("@/lib/session-tracker", () => ({ + SessionTracker: { refreshSession: vi.fn(async () => {}) }, +})); +vi.mock("@/lib/proxy-status-tracker", () => ({ + ProxyStatusTracker: { getInstance: () => ({ endRequest: vi.fn() }) }, +})); +vi.mock("@/lib/redis/live-chain-store", () => ({ deleteLiveChain: vi.fn(async () => {}) })); + +const CREATED_AT = new Date(0); +const USER = { + createdAt: CREATED_AT, + dailyQuota: null, + dailyResetMode: "fixed", + dailyResetTime: "00:00", + description: "hedge lifecycle user", + id: 21, + isEnabled: true, + limit5hResetMode: "fixed", + name: "hedge-user", + providerGroup: null, + role: "user", + rpm: null, + updatedAt: CREATED_AT, +} satisfies User; +const KEY = { + cacheTtlPreference: null, + canLoginWebUi: false, + createdAt: CREATED_AT, + dailyResetMode: "fixed", + dailyResetTime: "00:00", + id: 22, + isEnabled: true, + key: "sk-hedge-lifecycle", + limit5hResetMode: "fixed", + limit5hUsd: null, + limitConcurrentSessions: 0, + limitDailyUsd: null, + limitMonthlyUsd: null, + limitWeeklyUsd: null, + name: "hedge-key", + providerGroup: null, + updatedAt: CREATED_AT, + userId: USER.id, +} satisfies Key; +const MESSAGE = { + apiKey: KEY.key, + createdAt: CREATED_AT, + id: 51, + key: KEY, + user: USER, +} satisfies MessageContext; + +function createProvider(id: number, url: string, firstByteTimeoutStreamingMs: number): Provider { + return { + activeTimeEnd: null, + activeTimeStart: null, + allowedClients: [], + allowedModels: null, + anthropicAdaptiveThinking: null, + anthropicMaxTokensPreference: null, + anthropicThinkingBudgetPreference: null, + blockedClients: [], + cacheTtlPreference: null, + cc: 0, + circuitBreakerFailureThreshold: 5, + circuitBreakerHalfOpenSuccessThreshold: 2, + circuitBreakerOpenDuration: 1_800_000, + codexImageGenerationPreference: null, + codexParallelToolCallsPreference: null, + codexReasoningEffortPreference: null, + codexReasoningSummaryPreference: null, + codexServiceTierPreference: null, + codexTextVerbosityPreference: null, + context1mPreference: null, + costMultiplier: 1, + createdAt: CREATED_AT, + customHeaders: null, + dailyResetMode: "fixed", + dailyResetTime: "00:00", + disableSessionReuse: false, + faviconUrl: null, + firstByteTimeoutStreamingMs, + geminiGoogleSearchPreference: null, + groupPriorities: null, + groupTag: null, + id, + isEnabled: true, + key: `provider-key-${id}`, + limit5hResetMode: "fixed", + limit5hUsd: null, + limitConcurrentSessions: 0, + limitDailyUsd: null, + limitMonthlyUsd: null, + limitTotalUsd: null, + limitWeeklyUsd: null, + maxRetryAttempts: 1, + mcpPassthroughType: "none", + mcpPassthroughUrl: null, + modelRedirects: null, + name: `provider-${id}`, + preserveClientIp: false, + priority: id, + providerType: "claude", + providerVendorId: null, + proxyFallbackToDirect: false, + proxyUrl: null, + requestTimeoutNonStreamingMs: 0, + rpd: 0, + rpm: 0, + streamingIdleTimeoutMs: 0, + swapCacheTtlBilling: false, + totalCostResetAt: null, + tpm: 0, + updatedAt: CREATED_AT, + url, + websiteUrl: null, + weight: 1, + }; +} + +type Upstream = { + readonly abortCount: () => number; + readonly baseUrl: string; + readonly close: () => Promise; + readonly requestCount: () => number; + readonly response: Promise; + readonly send: (body: string) => Promise; + readonly terminated: Promise; +}; + +async function startUpstream(): Promise { + const sockets = new Set(); + const responseGate = Promise.withResolvers(); + const terminationGate = Promise.withResolvers(); + let requests = 0; + let aborts = 0; + const server = createServer((request, response) => { + requests += 1; + request.resume(); + response.writeHead(200, { "content-type": "text/event-stream" }); + response.flushHeaders(); + response.once("close", () => { + if (!response.writableEnded) aborts += 1; + terminationGate.resolve(); + }); + responseGate.resolve(response); + }); + server.on("connection", (socket) => { + sockets.add(socket); + socket.once("close", () => sockets.delete(socket)); + }); + const baseUrl = await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", () => { + const address = server.address(); + if (!address || typeof address === "string") { + reject(new Error("Loopback fixture did not receive a TCP address")); + return; + } + resolve(`http://127.0.0.1:${address.port}`); + }); + }); + return { + abortCount: () => aborts, + baseUrl, + close: async () => { + for (const socket of sockets) socket.destroy(); + sockets.clear(); + await new Promise((resolve) => server.close(() => resolve())); + }, + requestCount: () => requests, + response: responseGate.promise, + send: async (body) => { + const response = await responseGate.promise; + await new Promise((resolve) => response.end(body, resolve)); + }, + terminated: terminationGate.promise, + }; +} + +async function createSession( + provider: Provider, + pathname: string = "/v1/messages", + signal?: AbortSignal +): Promise { + const request = new Request(`https://hub.test${pathname}`, { + body: JSON.stringify({ + max_tokens: 32, + messages: [{ content: "integration", role: "user" }], + model: "claude-test", + stream: true, + }), + headers: { "content-type": "application/json" }, + method: "POST", + ...(signal ? { signal } : {}), + }); + const session = await ProxySession.fromContext(new Context(request)); + session.setAuthState({ apiKey: KEY.key, key: KEY, success: true, user: USER }); + session.setMessageContext(MESSAGE); + session.setOriginalFormat("claude"); + session.setOriginalModel("claude-test"); + session.setProvider(provider); + return session; +} + +function sse(inputTokens: number, outputTokens: number): string { + return `event: message_delta\ndata: ${JSON.stringify({ + usage: { input_tokens: inputTokens, output_tokens: outputTokens }, + })}\n\nevent: message_stop\ndata: {"type":"message_stop"}\n\n`; +} + +async function settleTasks(): Promise { + while (state.tasks.length > 0) { + const settlements = await Promise.allSettled(state.tasks.splice(0, state.tasks.length)); + const failures = settlements.flatMap((settlement) => + settlement.status === "rejected" ? [settlement.reason] : [] + ); + if (failures.length > 0) throw new AggregateError(failures, "Proxy lifecycle task failed"); + } +} + +function watchAgentReleases(expectedReleases: number) { + const pool = getGlobalAgentPool(); + const originalRelease = pool.releaseAgent.bind(pool); + const released = Promise.withResolvers(); + let releaseCount = 0; + const release = vi.spyOn(pool, "releaseAgent").mockImplementation((cacheKey, dispatcherId) => { + originalRelease(cacheKey, dispatcherId); + releaseCount += 1; + if (releaseCount === expectedReleases) released.resolve(); + }); + return { acquire: vi.spyOn(pool, "getAgent"), pool, release, released: released.promise }; +} + +beforeEach(async () => { + await resetGlobalAgentPool(); + vi.clearAllMocks(); + state.billHedgeLosers = false; + state.http2Error = null; + state.loserBilled = Promise.withResolvers(); + state.providers.length = 0; + state.tasks.length = 0; + state.addLoserCost.mockImplementation(async () => state.loserBilled.resolve()); + state.pickAlternative.mockImplementation(async (_session: unknown, excludedIds: number[]) => { + return state.providers.find((provider) => !excludedIds.includes(provider.id)) ?? null; + }); +}); + +afterEach(async () => { + vi.useRealTimers(); + await settleTasks(); + await resetGlobalAgentPool(); +}); + +describe("proxy hedge transport/lifecycle integration (persistence and control-plane seams mocked)", () => { + it("fences loser timers after winner settlement and releases each launched transport once", async () => { + const [slow, winner, fenced] = await Promise.all([ + startUpstream(), + startUpstream(), + startUpstream(), + ]); + const client = new AbortController(); + try { + const initialProvider = createProvider(1, slow.baseUrl, 50); + state.providers.push( + createProvider(2, winner.baseUrl, 50), + createProvider(3, fenced.baseUrl, 50) + ); + const session = await createSession(initialProvider, "/v1/messages", client.signal); + const agents = watchAgentReleases(2); + vi.useFakeTimers({ toFake: ["clearTimeout", "setTimeout"] }); + + const forwarded = ProxyForwarder.send(session); + await slow.response; + await vi.advanceTimersByTimeAsync(50); + await winner.response; + await winner.send(sse(8, 2)); + const downstream = await ProxyResponseHandler.dispatch(session, await forwarded); + await expect(downstream.text()).resolves.toContain("message_stop"); + await settleTasks(); + await slow.terminated; + await agents.released; + await vi.advanceTimersByTimeAsync(500); + + expect(fenced.requestCount()).toBe(0); + expect(slow.abortCount()).toBe(1); + expect(winner.abortCount()).toBe(0); + expect(agents.acquire).toHaveBeenCalledTimes(2); + expect(agents.release).toHaveBeenCalledTimes(2); + expect(new Set(agents.release.mock.calls.map(([key, id]) => `${key}|${id}`))).toHaveLength(2); + expect(agents.pool.getPoolStats().activeRequests).toBe(0); + } finally { + client.abort(new Error("fixture cleanup")); + await Promise.all([slow.close(), winner.close(), fenced.close()]); + } + }); + + it("classifies local database overload before upstream fanout", async () => { + const [initial, alternative] = await Promise.all([startUpstream(), startUpstream()]); + try { + const initialProvider = createProvider(1, initial.baseUrl, 50); + state.providers.push(createProvider(2, alternative.baseUrl, 50)); + const session = await createSession(initialProvider); + const wrapped = new Error("Failed query", { + cause: new DbPoolAdmissionError("data", 32), + }); + state.http2Error = wrapped; + + await expect(ProxyForwarder.send(session)).rejects.toBe(wrapped); + + expect(initial.requestCount()).toBe(0); + expect(alternative.requestCount()).toBe(0); + expect(state.pickAlternative).not.toHaveBeenCalled(); + expect(state.recordFailure).not.toHaveBeenCalled(); + expect(session.getProviderChain()).toEqual([ + expect.objectContaining({ + errorDetails: expect.objectContaining({ + system: expect.objectContaining({ errorCode: "DB_POOL_ADMISSION_EXCEEDED" }), + }), + reason: "system_error", + }), + ]); + } finally { + await Promise.all([initial.close(), alternative.close()]); + } + }); + + it("bills the hedge winner and naturally drained loser exactly once", async () => { + const [loser, winner] = await Promise.all([startUpstream(), startUpstream()]); + const client = new AbortController(); + try { + state.billHedgeLosers = true; + const initialProvider = createProvider(1, loser.baseUrl, 50); + state.providers.push(createProvider(2, winner.baseUrl, 50)); + const session = await createSession(initialProvider, "/v1/messages", client.signal); + const agents = watchAgentReleases(2); + vi.useFakeTimers({ toFake: ["clearTimeout", "setTimeout"] }); + + const forwarded = ProxyForwarder.send(session); + await loser.response; + await vi.advanceTimersByTimeAsync(50); + await winner.response; + await winner.send(sse(10, 3)); + const downstream = await ProxyResponseHandler.dispatch(session, await forwarded); + await downstream.text(); + await settleTasks(); + await loser.send(sse(7, 2)); + await state.loserBilled.promise; + await agents.released; + + expect(state.updateWinnerCost).toHaveBeenCalledTimes(1); + expect(state.updateWinnerCost.mock.calls[0]?.[0]).toBe(MESSAGE.id); + expect(String(state.updateWinnerCost.mock.calls[0]?.[1])).toBe("0.016"); + expect(state.updateMessageRequestCostWithBreakdown).not.toHaveBeenCalled(); + + expect(state.addLoserCost).toHaveBeenCalledTimes(1); + expect(state.addLoserCost.mock.calls[0]?.[0]).toBe(MESSAGE.id); + expect(String(state.addLoserCost.mock.calls[0]?.[1])).toBe("0.011"); + expect(state.addLoserCost.mock.calls[0]?.[2]).toEqual( + expect.objectContaining({ attemptNumber: 1, providerId: initialProvider.id }) + ); + + expect(state.durableTerminal).toHaveBeenCalledTimes(1); + expect(state.durableTerminal).toHaveBeenCalledWith( + MESSAGE.id, + expect.objectContaining({ + inputTokens: 10, + outputTokens: 3, + providerId: 2, + statusCode: 200, + }) + ); + expect(state.updateMessageRequestDetailsIfUnfinalized).not.toHaveBeenCalled(); + + expect(state.trackCost).toHaveBeenCalledTimes(2); + expect(state.trackCost).toHaveBeenNthCalledWith( + 1, + KEY.id, + 2, + "", + 0.016, + expect.objectContaining({ requestId: MESSAGE.id, userId: USER.id }) + ); + expect(state.trackCost).toHaveBeenNthCalledWith( + 2, + KEY.id, + initialProvider.id, + "", + 0.011, + expect.objectContaining({ + requestId: `${MESSAGE.id}:hedge-loser:${initialProvider.id}:1`, + userId: USER.id, + }) + ); + expect(state.settleLeaseBudgets).toHaveBeenCalledTimes(2); + expect(state.settleLeaseBudgets).toHaveBeenNthCalledWith( + 1, + expect.objectContaining({ + cost: 0.016, + entities: expect.objectContaining({ + provider: expect.objectContaining({ id: 2 }), + }), + requestId: MESSAGE.id, + }) + ); + expect(state.settleLeaseBudgets).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ + cost: 0.011, + entities: expect.objectContaining({ + provider: expect.objectContaining({ id: initialProvider.id }), + }), + requestId: `${MESSAGE.id}:hedge-loser:${initialProvider.id}:1`, + }) + ); + expect(loser.requestCount()).toBe(1); + expect(winner.requestCount()).toBe(1); + expect(loser.abortCount()).toBe(0); + expect(winner.abortCount()).toBe(0); + expect(agents.release).toHaveBeenCalledTimes(2); + expect(agents.pool.getPoolStats().activeRequests).toBe(0); + + await vi.advanceTimersByTimeAsync(500); + await settleTasks(); + + expect(state.pickAlternative).toHaveBeenCalledTimes(1); + expect(loser.requestCount()).toBe(1); + expect(winner.requestCount()).toBe(1); + expect(state.updateWinnerCost).toHaveBeenCalledTimes(1); + expect(state.addLoserCost).toHaveBeenCalledTimes(1); + expect(state.updateMessageRequestCostWithBreakdown).not.toHaveBeenCalled(); + expect(state.durableTerminal).toHaveBeenCalledTimes(1); + expect(state.updateMessageRequestDetailsIfUnfinalized).not.toHaveBeenCalled(); + expect(state.trackCost).toHaveBeenCalledTimes(2); + expect(state.settleLeaseBudgets).toHaveBeenCalledTimes(2); + expect(loser.abortCount()).toBe(0); + expect(winner.abortCount()).toBe(0); + expect(agents.release).toHaveBeenCalledTimes(2); + } finally { + client.abort(new Error("fixture cleanup")); + await Promise.all([loser.close(), winner.close()]); + } + }); + + it("keeps the first-byte deadline armed across the public response handoff", async () => { + const silent = await startUpstream(); + const client = new AbortController(); + try { + const provider = createProvider(1, silent.baseUrl, 50); + const session = await createSession(provider, "/v1/messages/count_tokens", client.signal); + const agents = watchAgentReleases(1); + vi.useFakeTimers({ toFake: ["clearTimeout", "setTimeout"] }); + + const forwarded = ProxyForwarder.send(session); + await silent.response; + const downstream = await ProxyResponseHandler.dispatch(session, await forwarded); + const bodyRejection = expect(downstream.text()).rejects.toThrow(); + await vi.advanceTimersByTimeAsync(50); + + await bodyRejection; + await settleTasks(); + await silent.terminated; + await agents.released; + expect(silent.abortCount()).toBe(1); + expect(state.durableTerminal).toHaveBeenCalledOnce(); + expect(state.durableTerminal).toHaveBeenCalledWith( + MESSAGE.id, + expect.objectContaining({ statusCode: 502 }) + ); + expect(agents.release).toHaveBeenCalledOnce(); + expect(agents.pool.getPoolStats().activeRequests).toBe(0); + } finally { + client.abort(new Error("fixture cleanup")); + await silent.close(); + } + }); +}); diff --git a/tests/integration/rolling-cost-redis.test.ts b/tests/integration/rolling-cost-redis.test.ts new file mode 100644 index 000000000..74c44a6a4 --- /dev/null +++ b/tests/integration/rolling-cost-redis.test.ts @@ -0,0 +1,130 @@ +import Redis from "ioredis"; +import { afterAll, afterEach, beforeAll, describe, expect, test } from "vitest"; +import { GET_COST_5H_ROLLING_WINDOW, TRACK_COST_ROLLING_WINDOW } from "@/lib/redis/lua-scripts"; + +const HAS_REDIS = Boolean(process.env.REDIS_URL); +const run = describe.skipIf(!HAS_REDIS); +const TEST_PREFIX = `it-rolling-cost-${Date.now()}-${Math.random().toString(16).slice(2)}`; +const WINDOW_MS = 5 * 60 * 60 * 1000; +const TTL_SECONDS = 60; + +run("rolling cost Lua integration", () => { + let redis: Redis; + const touchedKeys = new Set(); + + function nextKey(tag: string): string { + const key = `${TEST_PREFIX}:${tag}`; + touchedKeys.add(key); + return key; + } + + async function track(params: { + key: string; + cost: number; + nowMs: number; + requestId?: string; + }): Promise { + return redis.eval( + TRACK_COST_ROLLING_WINDOW, + 1, + params.key, + params.cost.toString(), + params.nowMs.toString(), + WINDOW_MS.toString(), + params.requestId ?? "", + TTL_SECONDS.toString() + ); + } + + async function getTotal(key: string, nowMs: number): Promise { + const result = await redis.eval( + GET_COST_5H_ROLLING_WINDOW, + 1, + key, + nowMs.toString(), + WINDOW_MS.toString() + ); + return Number(result); + } + + beforeAll(async () => { + redis = new Redis(process.env.REDIS_URL!, { + lazyConnect: true, + enableOfflineQueue: false, + maxRetriesPerRequest: 1, + }); + await redis.connect(); + await expect(redis.ping()).resolves.toBe("PONG"); + }); + + afterEach(async () => { + if (touchedKeys.size > 0) { + await redis.del(...touchedKeys); + touchedKeys.clear(); + } + }); + + afterAll(async () => { + if (redis?.status !== "end") { + await redis.quit(); + } + }); + + test("write-only tracking preserves valid state, replay cardinality, TTL, cutoff, and exact GET", async () => { + const key = nextKey("valid"); + const nowMs = 1_700_000_000_000; + const expiredAt = nowMs - WINDOW_MS - 1; + const retainedAt = nowMs - 1_000; + + await redis.zadd(key, expiredAt, `${expiredAt}:expired:4`); + await redis.zadd(key, retainedAt, `${retainedAt}:retained:1.5`); + + await expect(track({ key, cost: 2.5, nowMs, requestId: "request-1" })).resolves.toBe(1); + await expect(track({ key, cost: 2.5, nowMs, requestId: "request-1" })).resolves.toBe(1); + + expect(await redis.zcard(key)).toBe(2); + expect(await redis.zscore(key, `${expiredAt}:expired:4`)).toBeNull(); + expect(await redis.zscore(key, `${retainedAt}:retained:1.5`)).toBe(String(retainedAt)); + expect(await redis.zscore(key, `${nowMs}:request-1:2.5`)).toBe(String(nowMs)); + + const ttl = await redis.ttl(key); + expect(ttl).toBeGreaterThan(0); + expect(ttl).toBeLessThanOrEqual(TTL_SECONDS); + await expect(getTotal(key, nowMs)).resolves.toBeCloseTo(4, 10); + }); + + test("WRONGTYPE fails before mutation", async () => { + const key = nextKey("wrongtype"); + await redis.set(key, "not-a-zset"); + + await expect( + track({ key, cost: 2.5, nowMs: 1_700_000_000_000, requestId: "request-2" }) + ).rejects.toThrow(/WRONGTYPE/); + + await expect(redis.get(key)).resolves.toBe("not-a-zset"); + await expect(redis.ttl(key)).resolves.toBe(-1); + }); + + test("malformed members do not block writes and TTL repair, while exact GET remains strict", async () => { + const key = nextKey("malformed"); + const nowMs = 1_700_000_000_000; + const expiredAt = nowMs - WINDOW_MS - 1; + const retainedAt = nowMs - 1_000; + + await redis.zadd(key, expiredAt, "expired:non-numeric-cost"); + await redis.zadd(key, retainedAt, "retained:non-numeric-cost"); + await expect(redis.ttl(key)).resolves.toBe(-1); + + await expect(track({ key, cost: 3, nowMs, requestId: "request-3" })).resolves.toBe(1); + + expect(await redis.zscore(key, "expired:non-numeric-cost")).toBeNull(); + expect(await redis.zscore(key, "retained:non-numeric-cost")).toBe(String(retainedAt)); + expect(await redis.zscore(key, `${nowMs}:request-3:3`)).toBe(String(nowMs)); + expect(await redis.ttl(key)).toBeGreaterThan(0); + + await expect(getTotal(key, nowMs)).rejects.toThrow(); + + await redis.zrem(key, "retained:non-numeric-cost"); + await expect(getTotal(key, nowMs)).resolves.toBe(3); + }); +}); diff --git a/tests/unit/actions/providers-patch-contract.test.ts b/tests/unit/actions/providers-patch-contract.test.ts index 35a2e3047..fa0082ecc 100644 --- a/tests/unit/actions/providers-patch-contract.test.ts +++ b/tests/unit/actions/providers-patch-contract.test.ts @@ -850,21 +850,19 @@ describe("provider patch contract", () => { }); describe("MCP fields", () => { - it.each([ - "none", - "minimax", - "glm", - "custom", - ] as const)("accepts mcp_passthrough_type value: %s", (value) => { - const result = prepareProviderBatchApplyUpdates({ - mcp_passthrough_type: { set: value }, - }); - - expect(result.ok).toBe(true); - if (!result.ok) return; - - expect(result.data.mcp_passthrough_type).toBe(value); - }); + it.each(["none", "minimax", "glm", "custom"] as const)( + "accepts mcp_passthrough_type value: %s", + (value) => { + const result = prepareProviderBatchApplyUpdates({ + mcp_passthrough_type: { set: value }, + }); + + expect(result.ok).toBe(true); + if (!result.ok) return; + + expect(result.data.mcp_passthrough_type).toBe(value); + } + ); it("rejects invalid mcp_passthrough_type value", () => { const result = normalizeProviderBatchPatchDraft({ diff --git a/tests/unit/api/actions/legacy-deprecation.test.ts b/tests/unit/api/actions/legacy-deprecation.test.ts index 4292eac45..f1740c44c 100644 --- a/tests/unit/api/actions/legacy-deprecation.test.ts +++ b/tests/unit/api/actions/legacy-deprecation.test.ts @@ -75,20 +75,20 @@ describe("legacy actions API deprecation", () => { expectManagementSecurityHeaders(response); }); - test.each([ - "/api/actions/docs", - "/api/actions/scalar", - ])("keeps legacy docs UI %s available when execution is disabled but docs mode is deprecated", async (pathname) => { - vi.stubEnv("ENABLE_LEGACY_ACTIONS_API", "false"); - vi.stubEnv("LEGACY_ACTIONS_DOCS_MODE", "deprecated"); - - const response = await callFreshActionsRoute(pathname, "GET"); - - expect(response.status).toBe(200); - expect(response.headers.get("Deprecation")).toBe("@1777420800"); - expect(response.headers.get("Link")).toContain("/api/v1/openapi.json"); - expectManagementSecurityHeaders(response); - }); + test.each(["/api/actions/docs", "/api/actions/scalar"])( + "keeps legacy docs UI %s available when execution is disabled but docs mode is deprecated", + async (pathname) => { + vi.stubEnv("ENABLE_LEGACY_ACTIONS_API", "false"); + vi.stubEnv("LEGACY_ACTIONS_DOCS_MODE", "deprecated"); + + const response = await callFreshActionsRoute(pathname, "GET"); + + expect(response.status).toBe(200); + expect(response.headers.get("Deprecation")).toBe("@1777420800"); + expect(response.headers.get("Link")).toContain("/api/v1/openapi.json"); + expectManagementSecurityHeaders(response); + } + ); test("keeps deprecation date stable when sunset date is overridden", async () => { vi.stubEnv("LEGACY_ACTIONS_SUNSET_DATE", "2027-01-15"); @@ -114,21 +114,21 @@ describe("legacy actions API deprecation", () => { }); }); - test.each([ - "/api/actions/docs", - "/api/actions/scalar", - ])("can hide legacy docs UI %s independently with the docs mode flag", async (pathname) => { - vi.stubEnv("ENABLE_LEGACY_ACTIONS_API", "true"); - vi.stubEnv("LEGACY_ACTIONS_DOCS_MODE", "hidden"); - - const response = await callFreshActionsRoute(pathname, "GET"); - const body = await response.json(); - - expect(response.status).toBe(410); - expect(body).toMatchObject({ - status: 410, - errorCode: "api.legacy_actions_gone", - instance: pathname, - }); - }); + test.each(["/api/actions/docs", "/api/actions/scalar"])( + "can hide legacy docs UI %s independently with the docs mode flag", + async (pathname) => { + vi.stubEnv("ENABLE_LEGACY_ACTIONS_API", "true"); + vi.stubEnv("LEGACY_ACTIONS_DOCS_MODE", "hidden"); + + const response = await callFreshActionsRoute(pathname, "GET"); + const body = await response.json(); + + expect(response.status).toBe(410); + expect(body).toMatchObject({ + status: 410, + errorCode: "api.legacy_actions_gone", + instance: pathname, + }); + } + ); }); diff --git a/tests/unit/api/v1/status-code-map.test.ts b/tests/unit/api/v1/status-code-map.test.ts index e38074a4b..ec62a2e57 100644 --- a/tests/unit/api/v1/status-code-map.test.ts +++ b/tests/unit/api/v1/status-code-map.test.ts @@ -15,10 +15,11 @@ describe("v1 status code map", () => { [415, "Unsupported media type", "request.unsupported_media_type"], [429, "Too many requests", "rate_limit.exceeded"], [503, "Service unavailable", "dependency.unavailable"], - ] as Array< - [ProblemStatusCode, string, string] - >)("maps %s to defaults", (status, title, errorCode) => { - expect(getDefaultProblemTitle(status)).toBe(title); - expect(getDefaultErrorCode(status)).toBe(errorCode); - }); + ] as Array<[ProblemStatusCode, string, string]>)( + "maps %s to defaults", + (status, title, errorCode) => { + expect(getDefaultProblemTitle(status)).toBe(title); + expect(getDefaultErrorCode(status)).toBe(errorCode); + } + ); }); diff --git a/tests/unit/dashboard/user-insights-page.test.tsx b/tests/unit/dashboard/user-insights-page.test.tsx index 8d9873c56..8c7fa30d5 100644 --- a/tests/unit/dashboard/user-insights-page.test.tsx +++ b/tests/unit/dashboard/user-insights-page.test.tsx @@ -214,8 +214,8 @@ describe("UserInsightsView", () => { expect(heading!.textContent).toContain("TestUser"); expect(mockGetUserInsightsOverview).toHaveBeenCalledWith( 10, - resolveTimePresetDates("7days").startDate, - resolveTimePresetDates("7days").endDate + resolveTimePresetDates("7days", "UTC").startDate, + resolveTimePresetDates("7days", "UTC").endDate ); unmount(); @@ -265,7 +265,7 @@ describe("UserInsightsView", () => { await flushMicrotasks(); - const { startDate, endDate } = resolveTimePresetDates("30days"); + const { startDate, endDate } = resolveTimePresetDates("30days", "UTC"); expect(mockGetUserInsightsOverview).toHaveBeenLastCalledWith(10, startDate, endDate); unmount(); diff --git a/tests/unit/deploy-dockerfile-contract.test.ts b/tests/unit/deploy-dockerfile-contract.test.ts new file mode 100644 index 000000000..f63bd405c --- /dev/null +++ b/tests/unit/deploy-dockerfile-contract.test.ts @@ -0,0 +1,19 @@ +import { readFileSync } from "node:fs"; +import { resolve } from "node:path"; +import { describe, expect, it } from "vitest"; + +describe("deploy/Dockerfile runtime contract", () => { + it("runs as node with writable, environment-redacted diagnostic reports", () => { + const dockerfile = readFileSync(resolve(process.cwd(), "deploy/Dockerfile"), "utf8"); + const reportsDirectory = "RUN mkdir -p /app/reports && chown node:node /app/reports"; + const user = "USER node"; + const command = + 'CMD ["node", "--report-on-fatalerror", "--report-uncaught-exception", "--report-exclude-env", "--report-directory=/app/reports", "server.js"]'; + + expect(dockerfile).toContain(reportsDirectory); + expect(dockerfile).toContain(user); + expect(dockerfile).toContain(command); + expect(dockerfile.indexOf(reportsDirectory)).toBeLessThan(dockerfile.indexOf(user)); + expect(dockerfile.indexOf(user)).toBeLessThan(dockerfile.indexOf(command)); + }); +}); diff --git a/tests/unit/drizzle/db-admission.test.ts b/tests/unit/drizzle/db-admission.test.ts new file mode 100644 index 000000000..39e9bade4 --- /dev/null +++ b/tests/unit/drizzle/db-admission.test.ts @@ -0,0 +1,225 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +interface Deferred { + promise: Promise; + resolve: (value: T) => void; + reject: (error: unknown) => void; +} + +function deferred(): Deferred { + let resolve: (value: T) => void = () => {}; + let reject: (error: unknown) => void = () => {}; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + return { promise, resolve, reject }; +} + +interface RawClient { + unsafe: ReturnType; + begin: ReturnType; + end: ReturnType; +} + +describe("drizzle/db outstanding admission", () => { + const originalEnv = { + NODE_ENV: process.env.NODE_ENV, + DSN: process.env.DSN, + DB_POOL_MAX: process.env.DB_POOL_MAX, + }; + + let queryDeferreds: Deferred[]; + let transactionDeferreds: Deferred[]; + let rawClient: RawClient; + let postgresMock: ReturnType; + + beforeEach(() => { + vi.resetModules(); + queryDeferreds = []; + transactionDeferreds = []; + rawClient = { + unsafe: vi.fn(() => { + const pending = deferred(); + queryDeferreds.push(pending); + return pending.promise; + }), + begin: vi.fn(() => { + const pending = deferred(); + transactionDeferreds.push(pending); + return pending.promise; + }), + end: vi.fn(async () => {}), + }; + postgresMock = vi.fn(() => rawClient); + + process.env.NODE_ENV = "production"; + process.env.DSN = "postgres://postgres:postgres@localhost:5432/claude_code_hub_test"; + process.env.DB_POOL_MAX = "1"; + + vi.doMock("postgres", () => ({ default: postgresMock })); + vi.doMock("drizzle-orm/postgres-js", () => ({ + drizzle: (client: unknown) => ({ $client: client }), + })); + }); + + afterEach(() => { + for (const pending of queryDeferreds) pending.resolve([]); + for (const pending of transactionDeferreds) pending.resolve(undefined); + for (const [key, value] of Object.entries(originalEnv)) { + if (value === undefined) delete process.env[key]; + else process.env[key] = value; + } + }); + + it("在创建底层 unsafe query 前拒绝超额工作,并在 resolve/reject 后精确释放", async () => { + const { getDb } = await import("@/drizzle/db"); + const client = (getDb() as unknown as { $client: RawClient }).$client; + const observed: Promise[] = []; + const queries: PromiseLike[] = []; + let admissionError: unknown; + + for (let index = 0; index < 256; index += 1) { + try { + const query = client.unsafe("select 1"); + queries.push(query); + observed.push(Promise.resolve(query).catch(() => undefined)); + } catch (error) { + admissionError = error; + break; + } + } + + expect(admissionError).toMatchObject({ code: "DB_POOL_ADMISSION_EXCEEDED" }); + expect(rawClient.unsafe).toHaveBeenCalledTimes(queries.length); + expect(queries.length).toBeGreaterThan(0); + expect(queries.length).toBeLessThan(256); + + const duplicateObserver = Promise.resolve(queries[0]).catch(() => undefined); + queryDeferreds[0].resolve([]); + await Promise.all([observed[0], duplicateObserver]); + + const replacementAfterResolve = client.unsafe("select 2"); + const replacementResolveObserved = Promise.resolve(replacementAfterResolve).catch( + () => undefined + ); + expect(() => client.unsafe("select rejected after one exact release")).toThrowError( + expect.objectContaining({ code: "DB_POOL_ADMISSION_EXCEEDED" }) + ); + + queryDeferreds[1].reject(new Error("simulated query failure")); + await observed[1]; + + const replacementAfterReject = client.unsafe("select 3"); + const replacementRejectObserved = Promise.resolve(replacementAfterReject).catch( + () => undefined + ); + expect(() => client.unsafe("select still bounded")).toThrowError( + expect.objectContaining({ code: "DB_POOL_ADMISSION_EXCEEDED" }) + ); + + queryDeferreds.at(-2)?.resolve([]); + queryDeferreds.at(-1)?.resolve([]); + await Promise.all([replacementResolveObserved, replacementRejectObserved]); + }); + + it("transaction begin 与 raw unsafe 共用同一 admission,超额 begin 不创建底层事务", async () => { + const { getDb } = await import("@/drizzle/db"); + const client = (getDb() as unknown as { $client: RawClient }).$client; + const observed: Promise[] = []; + + for (let index = 0; index < 256; index += 1) { + try { + const query = client.unsafe("select saturate"); + observed.push(Promise.resolve(query).catch(() => undefined)); + } catch { + break; + } + } + expect(observed.length).toBeGreaterThan(0); + expect(observed.length).toBeLessThan(256); + + expect(() => client.begin(async () => undefined)).toThrowError( + expect.objectContaining({ code: "DB_POOL_ADMISSION_EXCEEDED" }) + ); + expect(rawClient.begin).not.toHaveBeenCalled(); + + queryDeferreds[0].resolve([]); + await observed[0]; + + const transaction = client.begin(async () => "done") as unknown as Promise; + expect(rawClient.begin).toHaveBeenCalledTimes(1); + expect(() => client.unsafe("select blocked by transaction")).toThrowError( + expect.objectContaining({ code: "DB_POOL_ADMISSION_EXCEEDED" }) + ); + + transactionDeferreds[0].resolve("done"); + await expect(transaction).resolves.toBe("done"); + + const acceptedAfterTransaction = client.unsafe("select after transaction"); + const acceptedObserved = Promise.resolve(acceptedAfterTransaction).catch(() => undefined); + queryDeferreds.at(-1)?.resolve([]); + await acceptedObserved; + }); + + it("cancel 未执行的 lazy query 时不应通过 then 重新启动查询", async () => { + const { createAdmittedSqlClient } = await import("@/drizzle/admitted-client"); + const makeLazyQuery = () => ({ + executed: false, + then: vi.fn(), + cancel: vi.fn(() => null), + }); + const first = makeLazyQuery(); + const second = makeLazyQuery(); + const client = { + unsafe: vi.fn().mockReturnValueOnce(first).mockReturnValueOnce(second), + begin: vi.fn(), + }; + const admitted = createAdmittedSqlClient(client, { pool: "data", maxOutstanding: 1 }); + + const firstQuery = admitted.unsafe() as typeof first; + firstQuery.cancel(); + + expect(first.then).not.toHaveBeenCalled(); + expect(first.cancel).toHaveBeenCalledOnce(); + expect(() => admitted.unsafe()).not.toThrow(); + }); + + it("tagged-template query 也受同一 admission 上限保护", async () => { + const { createAdmittedSqlClient } = await import("@/drizzle/admitted-client"); + const first = deferred(); + const second = deferred(); + const raw = Object.assign( + vi.fn().mockReturnValueOnce(first.promise).mockReturnValueOnce(second.promise), + { unsafe: vi.fn(), begin: vi.fn() } + ); + const admitted = createAdmittedSqlClient(raw, { pool: "data", maxOutstanding: 1 }); + + const firstQuery = admitted`select 1`; + expect(() => admitted`select 2`).toThrowError( + expect.objectContaining({ code: "DB_POOL_ADMISSION_EXCEEDED" }) + ); + expect(raw).toHaveBeenCalledTimes(1); + + first.resolve([]); + await firstQuery; + + const secondQuery = admitted`select 2`; + second.resolve([]); + await secondQuery; + expect(raw).toHaveBeenCalledTimes(2); + }); + + it("Drizzle query wrapper 从 cause 提取 SQLSTATE 但不暴露原始内容", async () => { + const { findSafeDatabaseError } = await import("@/drizzle/admitted-client"); + const details = findSafeDatabaseError({ + name: "DrizzleQueryError", + query: "select * from secrets where token = $1", + params: ["admission-canary"], + cause: { code: "55P03" }, + }); + + expect(details).toEqual({ kind: "query", code: "55P03", message: "Database query failed" }); + expect(JSON.stringify(details)).not.toContain("admission-canary"); + }); +}); diff --git a/tests/unit/drizzle/db-pool-config.test.ts b/tests/unit/drizzle/db-pool-config.test.ts index bc53b3d17..4e6508470 100644 --- a/tests/unit/drizzle/db-pool-config.test.ts +++ b/tests/unit/drizzle/db-pool-config.test.ts @@ -2,6 +2,12 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; type EnvSnapshot = Partial>; +interface MockSqlClient { + end: ReturnType; + unsafe: ReturnType; + begin: ReturnType; +} + function snapshotEnv(keys: string[]): EnvSnapshot { const snapshot: EnvSnapshot = {}; for (const key of keys) { @@ -27,25 +33,38 @@ describe("drizzle/db 连接池配置", () => { "DB_POOL_MAX", "DB_POOL_IDLE_TIMEOUT", "DB_POOL_CONNECT_TIMEOUT", + "DB_STATEMENT_TIMEOUT_MS", + "DB_LOCK_TIMEOUT_MS", "MESSAGE_REQUEST_WRITE_MODE", ]; - const postgresMock = vi.fn(); - const drizzleMock = vi.fn(() => ({ __db: true })); + const clients: MockSqlClient[] = []; + const postgresMock = vi.fn(() => { + const client: MockSqlClient = { + end: vi.fn(async () => {}), + unsafe: vi.fn(), + begin: vi.fn(), + }; + clients.push(client); + return client; + }); + const drizzleMock = vi.fn((client: MockSqlClient) => ({ client })); const originalEnv = snapshotEnv(envKeys); beforeEach(() => { vi.resetModules(); - postgresMock.mockReset(); - drizzleMock.mockReset(); + clients.length = 0; + postgresMock.mockClear(); + drizzleMock.mockClear(); - // 确保每个用例有一致的基础环境 process.env.DSN = "postgres://postgres:postgres@localhost:5432/claude_code_hub_test"; process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; delete process.env.DB_POOL_MAX; delete process.env.DB_POOL_IDLE_TIMEOUT; delete process.env.DB_POOL_CONNECT_TIMEOUT; + delete process.env.DB_STATEMENT_TIMEOUT_MS; + delete process.env.DB_LOCK_TIMEOUT_MS; vi.doMock("postgres", () => ({ default: postgresMock })); vi.doMock("drizzle-orm/postgres-js", () => ({ @@ -57,52 +76,235 @@ describe("drizzle/db 连接池配置", () => { restoreEnv(originalEnv); }); - it("生产环境默认 max=20、idle_timeout=20、connect_timeout=10", async () => { + it("生产环境把默认总预算 20 lazy 拆为 data=15、control=4、writer=1", async () => { process.env.NODE_ENV = "production"; - const { getDb } = await import("@/drizzle/db"); - getDb(); + const { getDb, getMessageWriterDb, runWithDataDbScope } = await import("@/drizzle/db"); + + expect(postgresMock).not.toHaveBeenCalled(); + + const controlDb = getDb(); + const writerDb = getMessageWriterDb(); + const dataDb = runWithDataDbScope(() => getDb()); + + expect(controlDb).not.toBe(writerDb); + expect(controlDb).not.toBe(dataDb); + expect(writerDb).not.toBe(dataDb); + expect(postgresMock).toHaveBeenCalledTimes(3); - expect(postgresMock).toHaveBeenCalledWith( - process.env.DSN, - expect.objectContaining({ - max: 20, - idle_timeout: 20, - connect_timeout: 10, - }) + const options = postgresMock.mock.calls.map((call) => call[1]); + expect(options).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + max: 15, + connection: expect.objectContaining({ + application_name: "claude-code-hub:data", + statement_timeout: 90_000, + lock_timeout: 5_000, + }), + }), + expect.objectContaining({ + max: 4, + connection: expect.objectContaining({ application_name: "claude-code-hub:control" }), + }), + expect.objectContaining({ + max: 1, + connection: expect.objectContaining({ application_name: "claude-code-hub:writer" }), + }), + ]) ); + expect(options.reduce((sum, option) => sum + option.max, 0)).toBe(20); + + getDb(); + getMessageWriterDb(); + runWithDataDbScope(() => getDb()); + expect(postgresMock).toHaveBeenCalledTimes(3); }); - it("开发环境默认 max=10", async () => { + it("开发和测试环境把默认总预算 10 拆为 data=7、control=2、writer=1", async () => { process.env.NODE_ENV = "development"; - const { getDb } = await import("@/drizzle/db"); + const { getDb, getMessageWriterDb, runWithDataDbScope } = await import("@/drizzle/db"); getDb(); + getMessageWriterDb(); + runWithDataDbScope(() => getDb()); - expect(postgresMock).toHaveBeenCalledWith( - process.env.DSN, - expect.objectContaining({ - max: 10, - }) - ); + const maxima = postgresMock.mock.calls.map((call) => call[1].max).sort((a, b) => a - b); + expect(maxima).toEqual([1, 2, 7]); + expect(maxima.reduce((sum, value) => sum + value, 0)).toBe(10); }); - it("支持通过 env 覆盖连接池参数", async () => { + it("自定义总预算仍只创建合计不超过 DB_POOL_MAX 的物理 pool", async () => { process.env.NODE_ENV = "production"; - process.env.DB_POOL_MAX = "50"; + process.env.DB_POOL_MAX = "24"; process.env.DB_POOL_IDLE_TIMEOUT = "30"; process.env.DB_POOL_CONNECT_TIMEOUT = "5"; + process.env.DB_STATEMENT_TIMEOUT_MS = "45000"; + process.env.DB_LOCK_TIMEOUT_MS = "1500"; - const { getDb } = await import("@/drizzle/db"); + const { getDb, getMessageWriterDb, runWithDataDbScope } = await import("@/drizzle/db"); getDb(); + getMessageWriterDb(); + runWithDataDbScope(() => getDb()); - expect(postgresMock).toHaveBeenCalledWith( - process.env.DSN, - expect.objectContaining({ - max: 50, - idle_timeout: 30, - connect_timeout: 5, - }) + const options = postgresMock.mock.calls.map((call) => call[1]); + expect(options.reduce((sum, option) => sum + option.max, 0)).toBe(24); + for (const option of options) { + expect(option).toEqual( + expect.objectContaining({ + idle_timeout: 30, + connect_timeout: 5, + connection: expect.objectContaining({ + statement_timeout: 45_000, + lock_timeout: 1_500, + }), + }) + ); + } + }); + + it("shutdown 只关闭已经 lazy 创建的物理 pool,且每条只关闭一次", async () => { + process.env.NODE_ENV = "production"; + + const { closeDbPools, getDb, getMessageWriterDb, runWithDataDbScope } = await import( + "@/drizzle/db" ); + getDb(); + getMessageWriterDb(); + runWithDataDbScope(() => getDb()); + + await closeDbPools(); + await closeDbPools(); + + expect(clients).toHaveLength(3); + for (const client of clients) { + expect(client.end).toHaveBeenCalledTimes(1); + expect(client.end).toHaveBeenCalledWith({ timeout: 5 }); + } + }); + + it("同步重入 shutdown 复用同一 pending Promise 且每个 pool 只关闭一次", async () => { + process.env.NODE_ENV = "production"; + + const { closeDbPools, getDb, getMessageWriterDb, runWithDataDbScope } = await import( + "@/drizzle/db" + ); + getDb(); + getMessageWriterDb(); + runWithDataDbScope(() => getDb()); + + let resolveClose!: () => void; + const closeBarrier = new Promise((resolve) => { + resolveClose = resolve; + }); + for (const client of clients) { + client.end.mockImplementation(() => closeBarrier); + } + + let reentrantClose: Promise | undefined; + clients[0].end.mockImplementationOnce(() => { + reentrantClose = closeDbPools(); + return closeBarrier; + }); + + let outerSettled = false; + let reentrantSettled = false; + const outerClose = closeDbPools(); + try { + reentrantClose?.then( + () => { + reentrantSettled = true; + }, + () => { + reentrantSettled = true; + } + ); + outerClose.then( + () => { + outerSettled = true; + }, + () => { + outerSettled = true; + } + ); + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect.soft(reentrantClose).toBeDefined(); + expect.soft(outerClose).toBe(reentrantClose); + expect.soft(clients.map((client) => client.end.mock.calls.length)).toEqual([1, 1, 1]); + expect.soft(outerSettled).toBe(false); + expect.soft(reentrantSettled).toBe(false); + } finally { + resolveClose(); + await Promise.all([outerClose, reentrantClose]); + } + }); + + it("pool closing 和 closed 状态都拒绝创建快照外连接", async () => { + process.env.NODE_ENV = "production"; + + const { closeDbPools, getDb, getMessageWriterDb, runWithDataDbScope } = await import( + "@/drizzle/db" + ); + getDb(); + getMessageWriterDb(); + runWithDataDbScope(() => getDb()); + + let resolveClose!: () => void; + const closeBarrier = new Promise((resolve) => { + resolveClose = resolve; + }); + for (const client of clients) { + client.end.mockImplementation(() => closeBarrier); + } + + const closing = closeDbPools(); + + expect(() => getDb()).toThrow("Database pools are closing"); + expect(() => getMessageWriterDb()).toThrow("Database pools are closing"); + expect(() => runWithDataDbScope(() => getDb())).toThrow("Database pools are closing"); + expect(postgresMock).toHaveBeenCalledTimes(3); + + resolveClose(); + await closing; + + expect(() => getDb()).toThrow("Database pools are closed"); + expect(() => getMessageWriterDb()).toThrow("Database pools are closed"); + expect(postgresMock).toHaveBeenCalledTimes(3); + await closeDbPools(); + + for (const client of clients) { + expect(client.end).toHaveBeenCalledTimes(1); + } + }); + + it.each([ + { total: 1, expectedPhysicalPools: 1 }, + { total: 2, expectedPhysicalPools: 2 }, + ])("极小总预算 $total 可共享 lane 且物理连接上限不超预算", async ({ + total, + expectedPhysicalPools, + }) => { + process.env.NODE_ENV = "production"; + process.env.DB_POOL_MAX = String(total); + + const { getDb, getMessageWriterDb, runWithDataDbScope } = await import("@/drizzle/db"); + const controlDb = getDb(); + const writerDb = getMessageWriterDb(); + const dataDb = runWithDataDbScope(() => getDb()); + + expect(controlDb).toBeDefined(); + expect(writerDb).toBeDefined(); + expect(dataDb).toBeDefined(); + expect(postgresMock).toHaveBeenCalledTimes(expectedPhysicalPools); + + const maxima = postgresMock.mock.calls.map((call) => call[1].max); + expect(maxima.every((value) => value >= 1)).toBe(true); + expect(maxima.reduce((sum, value) => sum + value, 0)).toBeLessThanOrEqual(total); + + vi.resetModules(); + clients.length = 0; + postgresMock.mockClear(); + drizzleMock.mockClear(); }); }); diff --git a/tests/unit/drizzle/db-scope.test.ts b/tests/unit/drizzle/db-scope.test.ts new file mode 100644 index 000000000..f7518a858 --- /dev/null +++ b/tests/unit/drizzle/db-scope.test.ts @@ -0,0 +1,93 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +describe("drizzle/db 数据面作用域", () => { + const originalEnv = { + NODE_ENV: process.env.NODE_ENV, + DSN: process.env.DSN, + DB_POOL_MAX: process.env.DB_POOL_MAX, + }; + + const postgresMock = vi.fn(() => ({ + unsafe: vi.fn(), + begin: vi.fn(), + end: vi.fn(async () => {}), + })); + const drizzleMock = vi.fn((client: unknown) => ({ client })); + + beforeEach(() => { + vi.resetModules(); + postgresMock.mockClear(); + drizzleMock.mockClear(); + process.env.NODE_ENV = "production"; + process.env.DSN = "postgres://postgres:postgres@localhost:5432/claude_code_hub_test"; + process.env.DB_POOL_MAX = "20"; + + vi.doMock("postgres", () => ({ default: postgresMock })); + vi.doMock("drizzle-orm/postgres-js", () => ({ drizzle: drizzleMock })); + }); + + afterEach(() => { + for (const [key, value] of Object.entries(originalEnv)) { + if (value === undefined) delete process.env[key]; + else process.env[key] = value; + } + }); + + it("默认使用 control pool,data scope 跨 await 和 timer 保持隔离", async () => { + const { getDb, runWithDataDbScope } = await import("@/drizzle/db"); + + const controlBefore = getDb(); + const dataObservations = await runWithDataDbScope(async () => { + const immediate = getDb(); + await Promise.resolve(); + const afterAwait = getDb(); + const afterTimer = await new Promise>((resolve) => { + setTimeout(() => resolve(getDb()), 0); + }); + return { immediate, afterAwait, afterTimer }; + }); + const controlAfter = getDb(); + + expect(dataObservations.immediate).toBe(dataObservations.afterAwait); + expect(dataObservations.immediate).toBe(dataObservations.afterTimer); + expect(dataObservations.immediate).not.toBe(controlBefore); + expect(controlAfter).toBe(controlBefore); + }); + + it("并行 data scope 不会把调用方的默认 control scope 泄漏为 data", async () => { + const { getDb, runWithDataDbScope } = await import("@/drizzle/db"); + const controlDb = getDb(); + + let releaseDataScope: () => void = () => {}; + const dataGate = new Promise((resolve) => { + releaseDataScope = resolve; + }); + const dataTask = runWithDataDbScope(async () => { + const scopedDb = getDb(); + await dataGate; + return { scopedDb, afterGate: getDb() }; + }); + + expect(getDb()).toBe(controlDb); + releaseDataScope(); + const dataResult = await dataTask; + expect(dataResult.scopedDb).toBe(dataResult.afterGate); + expect(dataResult.scopedDb).not.toBe(controlDb); + expect(getDb()).toBe(controlDb); + }); + + it("route handler wrapper 在完整异步 handler 生命周期内保持 data scope", async () => { + const { getDb, withDataDbScope } = await import("@/drizzle/db"); + const controlDb = getDb(); + const handler = withDataDbScope(async (value: string) => { + await new Promise((resolve) => setTimeout(resolve, 0)); + return { value, db: getDb() }; + }); + + const result = await handler("route"); + + expect(result.value).toBe("route"); + expect(result.db).not.toBe(controlDb); + expect(getDb()).toBe(controlDb); + }); +}); diff --git a/tests/unit/drizzle/db-shutdown.test.ts b/tests/unit/drizzle/db-shutdown.test.ts new file mode 100644 index 000000000..ef07ac7ef --- /dev/null +++ b/tests/unit/drizzle/db-shutdown.test.ts @@ -0,0 +1,282 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("@/lib/logger", () => ({ + logger: { + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + }, +})); + +describe.sequential("数据库连接池 shutdown", () => { + beforeEach(() => { + vi.resetModules(); + vi.useRealTimers(); + }); + + it("message writer flush 完成后关闭所有数据库 pool", async () => { + const lifecycle: string[] = []; + const closeDbPools = vi.fn(async () => { + lifecycle.push("db"); + }); + + vi.doMock("@/lib/cache/session-cache", () => ({ stopCacheCleanup: () => {} })); + vi.doMock("@/lib/provider-endpoints/probe-scheduler", () => ({ + stopEndpointProbeScheduler: () => {}, + })); + vi.doMock("@/lib/public-status/scheduler", () => ({ + stopPublicStatusRebuildScheduler: async () => {}, + })); + vi.doMock("@/lib/provider-endpoints/probe-log-cleanup", () => ({ + stopEndpointProbeLogCleanup: () => {}, + })); + vi.doMock("@/lib/async-task-manager", () => ({ shutdownAllAsyncTasks: () => {} })); + vi.doMock("@/repository/message-write-buffer", () => ({ + stopMessageRequestWriteBuffer: async () => { + lifecycle.push("writer"); + }, + })); + vi.doMock("@/lib/langfuse", () => ({ shutdownLangfuse: async () => {} })); + vi.doMock("@/lib/redis", () => ({ closeRedis: async () => {} })); + vi.doMock("@/drizzle/db", () => ({ closeDbPools })); + + const { runApplicationCleanup } = await import("@/lib/lifecycle/shutdown"); + await runApplicationCleanup("SIGTERM", { totalTimeoutMs: 5_000, perStepTimeoutMs: 500 }); + + expect(closeDbPools).toHaveBeenCalledTimes(1); + expect(lifecycle).toEqual(["writer", "db"]); + }); + + it("writer 前整体 deadline 已到也不能提前完成 cleanup", async () => { + vi.useFakeTimers(); + let resolveTasks!: () => void; + const tasksSettled = new Promise((resolve) => { + resolveTasks = resolve; + }); + const stopMessageRequestWriteBuffer = vi.fn(async () => {}); + const closeDbPools = vi.fn(async () => {}); + + vi.doMock("@/lib/cache/session-cache", () => ({ stopCacheCleanup: () => {} })); + vi.doMock("@/lib/provider-endpoints/probe-scheduler", () => ({ + stopEndpointProbeScheduler: () => {}, + })); + vi.doMock("@/lib/public-status/scheduler", () => ({ + stopPublicStatusRebuildScheduler: async () => {}, + })); + vi.doMock("@/lib/provider-endpoints/probe-log-cleanup", () => ({ + stopEndpointProbeLogCleanup: () => {}, + })); + vi.doMock("@/lib/async-task-manager", () => ({ + shutdownAllAsyncTasks: () => tasksSettled, + })); + vi.doMock("@/repository/message-write-buffer", () => ({ + stopMessageRequestWriteBuffer, + })); + vi.doMock("@/lib/langfuse", () => ({ shutdownLangfuse: async () => {} })); + vi.doMock("@/lib/redis", () => ({ closeRedis: async () => {} })); + vi.doMock("@/drizzle/db", () => ({ closeDbPools })); + + const { runApplicationCleanup } = await import("@/lib/lifecycle/shutdown"); + let cleanupSettled = false; + const cleanup = runApplicationCleanup("SIGTERM", { + totalTimeoutMs: 100, + perStepTimeoutMs: 1_000, + }); + void cleanup.then( + () => { + cleanupSettled = true; + }, + () => { + cleanupSettled = true; + } + ); + + await vi.advanceTimersByTimeAsync(100); + + expect(cleanupSettled).toBe(false); + expect(stopMessageRequestWriteBuffer).not.toHaveBeenCalled(); + + resolveTasks(); + await vi.advanceTimersByTimeAsync(0); + await cleanup; + + expect(stopMessageRequestWriteBuffer).toHaveBeenCalledTimes(1); + expect(closeDbPools).toHaveBeenCalledTimes(1); + }); + + it("message writer 超过单步 timeout 时继续等待,settled 后才关闭数据库 pool", async () => { + vi.useFakeTimers(); + const lifecycle: string[] = []; + let resolveWriter!: () => void; + const writerStopped = new Promise((resolve) => { + resolveWriter = () => { + lifecycle.push("writer"); + resolve(); + }; + }); + const closeDbPools = vi.fn(async () => { + lifecycle.push("db"); + }); + + vi.doMock("@/lib/cache/session-cache", () => ({ stopCacheCleanup: () => {} })); + vi.doMock("@/lib/provider-endpoints/probe-scheduler", () => ({ + stopEndpointProbeScheduler: () => {}, + })); + vi.doMock("@/lib/public-status/scheduler", () => ({ + stopPublicStatusRebuildScheduler: async () => {}, + })); + vi.doMock("@/lib/provider-endpoints/probe-log-cleanup", () => ({ + stopEndpointProbeLogCleanup: () => {}, + })); + vi.doMock("@/lib/async-task-manager", () => ({ shutdownAllAsyncTasks: () => {} })); + vi.doMock("@/repository/message-write-buffer", () => ({ + stopMessageRequestWriteBuffer: () => writerStopped, + })); + vi.doMock("@/lib/langfuse", () => ({ shutdownLangfuse: async () => {} })); + vi.doMock("@/lib/redis", () => ({ closeRedis: async () => {} })); + vi.doMock("@/drizzle/db", () => ({ closeDbPools })); + + const { runApplicationCleanup } = await import("@/lib/lifecycle/shutdown"); + let cleanupSettled = false; + const cleanup = runApplicationCleanup("SIGTERM", { + totalTimeoutMs: 1_000, + perStepTimeoutMs: 100, + }); + void cleanup.then( + () => { + cleanupSettled = true; + }, + () => { + cleanupSettled = true; + } + ); + + await vi.advanceTimersByTimeAsync(100); + + expect(cleanupSettled).toBe(false); + expect(closeDbPools).not.toHaveBeenCalled(); + + resolveWriter(); + await vi.advanceTimersByTimeAsync(0); + await cleanup; + + expect(closeDbPools).toHaveBeenCalledTimes(1); + expect(lifecycle).toEqual(["writer", "db"]); + }); + + it("message writer 到整体 deadline 仍未 settled 时 cleanup 保持 pending,settled 后才关闭 pool", async () => { + vi.useFakeTimers(); + const lifecycle: string[] = []; + let resolveWriter!: () => void; + const writerStopped = new Promise((resolve) => { + resolveWriter = () => { + lifecycle.push("writer"); + resolve(); + }; + }); + const closeDbPools = vi.fn(async () => { + lifecycle.push("db"); + }); + + vi.doMock("@/lib/cache/session-cache", () => ({ stopCacheCleanup: () => {} })); + vi.doMock("@/lib/provider-endpoints/probe-scheduler", () => ({ + stopEndpointProbeScheduler: () => {}, + })); + vi.doMock("@/lib/public-status/scheduler", () => ({ + stopPublicStatusRebuildScheduler: async () => {}, + })); + vi.doMock("@/lib/provider-endpoints/probe-log-cleanup", () => ({ + stopEndpointProbeLogCleanup: () => {}, + })); + vi.doMock("@/lib/async-task-manager", () => ({ shutdownAllAsyncTasks: () => {} })); + vi.doMock("@/repository/message-write-buffer", () => ({ + stopMessageRequestWriteBuffer: () => writerStopped, + })); + vi.doMock("@/lib/langfuse", () => ({ shutdownLangfuse: async () => {} })); + vi.doMock("@/lib/redis", () => ({ closeRedis: async () => {} })); + vi.doMock("@/drizzle/db", () => ({ closeDbPools })); + + const { runApplicationCleanup } = await import("@/lib/lifecycle/shutdown"); + const cleanup = runApplicationCleanup("SIGTERM", { + totalTimeoutMs: 500, + perStepTimeoutMs: 100, + }); + let cleanupSettled = false; + void cleanup.then( + () => { + cleanupSettled = true; + }, + () => { + cleanupSettled = true; + } + ); + + await vi.advanceTimersByTimeAsync(500); + + expect(cleanupSettled).toBe(false); + expect(closeDbPools).not.toHaveBeenCalled(); + + resolveWriter(); + await vi.advanceTimersByTimeAsync(0); + await cleanup; + + expect(closeDbPools).toHaveBeenCalledTimes(1); + expect(lifecycle).toEqual(["writer", "db"]); + }); + + it("数据库 pool close 超过单步 timeout 时不得 detach", async () => { + vi.useFakeTimers(); + let resolveClose!: () => void; + const poolClosed = new Promise((resolve) => { + resolveClose = resolve; + }); + const closeDbPools = vi.fn(() => poolClosed); + const shutdownLangfuse = vi.fn(async () => {}); + + vi.doMock("@/lib/cache/session-cache", () => ({ stopCacheCleanup: () => {} })); + vi.doMock("@/lib/provider-endpoints/probe-scheduler", () => ({ + stopEndpointProbeScheduler: () => {}, + })); + vi.doMock("@/lib/public-status/scheduler", () => ({ + stopPublicStatusRebuildScheduler: async () => {}, + })); + vi.doMock("@/lib/provider-endpoints/probe-log-cleanup", () => ({ + stopEndpointProbeLogCleanup: () => {}, + })); + vi.doMock("@/lib/async-task-manager", () => ({ shutdownAllAsyncTasks: async () => {} })); + vi.doMock("@/repository/message-write-buffer", () => ({ + stopMessageRequestWriteBuffer: async () => {}, + })); + vi.doMock("@/lib/langfuse", () => ({ shutdownLangfuse })); + vi.doMock("@/lib/redis", () => ({ closeRedis: async () => {} })); + vi.doMock("@/drizzle/db", () => ({ closeDbPools })); + + const { runApplicationCleanup } = await import("@/lib/lifecycle/shutdown"); + let cleanupSettled = false; + const cleanup = runApplicationCleanup("SIGTERM", { + totalTimeoutMs: 500, + perStepTimeoutMs: 100, + }); + void cleanup.then( + () => { + cleanupSettled = true; + }, + () => { + cleanupSettled = true; + } + ); + + await vi.advanceTimersByTimeAsync(100); + + expect(closeDbPools).toHaveBeenCalledTimes(1); + expect(cleanupSettled).toBe(false); + expect(shutdownLangfuse).not.toHaveBeenCalled(); + + resolveClose(); + await vi.advanceTimersByTimeAsync(0); + await cleanup; + + expect(shutdownLangfuse).toHaveBeenCalledTimes(1); + }); +}); diff --git a/tests/unit/i18n/key-created-copy.test.ts b/tests/unit/i18n/key-created-copy.test.ts index 5bcdaf7c0..1ba546dfb 100644 --- a/tests/unit/i18n/key-created-copy.test.ts +++ b/tests/unit/i18n/key-created-copy.test.ts @@ -65,31 +65,32 @@ function getString(messages: Record, keyPath: readonly string[] describe.each(LOCALES)("key creation copy (%s)", (locale) => { const dashboard = loadMessages(locale, "dashboard.json"); - test.each( - COPY_PATHS.map((p) => [p.join("."), p] as const) - )("%s matches the actual reveal behavior", (_label, keyPath) => { - const copy = getString(dashboard, keyPath); + test.each(COPY_PATHS.map((p) => [p.join("."), p] as const))( + "%s matches the actual reveal behavior", + (_label, keyPath) => { + const copy = getString(dashboard, keyPath); - expect(copy.trim().length).toBeGreaterThan(0); - for (const pattern of ONE_TIME_CLAIM_PATTERNS) { - expect(copy).not.toMatch(pattern); + expect(copy.trim().length).toBeGreaterThan(0); + for (const pattern of ONE_TIME_CLAIM_PATTERNS) { + expect(copy).not.toMatch(pattern); + } + expect(copy).toMatch(REVIEWABLE_MARKERS[locale]); } - expect(copy).toMatch(REVIEWABLE_MARKERS[locale]); - }); + ); }); describe.each(LOCALES)("removeKey error code translations (%s)", (locale) => { const errors = loadMessages(locale, "errors.json"); - test.each([ - "CANNOT_DELETE_LAST_KEY", - "CANNOT_DELETE_LAST_GROUP_KEY", - ])("errors namespace translates %s", (code) => { - const value = errors[code]; - expect(value, `${locale}/errors.json must define ${code}`).toBeTypeOf("string"); - expect((value as string).trim().length).toBeGreaterThan(0); - // Must be a distinct, specific message rather than a copy of a generic one. - expect(value).not.toBe(errors.OPERATION_FAILED); - expect(value).not.toBe(errors.DELETE_FAILED); - }); + test.each(["CANNOT_DELETE_LAST_KEY", "CANNOT_DELETE_LAST_GROUP_KEY"])( + "errors namespace translates %s", + (code) => { + const value = errors[code]; + expect(value, `${locale}/errors.json must define ${code}`).toBeTypeOf("string"); + expect((value as string).trim().length).toBeGreaterThan(0); + // Must be a distinct, specific message rather than a copy of a generic one. + expect(value).not.toBe(errors.OPERATION_FAILED); + expect(value).not.toBe(errors.DELETE_FAILED); + } + ); }); diff --git a/tests/unit/instrumentation-crash-handler.test.ts b/tests/unit/instrumentation-crash-handler.test.ts index bda70367c..dc3016174 100644 --- a/tests/unit/instrumentation-crash-handler.test.ts +++ b/tests/unit/instrumentation-crash-handler.test.ts @@ -9,6 +9,7 @@ * 防止未来重构(删掉早返回、反转判断、移动谓词调用)在谓词测试全绿的情况下重新引入崩溃。 */ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { DrizzleQueryError } from "drizzle-orm"; vi.mock("@/lib/logger", () => ({ logger: { @@ -121,37 +122,67 @@ describe("registerCrashDiagnostics", () => { describe("genuine errors (must still fail-fast)", () => { it("uncaughtException: a generic Error exits with code 1 and writes diagnostics", () => { const { uncaughtException } = captureHandlers(); - uncaughtException(new Error("real bug")); + const error = new Error("real bug"); + uncaughtException(error); expect(exitSpy).toHaveBeenCalledWith(1); expect(logger.fatal).toHaveBeenCalledTimes(1); expect(logger.warn).not.toHaveBeenCalled(); // fatal 路径必须写出同步 stderr 兜底诊断,防止回归静默吞掉致命错误 expect(stderrSpy).toHaveBeenCalled(); + expect(process.report?.excludeEnv).toBe(true); + const reportCalls = (process.report!.writeReport as unknown as ReturnType).mock + .calls; + expect(reportCalls[0]?.[1]).toBe(error); }); - it("uncaughtException: a non-benign transport code (ECONNREFUSED) exits with code 1", () => { + it("uncaughtException: database wrappers are redacted at every crash sink", () => { const { uncaughtException } = captureHandlers(); - uncaughtException(makeError("ECONNREFUSED", "connect ECONNREFUSED")); - + const error = new DrizzleQueryError( + "select * from secrets where token = $1", + ["crash-report-canary"], + new Error("driver exposed crash-report-canary") + ); + + uncaughtException(error); + + const reportCalls = (process.report!.writeReport as unknown as ReturnType).mock + .calls; + const reportError = reportCalls[0]?.[1] as Error; + expect(reportError).not.toBe(error); + expect(reportError.message).toBe("Database query failed"); + expect( + JSON.stringify({ + reportCalls, + stderrCalls: stderrSpy.mock.calls, + fatalCalls: (logger.fatal as unknown as ReturnType).mock.calls, + }) + ).not.toContain("crash-report-canary"); expect(exitSpy).toHaveBeenCalledWith(1); - expect(logger.fatal).toHaveBeenCalledTimes(1); }); - it.each([ - "ECONNRESET", - "ERR_STREAM_PREMATURE_CLOSE", - ])("uncaughtException: ambiguous code %s is NOT suppressed and still exits with code 1", (code) => { - // 这些码方向不明(可能来自上游 DB/Redis/provider),进程级无上下文区分, - // 必须保持 fail-fast,避免误吞真正的基础设施故障。 + it("uncaughtException: a non-benign transport code (ECONNREFUSED) exits with code 1", () => { const { uncaughtException } = captureHandlers(); - uncaughtException(makeError(code)); + uncaughtException(makeError("ECONNREFUSED", "connect ECONNREFUSED")); expect(exitSpy).toHaveBeenCalledWith(1); expect(logger.fatal).toHaveBeenCalledTimes(1); - expect(logger.warn).not.toHaveBeenCalled(); }); + it.each(["ECONNRESET", "ERR_STREAM_PREMATURE_CLOSE"])( + "uncaughtException: ambiguous code %s is NOT suppressed and still exits with code 1", + (code) => { + // 这些码方向不明(可能来自上游 DB/Redis/provider),进程级无上下文区分, + // 必须保持 fail-fast,避免误吞真正的基础设施故障。 + const { uncaughtException } = captureHandlers(); + uncaughtException(makeError(code)); + + expect(exitSpy).toHaveBeenCalledWith(1); + expect(logger.fatal).toHaveBeenCalledTimes(1); + expect(logger.warn).not.toHaveBeenCalled(); + } + ); + it("unhandledRejection: a generic rejection exits with code 1", () => { const { unhandledRejection } = captureHandlers(); unhandledRejection(new Error("real rejection")); diff --git a/tests/unit/langfuse/emit-proxy-trace.test.ts b/tests/unit/langfuse/emit-proxy-trace.test.ts new file mode 100644 index 000000000..9fbcd5dc9 --- /dev/null +++ b/tests/unit/langfuse/emit-proxy-trace.test.ts @@ -0,0 +1,43 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import type { ProxySession } from "@/app/v1/_lib/proxy/session"; + +vi.mock("@/lib/logger", () => ({ + logger: { + warn: vi.fn(), + }, +})); + +import { emitProxyLangfuseTrace } from "@/lib/langfuse/emit-proxy-trace"; +import { logger } from "@/lib/logger"; + +describe("emitProxyLangfuseTrace", () => { + afterEach(() => { + vi.unstubAllEnvs(); + vi.clearAllMocks(); + }); + + it("never lets a synchronous session snapshot failure escape", () => { + vi.stubEnv("LANGFUSE_PUBLIC_KEY", "test-public"); + vi.stubEnv("LANGFUSE_SECRET_KEY", "test-secret"); + const session = { + getProviderChain() { + throw new Error("snapshot exploded"); + }, + } as unknown as ProxySession; + + expect(() => + emitProxyLangfuseTrace(session, { + responseHeaders: new Headers(), + responseText: "", + usageMetrics: null, + costUsd: undefined, + statusCode: 500, + durationMs: 1, + isStreaming: false, + }) + ).not.toThrow(); + expect(logger.warn).toHaveBeenCalledWith("[Langfuse] Proxy trace snapshot failed", { + error: "snapshot exploded", + }); + }); +}); diff --git a/tests/unit/langfuse/langfuse-trace.test.ts b/tests/unit/langfuse/langfuse-trace.test.ts index 91bd63899..92dcea2f8 100644 --- a/tests/unit/langfuse/langfuse-trace.test.ts +++ b/tests/unit/langfuse/langfuse-trace.test.ts @@ -234,12 +234,30 @@ describe("traceProxyRequest", () => { expect(llmCall[1].output).toEqual(responseBody); }); - test("should pass raw headers without redaction", async () => { + test("redacts credential headers and preserves benign headers at the Langfuse boundary", async () => { const { traceProxyRequest } = await import("@/lib/langfuse/trace-proxy-request"); + const authorizationSecret = "Bearer request-authorization-secret"; + const apiKeySecret = "request-api-key-secret"; + const cookieSecret = "request-cookie-secret"; + const setCookieSecret = "response-set-cookie-secret"; await traceProxyRequest({ - session: createMockSession(), - responseHeaders: new Headers({ "x-api-key": "secret-mock" }), + session: createMockSession({ + headers: new Headers({ + authorization: authorizationSecret, + cookie: cookieSecret, + "content-type": "application/json", + "x-api-key": apiKeySecret, + "x-cch-future-marker": "future-internal-canary", + "x-cch-responses-ws-session": "ws-session-canary", + "x-request-id": "request-123", + }), + }), + responseHeaders: new Headers({ + "content-type": "text/event-stream", + "set-cookie": setCookieSecret, + "x-response-id": "response-456", + }), durationMs: 500, statusCode: 200, isStreaming: false, @@ -249,9 +267,36 @@ describe("traceProxyRequest", () => { (c: unknown[]) => c[0] === "llm-call" ); const metadata = llmCall[1].metadata; - expect(metadata.requestHeaders["x-api-key"]).toBe("test-mock-key-not-real"); - expect(metadata.requestHeaders["content-type"]).toBe("application/json"); - expect(metadata.responseHeaders["x-api-key"]).toBe("secret-mock"); + expect(metadata.requestHeaders).toEqual({ + authorization: "[REDACTED]", + cookie: "[REDACTED]", + "content-type": "application/json", + "x-api-key": "[REDACTED]", + "x-request-id": "request-123", + }); + expect(metadata.responseHeaders).toEqual({ + "content-type": "text/event-stream", + "set-cookie": "[REDACTED]", + "x-response-id": "response-456", + }); + + const serializedSdkArguments = JSON.stringify({ + rootObservation: mockStartObservation.mock.calls, + propagatedAttributes: mockPropagateAttributes.mock.calls, + childObservations: mockRootSpan.startObservation.mock.calls, + generationUpdates: mockGenerationUpdate.mock.calls, + generationEnds: mockGenerationEnd.mock.calls, + guardEnds: mockGuardSpanEnd.mock.calls, + eventEnds: mockEventEnd.mock.calls, + traceIo: mockSetTraceIO.mock.calls, + rootEnds: mockSpanEnd.mock.calls, + }); + for (const secret of [authorizationSecret, apiKeySecret, cookieSecret, setCookieSecret]) { + expect(serializedSdkArguments).not.toContain(secret); + } + expect(serializedSdkArguments).not.toContain("x-cch-"); + expect(serializedSdkArguments).not.toContain("future-internal-canary"); + expect(serializedSdkArguments).not.toContain("ws-session-canary"); }); test("should include provider name and model in tags", async () => { diff --git a/tests/unit/lib/async-task-manager-edge-runtime.test.ts b/tests/unit/lib/async-task-manager-edge-runtime.test.ts index dbc21f36f..f5db6451a 100644 --- a/tests/unit/lib/async-task-manager-edge-runtime.test.ts +++ b/tests/unit/lib/async-task-manager-edge-runtime.test.ts @@ -47,7 +47,7 @@ describe.sequential("AsyncTaskManager edge runtime", () => { process.env.NEXT_RUNTIME = "edge"; const { AsyncTaskManager } = await import("@/lib/async-task-manager"); - AsyncTaskManager.register("t1", Promise.resolve()); + AsyncTaskManager.register("t1", async () => {}); expect(processOnceSpy).not.toHaveBeenCalled(); }); @@ -59,7 +59,7 @@ describe.sequential("AsyncTaskManager edge runtime", () => { process.env.NEXT_RUNTIME = "nodejs"; const { AsyncTaskManager } = await import("@/lib/async-task-manager"); - AsyncTaskManager.register("t1", Promise.resolve()); + AsyncTaskManager.register("t1", async () => {}); const signals = processOnceSpy.mock.calls.map((c) => c[0]); expect(signals).toContain("beforeExit"); @@ -80,7 +80,7 @@ describe.sequential("AsyncTaskManager edge runtime", () => { const taskPromise = new Promise((resolve) => { resolveTask = resolve; }); - const controller = AsyncTaskManager.register("t1", taskPromise); + const controller = AsyncTaskManager.register("t1", () => taskPromise); const beforeExitHandler = processOnceSpy.mock.calls.find((c) => c[0] === "beforeExit")?.[1]; expect(beforeExitHandler).toBeTypeOf("function"); @@ -104,7 +104,7 @@ describe.sequential("AsyncTaskManager edge runtime", () => { const taskPromise = new Promise((resolve) => { resolveTask = resolve; }); - const controller = AsyncTaskManager.register("t1", taskPromise); + const controller = AsyncTaskManager.register("t1", () => taskPromise); expect(controller.signal.aborted).toBe(false); shutdownAllAsyncTasks(); @@ -115,6 +115,112 @@ describe.sequential("AsyncTaskManager edge runtime", () => { await taskPromise; }); + it("does not start tasks registered after shutdown observes an empty task snapshot", async () => { + process.env.CI = "true"; + process.env.NEXT_RUNTIME = "nodejs"; + + const { AsyncTaskManager, shutdownAllAsyncTasks } = await import("@/lib/async-task-manager"); + + const shutdownPromise = shutdownAllAsyncTasks(); + let taskStarted = false; + const lateController = AsyncTaskManager.register("late-task", async () => { + taskStarted = true; + await new Promise(() => {}); + }); + const lateRegistrationState = { + aborted: lateController.signal.aborted, + active: AsyncTaskManager.getActiveTaskCount(), + taskStarted, + }; + const repeatedShutdownPromise = shutdownAllAsyncTasks(); + + expect({ + samePromise: repeatedShutdownPromise === shutdownPromise, + ...lateRegistrationState, + }).toEqual({ + samePromise: true, + aborted: true, + active: 0, + taskStarted: false, + }); + }); + + it("joins a tail task registered synchronously by an abort listener", async () => { + process.env.CI = "true"; + process.env.NEXT_RUNTIME = "nodejs"; + + const { AsyncTaskManager, shutdownAllAsyncTasks } = await import("@/lib/async-task-manager"); + const firstController = new AbortController(); + let resolveFirst!: () => void; + const firstTask = new Promise((resolve) => { + resolveFirst = resolve; + }); + let tailStarted = false; + let tailAborted = false; + + firstController.signal.addEventListener( + "abort", + () => { + AsyncTaskManager.register("tail-task", async (signal) => { + tailStarted = true; + await new Promise((resolve) => { + signal.addEventListener( + "abort", + () => { + tailAborted = true; + resolve(); + }, + { once: true } + ); + }); + }); + resolveFirst(); + }, + { once: true } + ); + + AsyncTaskManager.register("first-task", () => firstTask, { + abortController: firstController, + }); + + await shutdownAllAsyncTasks(); + + expect(tailStarted).toBe(true); + expect(tailAborted).toBe(true); + expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + }); + + it("publishes task admission before a factory synchronously reenters shutdown", async () => { + process.env.CI = "true"; + process.env.NEXT_RUNTIME = "nodejs"; + + const { AsyncTaskManager, shutdownAllAsyncTasks } = await import("@/lib/async-task-manager"); + let resolveTask!: () => void; + let reentrantShutdown: Promise | undefined; + + const controller = AsyncTaskManager.register("reentrant-task", async () => { + reentrantShutdown = shutdownAllAsyncTasks(); + await new Promise((resolve) => { + resolveTask = resolve; + }); + }); + const repeatedShutdown = shutdownAllAsyncTasks(); + let shutdownSettled = false; + repeatedShutdown.then(() => { + shutdownSettled = true; + }); + + expect(reentrantShutdown).toBe(repeatedShutdown); + expect(controller.signal.aborted).toBe(true); + await new Promise((resolve) => queueMicrotask(() => resolve())); + expect(shutdownSettled).toBe(false); + + resolveTask(); + await repeatedShutdown; + + expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + }); + it("runs cleanupCompletedTasks on interval tick", async () => { vi.useFakeTimers(); process.env.NEXT_RUNTIME = "nodejs"; @@ -125,7 +231,7 @@ describe.sequential("AsyncTaskManager edge runtime", () => { "cleanupCompletedTasks" ); - AsyncTaskManager.register("t1", new Promise(() => {})); + AsyncTaskManager.register("t1", async () => new Promise(() => {})); vi.advanceTimersByTime(60_000); expect(cleanupSpy).toHaveBeenCalledTimes(1); @@ -142,15 +248,15 @@ describe.sequential("AsyncTaskManager edge runtime", () => { resolveTask = resolve; }); - const controller = AsyncTaskManager.register("t1", taskPromise); + const controller = AsyncTaskManager.register("t1", () => taskPromise); expect(controller.signal.aborted).toBe(false); expect(AsyncTaskManager.getActiveTaskCount()).toBe(1); resolveTask!(); await taskPromise; - await new Promise((resolve) => queueMicrotask(() => resolve())); - - expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + await vi.waitFor(() => { + expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + }); }); it("does nothing when cancelling unknown taskId", async () => { @@ -176,7 +282,7 @@ describe.sequential("AsyncTaskManager edge runtime", () => { resolveTask = resolve; }); - AsyncTaskManager.register("t1", taskPromise, "custom_type"); + AsyncTaskManager.register("t1", () => taskPromise, "custom_type"); const tasks = AsyncTaskManager.getActiveTasks(); expect(tasks).toHaveLength(1); @@ -198,7 +304,7 @@ describe.sequential("AsyncTaskManager edge runtime", () => { resolveFirst = resolve; }); - const firstController = AsyncTaskManager.register("t1", firstPromise); + const firstController = AsyncTaskManager.register("t1", () => firstPromise); expect(firstController.signal.aborted).toBe(false); let resolveSecond: () => void; @@ -206,7 +312,7 @@ describe.sequential("AsyncTaskManager edge runtime", () => { resolveSecond = resolve; }); - AsyncTaskManager.register("t1", secondPromise); + AsyncTaskManager.register("t1", () => secondPromise); expect(firstController.signal.aborted).toBe(true); @@ -215,6 +321,244 @@ describe.sequential("AsyncTaskManager edge runtime", () => { await Promise.all([firstPromise, secondPromise]); }); + it("does not start B when aborting same-ID A synchronously registers C", async () => { + process.env.CI = "true"; + process.env.NEXT_RUNTIME = "nodejs"; + + const { AsyncTaskManager, shutdownAllAsyncTasks } = await import("@/lib/async-task-manager"); + + let resolveA!: () => void; + let resolveC!: () => void; + let cStartCount = 0; + const aController = AsyncTaskManager.register( + "t1", + async (signal) => { + signal.addEventListener( + "abort", + () => { + AsyncTaskManager.register( + "t1", + async () => { + cStartCount += 1; + await new Promise((resolve) => { + resolveC = resolve; + }); + }, + "generation-c" + ); + }, + { once: true } + ); + + await new Promise((resolve) => { + resolveA = resolve; + }); + }, + "generation-a" + ); + + let bStartCount = 0; + const bController = AsyncTaskManager.register( + "t1", + async () => { + bStartCount += 1; + }, + "generation-b" + ); + + expect({ + bAborted: bController.signal.aborted, + bStartCount, + cStartCount, + }).toEqual({ + bAborted: true, + bStartCount: 0, + cStartCount: 1, + }); + + await vi.waitFor(() => { + expect(AsyncTaskManager.getActiveTasks().map(({ taskType }) => taskType)).toEqual([ + "generation-a", + "generation-c", + ]); + }); + + const shutdownPromise = shutdownAllAsyncTasks(); + let shutdownSettled = false; + shutdownPromise.then(() => { + shutdownSettled = true; + }); + await Promise.resolve(); + + expect(aController.signal.aborted).toBe(true); + expect(shutdownSettled).toBe(false); + expect(AsyncTaskManager.getActiveTaskCount()).toBe(2); + + resolveA(); + resolveC(); + await shutdownPromise; + + expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + }); + + it("keeps an aborted duplicate task generation joinable until its factory settles", async () => { + process.env.CI = "true"; + process.env.NEXT_RUNTIME = "nodejs"; + + const { AsyncTaskManager, shutdownAllAsyncTasks } = await import("@/lib/async-task-manager"); + + let resolveFirst!: () => void; + let firstAborted = false; + let firstFinalized = false; + const firstController = AsyncTaskManager.register("t1", async (signal) => { + signal.addEventListener("abort", () => { + firstAborted = true; + }); + + try { + await new Promise((resolve) => { + resolveFirst = resolve; + }); + } finally { + firstFinalized = true; + } + }); + + let resolveSecond!: () => void; + let secondAborted = false; + AsyncTaskManager.register("t1", async (signal) => { + signal.addEventListener("abort", () => { + secondAborted = true; + }); + + await new Promise((resolve) => { + resolveSecond = resolve; + }); + }); + + expect(firstController.signal.aborted).toBe(true); + expect(firstAborted).toBe(true); + expect(firstFinalized).toBe(false); + expect(AsyncTaskManager.getActiveTaskCount()).toBe(2); + + const shutdownPromise = shutdownAllAsyncTasks(); + let shutdownSettled = false; + shutdownPromise.then(() => { + shutdownSettled = true; + }); + + expect(secondAborted).toBe(true); + resolveSecond(); + await new Promise((resolve) => setImmediate(resolve)); + + expect(shutdownSettled).toBe(false); + expect(AsyncTaskManager.getActiveTaskCount()).toBe(1); + + resolveFirst(); + await shutdownPromise; + + expect(firstFinalized).toBe(true); + expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + }); + + it("retains the prior generation when a duplicate factory throws synchronously", async () => { + process.env.CI = "true"; + process.env.NEXT_RUNTIME = "nodejs"; + + const { AsyncTaskManager, shutdownAllAsyncTasks } = await import("@/lib/async-task-manager"); + + let resolveFirst!: () => void; + const firstController = AsyncTaskManager.register("t1", async () => { + await new Promise((resolve) => { + resolveFirst = resolve; + }); + }); + + AsyncTaskManager.register("t1", () => { + throw new Error("replacement failed before returning a promise"); + }); + + expect(firstController.signal.aborted).toBe(true); + await vi.waitFor(() => { + expect(AsyncTaskManager.getActiveTaskCount()).toBe(1); + }); + + const shutdownPromise = shutdownAllAsyncTasks(); + let shutdownSettled = false; + shutdownPromise.then(() => { + shutdownSettled = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + + expect(shutdownSettled).toBe(false); + + resolveFirst(); + await shutdownPromise; + + expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + }); + + it("joins a same-id generation registered by an abort listener during shutdown", async () => { + process.env.CI = "true"; + process.env.NEXT_RUNTIME = "nodejs"; + + const { AsyncTaskManager, shutdownAllAsyncTasks } = await import("@/lib/async-task-manager"); + + let resolveFirst!: () => void; + let resolveTail!: () => void; + let firstFinalized = false; + let tailStarted = false; + let tailAborted = false; + + AsyncTaskManager.register("t1", async (signal) => { + signal.addEventListener( + "abort", + () => { + AsyncTaskManager.register("t1", async (tailSignal) => { + tailStarted = true; + tailSignal.addEventListener("abort", () => { + tailAborted = true; + }); + await new Promise((resolve) => { + resolveTail = resolve; + }); + }); + resolveFirst(); + }, + { once: true } + ); + + try { + await new Promise((resolve) => { + resolveFirst = resolve; + }); + } finally { + firstFinalized = true; + } + }); + + const shutdownPromise = shutdownAllAsyncTasks(); + let shutdownSettled = false; + shutdownPromise.then(() => { + shutdownSettled = true; + }); + + expect(tailStarted).toBe(true); + await vi.waitFor(() => { + expect(tailAborted).toBe(true); + expect(firstFinalized).toBe(true); + }); + await new Promise((resolve) => setImmediate(resolve)); + + expect(shutdownSettled).toBe(false); + expect(AsyncTaskManager.getActiveTaskCount()).toBe(1); + + resolveTail(); + await shutdownPromise; + + expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + }); + it("does not let an old task finalizer remove a newer task with the same taskId", async () => { process.env.CI = "true"; process.env.NEXT_RUNTIME = "nodejs"; @@ -225,25 +569,25 @@ describe.sequential("AsyncTaskManager edge runtime", () => { const firstPromise = new Promise((resolve) => { resolveFirst = resolve; }); - AsyncTaskManager.register("t1", firstPromise); + AsyncTaskManager.register("t1", () => firstPromise); let resolveSecond: () => void; const secondPromise = new Promise((resolve) => { resolveSecond = resolve; }); - AsyncTaskManager.register("t1", secondPromise); + AsyncTaskManager.register("t1", () => secondPromise); resolveFirst!(); await firstPromise; - await new Promise((resolve) => queueMicrotask(() => resolve())); - - expect(AsyncTaskManager.getActiveTaskCount()).toBe(1); + await vi.waitFor(() => { + expect(AsyncTaskManager.getActiveTaskCount()).toBe(1); + }); resolveSecond!(); await secondPromise; - await new Promise((resolve) => queueMicrotask(() => resolve())); - - expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + await vi.waitFor(() => { + expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + }); }); it("logs task cancelled when isClientAbortError returns true", async () => { @@ -257,11 +601,13 @@ describe.sequential("AsyncTaskManager edge runtime", () => { const { AsyncTaskManager } = await import("@/lib/async-task-manager"); const taskPromise = Promise.reject(new Error("aborted")); - AsyncTaskManager.register("t1", taskPromise); + AsyncTaskManager.register("t1", () => taskPromise); await taskPromise.catch(() => {}); - expect(vi.mocked(logger.info)).toHaveBeenCalled(); + await vi.waitFor(() => { + expect(vi.mocked(logger.info)).toHaveBeenCalled(); + }); }); it("logs task failed when isClientAbortError returns false", async () => { @@ -275,11 +621,60 @@ describe.sequential("AsyncTaskManager edge runtime", () => { const { AsyncTaskManager } = await import("@/lib/async-task-manager"); const taskPromise = Promise.reject(new Error("boom")); - AsyncTaskManager.register("t1", taskPromise); + AsyncTaskManager.register("t1", () => taskPromise); await taskPromise.catch(() => {}); - expect(vi.mocked(logger.error)).toHaveBeenCalled(); + await vi.waitFor(() => { + expect(vi.mocked(logger.error)).toHaveBeenCalled(); + }); + }); + + it("keeps a stale task tracked until its promise settles during shutdown", async () => { + vi.useFakeTimers(); + process.env.NEXT_RUNTIME = "nodejs"; + + const { AsyncTaskManager, shutdownAllAsyncTasks } = await import("@/lib/async-task-manager"); + + let resolveTask!: () => void; + let taskSettled = false; + const taskPromise = new Promise((resolve) => { + resolveTask = resolve; + }).finally(() => { + taskSettled = true; + }); + const controller = AsyncTaskManager.register("stale-task", () => taskPromise, { + staleTimeoutMs: 1, + }); + + vi.advanceTimersByTime(60_000); + + const shutdownPromise = shutdownAllAsyncTasks(); + let shutdownSettled = false; + shutdownPromise.then(() => { + shutdownSettled = true; + }); + await Promise.resolve(); + await Promise.resolve(); + + const stateBeforeSettlement = { + aborted: controller.signal.aborted, + active: AsyncTaskManager.getActiveTaskCount(), + shutdownSettled, + taskSettled, + }; + + resolveTask(); + await taskPromise; + await shutdownPromise; + + expect(stateBeforeSettlement).toEqual({ + aborted: true, + active: 1, + shutdownSettled: false, + taskSettled: false, + }); + expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); }); it("cleanupCompletedTasks cancels stale tasks", async () => { @@ -294,7 +689,7 @@ describe.sequential("AsyncTaskManager edge runtime", () => { resolveTask = resolve; }); - const controller = AsyncTaskManager.register("stale-task", taskPromise, "custom_type"); + const controller = AsyncTaskManager.register("stale-task", () => taskPromise, "custom_type"); const managerAny = AsyncTaskManager as unknown as { tasks: Map; @@ -310,18 +705,25 @@ describe.sequential("AsyncTaskManager edge runtime", () => { const freshPromise = new Promise((resolve) => { resolveFresh = resolve; }); - const freshController = AsyncTaskManager.register("fresh-task", freshPromise, "custom_type"); + const freshController = AsyncTaskManager.register( + "fresh-task", + () => freshPromise, + "custom_type" + ); managerAny.cleanupCompletedTasks(); expect(controller.signal.aborted).toBe(true); expect(freshController.signal.aborted).toBe(false); - expect(AsyncTaskManager.getActiveTaskCount()).toBe(1); + expect(AsyncTaskManager.getActiveTaskCount()).toBe(2); expect(vi.mocked(logger.warn)).toHaveBeenCalled(); resolveTask!(); resolveFresh!(); await Promise.all([taskPromise, freshPromise]); + await vi.waitFor(() => { + expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + }); }); it("does not cancel a long-running task that was recently touched", async () => { @@ -335,7 +737,11 @@ describe.sequential("AsyncTaskManager edge runtime", () => { resolveTask = resolve; }); - const controller = AsyncTaskManager.register("active-stream", taskPromise, "stream-processing"); + const controller = AsyncTaskManager.register( + "active-stream", + () => taskPromise, + "stream-processing" + ); const managerAny = AsyncTaskManager as unknown as { tasks: Map; @@ -358,7 +764,7 @@ describe.sequential("AsyncTaskManager edge runtime", () => { await taskPromise; }); - it("cleanupCompletedTasks aborts a provided controller and detaches stale tasks", async () => { + it("cleanupCompletedTasks retains a stale task until its provided promise settles", async () => { process.env.CI = "true"; process.env.NEXT_RUNTIME = "nodejs"; @@ -370,7 +776,7 @@ describe.sequential("AsyncTaskManager edge runtime", () => { }); const controller = new AbortController(); - const returnedController = AsyncTaskManager.register("stale-task", taskPromise, { + const returnedController = AsyncTaskManager.register("stale-task", () => taskPromise, { taskType: "stream-processing", abortController: controller, }); @@ -389,10 +795,13 @@ describe.sequential("AsyncTaskManager edge runtime", () => { managerAny.cleanupCompletedTasks(); expect(controller.signal.aborted).toBe(true); - expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + expect(AsyncTaskManager.getActiveTaskCount()).toBe(1); resolveTask!(); await taskPromise; + await vi.waitFor(() => { + expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + }); }); it("cleanupAll cancels tasks and clears interval", async () => { @@ -405,7 +814,7 @@ describe.sequential("AsyncTaskManager edge runtime", () => { const taskPromise = new Promise((resolve) => { resolveTask = resolve; }); - const controller = AsyncTaskManager.register("t1", taskPromise); + const controller = AsyncTaskManager.register("t1", () => taskPromise); const clearIntervalSpy = vi.spyOn(globalThis, "clearInterval"); const intervalId = setInterval(() => {}, 1_000); @@ -418,12 +827,15 @@ describe.sequential("AsyncTaskManager edge runtime", () => { managerAny.cleanupAll(); expect(controller.signal.aborted).toBe(true); - expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + expect(AsyncTaskManager.getActiveTaskCount()).toBe(1); expect(clearIntervalSpy).toHaveBeenCalledWith(intervalId); expect(managerAny.cleanupInterval).toBeNull(); resolveTask!(); await taskPromise; + await vi.waitFor(() => { + expect(AsyncTaskManager.getActiveTaskCount()).toBe(0); + }); clearInterval(intervalId); }); }); diff --git a/tests/unit/lib/provider-allowed-model-schema.test.ts b/tests/unit/lib/provider-allowed-model-schema.test.ts index 9787daab8..44c2deddc 100644 --- a/tests/unit/lib/provider-allowed-model-schema.test.ts +++ b/tests/unit/lib/provider-allowed-model-schema.test.ts @@ -51,20 +51,17 @@ describe("provider-allowed-model-schema", () => { }); describe("regex 模式的 glob 通配符兼容", () => { - it.each<[string]>([ - ["*"], - ["*."], - ["claude-*"], - ["*-opus-*"], - ["?"], - ])("接受 glob 风格的 pattern: %s", (pattern) => { - const result = PROVIDER_ALLOWED_MODEL_RULE_SCHEMA.safeParse({ - matchType: "regex", - pattern, - }); + it.each<[string]>([["*"], ["*."], ["claude-*"], ["*-opus-*"], ["?"]])( + "接受 glob 风格的 pattern: %s", + (pattern) => { + const result = PROVIDER_ALLOWED_MODEL_RULE_SCHEMA.safeParse({ + matchType: "regex", + pattern, + }); - expect(result.success).toBe(true); - }); + expect(result.success).toBe(true); + } + ); it("仍然拒绝纯粹无法解析的正则", () => { const result = PROVIDER_ALLOWED_MODEL_RULE_SCHEMA.safeParse({ diff --git a/tests/unit/lib/provider-model-redirect-schema.test.ts b/tests/unit/lib/provider-model-redirect-schema.test.ts index a29461437..3766b181f 100644 --- a/tests/unit/lib/provider-model-redirect-schema.test.ts +++ b/tests/unit/lib/provider-model-redirect-schema.test.ts @@ -53,21 +53,18 @@ describe("provider-model-redirect-schema", () => { }); describe("regex 模式的 glob 通配符兼容", () => { - it.each<[string]>([ - ["*"], - ["*."], - ["claude-*"], - ["*-opus-*"], - ["?"], - ])("接受 glob 风格的 source: %s", (source) => { - const result = PROVIDER_MODEL_REDIRECT_RULE_SCHEMA.safeParse({ - matchType: "regex", - source, - target: "claude-sonnet-4-6", - }); + it.each<[string]>([["*"], ["*."], ["claude-*"], ["*-opus-*"], ["?"]])( + "接受 glob 风格的 source: %s", + (source) => { + const result = PROVIDER_MODEL_REDIRECT_RULE_SCHEMA.safeParse({ + matchType: "regex", + source, + target: "claude-sonnet-4-6", + }); - expect(result.success).toBe(true); - }); + expect(result.success).toBe(true); + } + ); it("仍然拒绝纯粹无法解析的正则", () => { const result = PROVIDER_MODEL_REDIRECT_RULE_SCHEMA.safeParse({ diff --git a/tests/unit/lib/rate-limit/lease-service.test.ts b/tests/unit/lib/rate-limit/lease-service.test.ts index d0c0ef9c7..d6af232e1 100644 --- a/tests/unit/lib/rate-limit/lease-service.test.ts +++ b/tests/unit/lib/rate-limit/lease-service.test.ts @@ -63,6 +63,7 @@ describe("LeaseService", () => { vi.useFakeTimers(); vi.setSystemTime(new Date(nowMs)); vi.clearAllMocks(); + mockRedis.status = "ready"; }); afterEach(() => { @@ -955,6 +956,220 @@ describe("LeaseService", () => { }); }); + describe("settleLeaseBudgets", () => { + const settlementParams = { + requestId: 9001, + cost: 1.25, + entities: { + key: { + id: 101, + resetModes: { "5h": "rolling", daily: "fixed" }, + }, + user: { + id: 202, + resetModes: { "5h": "fixed", daily: "rolling" }, + }, + provider: { + id: 303, + resetModes: { "5h": "rolling", daily: "fixed" }, + }, + }, + } as const; + + const encodedSettlements = JSON.stringify([ + [1, 8.75], + [0, -1], + [-1, 0.5], + [1, 18.75], + [1, 28.75], + [1, 38.75], + [1, 48.75], + [1, 58.75], + [1, 68.75], + [1, 78.75], + [1, 88.75], + [1, 98.75], + ]); + + it("settles all four windows for key, user, and provider in one bounded eval", async () => { + const { LeaseService } = await import("@/lib/rate-limit/lease-service"); + const { buildLeaseKey } = await import("@/lib/rate-limit/lease"); + + mockRedis.eval.mockResolvedValue([0, encodedSettlements]); + + const result = await LeaseService.settleLeaseBudgets(settlementParams); + + expect(result).toEqual({ + requestId: "9001", + status: "settled", + settlements: [ + { + entityType: "key", + entityId: 101, + window: "5h", + status: "decremented", + newRemaining: 8.75, + }, + { + entityType: "key", + entityId: 101, + window: "daily", + status: "missing", + newRemaining: -1, + }, + { + entityType: "key", + entityId: 101, + window: "weekly", + status: "insufficient", + newRemaining: 0.5, + }, + { + entityType: "key", + entityId: 101, + window: "monthly", + status: "decremented", + newRemaining: 18.75, + }, + { + entityType: "user", + entityId: 202, + window: "5h", + status: "decremented", + newRemaining: 28.75, + }, + { + entityType: "user", + entityId: 202, + window: "daily", + status: "decremented", + newRemaining: 38.75, + }, + { + entityType: "user", + entityId: 202, + window: "weekly", + status: "decremented", + newRemaining: 48.75, + }, + { + entityType: "user", + entityId: 202, + window: "monthly", + status: "decremented", + newRemaining: 58.75, + }, + { + entityType: "provider", + entityId: 303, + window: "5h", + status: "decremented", + newRemaining: 68.75, + }, + { + entityType: "provider", + entityId: 303, + window: "daily", + status: "decremented", + newRemaining: 78.75, + }, + { + entityType: "provider", + entityId: 303, + window: "weekly", + status: "decremented", + newRemaining: 88.75, + }, + { + entityType: "provider", + entityId: 303, + window: "monthly", + status: "decremented", + newRemaining: 98.75, + }, + ], + }); + + expect(mockRedis.eval).toHaveBeenCalledTimes(1); + expect(String(mockRedis.eval.mock.calls[0]?.[0])).not.toMatch(/\bSCAN\b/i); + expect(mockRedis.eval).toHaveBeenCalledWith( + expect.any(String), + 13, + "lease:settlement:9001", + buildLeaseKey("key", 101, "5h", "rolling"), + buildLeaseKey("key", 101, "daily", "fixed"), + buildLeaseKey("key", 101, "weekly"), + buildLeaseKey("key", 101, "monthly"), + buildLeaseKey("user", 202, "5h", "fixed"), + buildLeaseKey("user", 202, "daily", "rolling"), + buildLeaseKey("user", 202, "weekly"), + buildLeaseKey("user", 202, "monthly"), + buildLeaseKey("provider", 303, "5h", "rolling"), + buildLeaseKey("provider", 303, "daily", "fixed"), + buildLeaseKey("provider", 303, "weekly"), + buildLeaseKey("provider", 303, "monthly"), + "1.25", + "300" + ); + }); + + it("reports a replay as duplicate while preserving the original settlement details", async () => { + const { LeaseService } = await import("@/lib/rate-limit/lease-service"); + + mockRedis.eval + .mockResolvedValueOnce([0, encodedSettlements]) + .mockResolvedValueOnce([1, encodedSettlements]); + + const first = await LeaseService.settleLeaseBudgets(settlementParams); + const replay = await LeaseService.settleLeaseBudgets(settlementParams); + + expect(first.status).toBe("settled"); + expect(replay).toEqual({ + ...first, + status: "duplicate", + }); + + const script = String(mockRedis.eval.mock.calls[0]?.[0]); + expect(script).toContain('redis.call("GET", markerKey)'); + expect(script).toContain("return {1, previousSettlement}"); + expect(script).toContain('redis.call("SETEX", markerKey, markerTtlSeconds, encoded)'); + expect(script.indexOf("return {1, previousSettlement}")).toBeLessThan( + script.indexOf("for keyIndex = 2, #KEYS do") + ); + }); + + it("validates every lease before mutating any budget", async () => { + const { LeaseService } = await import("@/lib/rate-limit/lease-service"); + + mockRedis.eval.mockResolvedValue([0, encodedSettlements]); + await LeaseService.settleLeaseBudgets(settlementParams); + + const script = String(mockRedis.eval.mock.calls[0]?.[0]); + expect(script).toContain('local leaseReply = redis.pcall("GET", leaseKey)'); + expect(script).toContain("local pendingWrites = {}"); + expect(script).not.toContain('redis.call("SETEX", leaseKey'); + expect(script.indexOf("for writeIndex = 1, #pendingWrites do")).toBeLessThan( + script.indexOf('redis.call("SETEX", pendingWrite[1]') + ); + }); + + it("returns a structured fail-open result when Redis is unavailable", async () => { + const { LeaseService } = await import("@/lib/rate-limit/lease-service"); + + mockRedis.status = "connecting"; + + await expect(LeaseService.settleLeaseBudgets(settlementParams)).resolves.toEqual({ + requestId: "9001", + status: "fail_open", + settlements: [], + failOpen: true, + }); + expect(mockRedis.eval).not.toHaveBeenCalled(); + + mockRedis.status = "ready"; + }); + }); + describe("getCostLease - limit change detection", () => { const nowMs = 1706400000000; diff --git a/tests/unit/lib/rate-limit/rolling-window-5h.test.ts b/tests/unit/lib/rate-limit/rolling-window-5h.test.ts index 1fe0cd1d2..18240b3d2 100644 --- a/tests/unit/lib/rate-limit/rolling-window-5h.test.ts +++ b/tests/unit/lib/rate-limit/rolling-window-5h.test.ts @@ -20,6 +20,10 @@ vi.mock("@/lib/utils/timezone", () => ({ const pipelineCommands: Array = []; const pipeline = { + eval: vi.fn((...args: unknown[]) => { + pipelineCommands.push(["eval", ...args]); + return pipeline; + }), zadd: vi.fn((...args: unknown[]) => { pipelineCommands.push(["zadd", ...args]); return pipeline; @@ -88,10 +92,6 @@ describe("RateLimitService - 5h rolling window behavior", () => { it("T0: consume $10, window should be $10", async () => { const { RateLimitService } = await import("@/lib/rate-limit"); - // trackCost calls eval twice (key + provider) - redisClient.eval.mockResolvedValueOnce("10"); // TRACK key - redisClient.eval.mockResolvedValueOnce("10"); // TRACK provider - await RateLimitService.trackCost(1, 2, "sess", 10, { requestId: 1, createdAtMs: baseTime }); // getCurrentCost calls eval once, then exists @@ -105,18 +105,12 @@ describe("RateLimitService - 5h rolling window behavior", () => { it("T1 (3h later): consume $20, window should be $30", async () => { const { RateLimitService } = await import("@/lib/rate-limit"); - // T0: Track $10 (2 evals: key + provider) - redisClient.eval.mockResolvedValueOnce("10"); - redisClient.eval.mockResolvedValueOnce("10"); await RateLimitService.trackCost(1, 2, "sess", 10, { requestId: 1, createdAtMs: baseTime }); // T1: Move to 3h later const t1 = baseTime + 3 * 60 * 60 * 1000; vi.setSystemTime(new Date(t1)); - // Track $20 (2 evals: key + provider) - redisClient.eval.mockResolvedValueOnce("20"); - redisClient.eval.mockResolvedValueOnce("20"); await RateLimitService.trackCost(1, 2, "sess", 20, { requestId: 2, createdAtMs: t1 }); // getCurrentCost: eval returns sum @@ -130,16 +124,11 @@ describe("RateLimitService - 5h rolling window behavior", () => { it("T2 (6h later): query cost, should only include T1 ($20) as T0 expired", async () => { const { RateLimitService } = await import("@/lib/rate-limit"); - // T0: Track $10 (2 evals) - redisClient.eval.mockResolvedValueOnce("10"); - redisClient.eval.mockResolvedValueOnce("10"); await RateLimitService.trackCost(1, 2, "sess", 10, { requestId: 1, createdAtMs: baseTime }); - // T1: 3h later, track $20 (2 evals) + // T1: 3h later, track $20 const t1 = baseTime + 3 * 60 * 60 * 1000; vi.setSystemTime(new Date(t1)); - redisClient.eval.mockResolvedValueOnce("30"); - redisClient.eval.mockResolvedValueOnce("30"); await RateLimitService.trackCost(1, 2, "sess", 20, { requestId: 2, createdAtMs: t1 }); // T2: 6h after T0 (3h after T1) @@ -164,16 +153,11 @@ describe("RateLimitService - 5h rolling window behavior", () => { it("T0: consume $5, T1 (4h59m later): consume $10, window = $15", async () => { const { RateLimitService } = await import("@/lib/rate-limit"); - // T0: Track $5 (2 evals) - redisClient.eval.mockResolvedValueOnce("5"); - redisClient.eval.mockResolvedValueOnce("5"); await RateLimitService.trackCost(1, 2, "sess", 5, { requestId: 1, createdAtMs: baseTime }); // T1: 4h59m later (still within 5h) const t1 = baseTime + (4 * 60 + 59) * 60 * 1000; vi.setSystemTime(new Date(t1)); - redisClient.eval.mockResolvedValueOnce("15"); - redisClient.eval.mockResolvedValueOnce("15"); await RateLimitService.trackCost(1, 2, "sess", 10, { requestId: 2, createdAtMs: t1 }); // Both entries should be in window @@ -187,16 +171,11 @@ describe("RateLimitService - 5h rolling window behavior", () => { it("T2 (5h01m after T0): query, window = $10 (T0 expired)", async () => { const { RateLimitService } = await import("@/lib/rate-limit"); - // T0: Track $5 (2 evals) - redisClient.eval.mockResolvedValueOnce("5"); - redisClient.eval.mockResolvedValueOnce("5"); await RateLimitService.trackCost(1, 2, "sess", 5, { requestId: 1, createdAtMs: baseTime }); - // T1: 4h59m later (2 evals) + // T1: 4h59m later const t1 = baseTime + (4 * 60 + 59) * 60 * 1000; vi.setSystemTime(new Date(t1)); - redisClient.eval.mockResolvedValueOnce("15"); - redisClient.eval.mockResolvedValueOnce("15"); await RateLimitService.trackCost(1, 2, "sess", 10, { requestId: 2, createdAtMs: t1 }); // T2: 5h01m after T0 @@ -216,30 +195,21 @@ describe("RateLimitService - 5h rolling window behavior", () => { it("should correctly calculate window with multiple entries at different times", async () => { const { RateLimitService } = await import("@/lib/rate-limit"); - // T0: $10 (2 evals) - redisClient.eval.mockResolvedValueOnce("10"); - redisClient.eval.mockResolvedValueOnce("10"); await RateLimitService.trackCost(1, 2, "sess", 10, { requestId: 1, createdAtMs: baseTime }); - // T1: 1h later, $20 (2 evals) + // T1: 1h later, $20 const t1 = baseTime + 1 * 60 * 60 * 1000; vi.setSystemTime(new Date(t1)); - redisClient.eval.mockResolvedValueOnce("30"); - redisClient.eval.mockResolvedValueOnce("30"); await RateLimitService.trackCost(1, 2, "sess", 20, { requestId: 2, createdAtMs: t1 }); - // T2: 2h later, $15 (2 evals) + // T2: 2h later, $15 const t2 = baseTime + 2 * 60 * 60 * 1000; vi.setSystemTime(new Date(t2)); - redisClient.eval.mockResolvedValueOnce("45"); - redisClient.eval.mockResolvedValueOnce("45"); await RateLimitService.trackCost(1, 2, "sess", 15, { requestId: 3, createdAtMs: t2 }); - // T3: 3h after T0, $25 (2 evals) + // T3: 3h after T0, $25 const t3 = baseTime + 3 * 60 * 60 * 1000; vi.setSystemTime(new Date(t3)); - redisClient.eval.mockResolvedValueOnce("70"); - redisClient.eval.mockResolvedValueOnce("70"); await RateLimitService.trackCost(1, 2, "sess", 25, { requestId: 4, createdAtMs: t3 }); // At T3: all 4 entries within window = $70 @@ -264,9 +234,6 @@ describe("RateLimitService - 5h rolling window behavior", () => { it("should reject request when rolling window exceeds limit", async () => { const { RateLimitService } = await import("@/lib/rate-limit"); - // T0: consume $40 (2 evals for trackCost) - redisClient.eval.mockResolvedValueOnce("40"); - redisClient.eval.mockResolvedValueOnce("40"); await RateLimitService.trackCost(1, 2, "sess", 40, { requestId: 1, createdAtMs: baseTime }); // Check limit (5h = $50) - checkCostLimits calls eval @@ -298,9 +265,7 @@ describe("RateLimitService - 5h rolling window behavior", () => { // Current is $40, limit is $50, should still be allowed expect(checkT1.allowed).toBe(true); - // After adding $20, would be $60 - trackCost (2 evals) - redisClient.eval.mockResolvedValueOnce("60"); - redisClient.eval.mockResolvedValueOnce("60"); + // After adding $20, would be $60. await RateLimitService.trackCost(1, 2, "sess", 20, { requestId: 2, createdAtMs: t1 }); // Verify window now shows $60 @@ -333,18 +298,12 @@ describe("RateLimitService - 5h rolling window behavior", () => { const day1_22h = new Date("2024-01-15T22:00:00.000Z").getTime(); vi.setSystemTime(new Date(day1_22h)); - // Track $10 (2 evals) - redisClient.eval.mockResolvedValueOnce("10"); - redisClient.eval.mockResolvedValueOnce("10"); await RateLimitService.trackCost(1, 2, "sess", 10, { requestId: 1, createdAtMs: day1_22h }); // Day2 01:00 UTC (3h later, crossed midnight) const day2_01h = new Date("2024-01-16T01:00:00.000Z").getTime(); vi.setSystemTime(new Date(day2_01h)); - // Track $20 (2 evals) - redisClient.eval.mockResolvedValueOnce("30"); - redisClient.eval.mockResolvedValueOnce("30"); await RateLimitService.trackCost(1, 2, "sess", 20, { requestId: 2, createdAtMs: day2_01h }); // Both entries in window = $30 @@ -401,16 +360,11 @@ describe("RateLimitService - 5h rolling window behavior", () => { it("should work identically for provider entities", async () => { const { RateLimitService } = await import("@/lib/rate-limit"); - // T0: provider consumes $15 (2 evals) - redisClient.eval.mockResolvedValueOnce("15"); - redisClient.eval.mockResolvedValueOnce("15"); await RateLimitService.trackCost(1, 2, "sess", 15, { requestId: 1, createdAtMs: baseTime }); - // T1: 4h later, consume $25 (2 evals) + // T1: 4h later, consume $25 const t1 = baseTime + 4 * 60 * 60 * 1000; vi.setSystemTime(new Date(t1)); - redisClient.eval.mockResolvedValueOnce("40"); - redisClient.eval.mockResolvedValueOnce("40"); await RateLimitService.trackCost(1, 2, "sess", 25, { requestId: 2, createdAtMs: t1 }); // Window = $40 diff --git a/tests/unit/lib/rate-limit/rolling-window-cache-warm.test.ts b/tests/unit/lib/rate-limit/rolling-window-cache-warm.test.ts index 5493e83ae..e6c2080c0 100644 --- a/tests/unit/lib/rate-limit/rolling-window-cache-warm.test.ts +++ b/tests/unit/lib/rate-limit/rolling-window-cache-warm.test.ts @@ -3,6 +3,10 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; const pipelineCommands: Array = []; const pipeline = { + eval: vi.fn((...args: unknown[]) => { + pipelineCommands.push(["eval", ...args]); + return pipeline; + }), zadd: vi.fn((...args: unknown[]) => { pipelineCommands.push(["zadd", ...args]); return pipeline; @@ -119,12 +123,13 @@ describe("RateLimitService rolling window cache warm", () => { providerResetMode: "fixed", }); - const evalCalls = redisClient.eval.mock.calls; + const evalCalls = pipelineCommands.filter((call) => call[0] === "eval"); expect(evalCalls.length).toBeGreaterThanOrEqual(2); const [firstCall] = evalCalls; - expect(firstCall[2]).toBe("key:1:cost_5h_rolling"); - expect(firstCall[4]).toBe(String(nowMs - 1000)); - expect(firstCall[6]).toBe("123"); + expect(firstCall[3]).toBe("key:1:cost_5h_rolling"); + expect(firstCall[5]).toBe(String(nowMs - 1000)); + expect(firstCall[7]).toBe("123"); + expect(redisClient.eval).not.toHaveBeenCalled(); }); }); diff --git a/tests/unit/lib/rate-limit/service-extra.test.ts b/tests/unit/lib/rate-limit/service-extra.test.ts index 46d411f77..c75e259dc 100644 --- a/tests/unit/lib/rate-limit/service-extra.test.ts +++ b/tests/unit/lib/rate-limit/service-extra.test.ts @@ -297,7 +297,11 @@ describe("RateLimitService - other quota paths", () => { await RateLimitService.trackUserDailyCost(1, 1.25, "00:00", "rolling", { requestId: 123 }); - expect(redisClientRef.eval).toHaveBeenCalled(); + expect(redisClientRef.pipeline).toHaveBeenCalledTimes(1); + expect(redisClientRef.eval).not.toHaveBeenCalled(); + expect( + pipelineCalls.some((call) => call[0] === "eval" && call[3] === "user:1:cost_daily_rolling") + ).toBe(true); }); it("checkUserRPM:达到上限时应拦截", async () => { @@ -364,6 +368,43 @@ describe("RateLimitService - other quota paths", () => { expect(result.size).toBe(0); }); + it("settleLeaseBudgets:应通过单次 Redis 结算返回 12 个显式窗口结果", async () => { + const { RateLimitService } = await import("@/lib/rate-limit"); + + redisClientRef.eval.mockResolvedValueOnce([ + 0, + JSON.stringify(Array.from({ length: 12 }, (_, index) => [1, 100 - index])), + ]); + + const result = await RateLimitService.settleLeaseBudgets({ + requestId: 7001, + cost: 0.75, + entities: { + key: { id: 11 }, + user: { id: 22 }, + provider: { id: 33 }, + }, + }); + + expect(result.status).toBe("settled"); + expect(result.settlements).toHaveLength(12); + expect(result.settlements.map(({ entityType, window }) => `${entityType}:${window}`)).toEqual([ + "key:5h", + "key:daily", + "key:weekly", + "key:monthly", + "user:5h", + "user:daily", + "user:weekly", + "user:monthly", + "provider:5h", + "provider:daily", + "provider:weekly", + "provider:monthly", + ]); + expect(redisClientRef.eval).toHaveBeenCalledTimes(1); + }); + it("getCurrentCostBatch:Redis 非 ready 时应返回默认 0", async () => { const { RateLimitService } = await import("@/lib/rate-limit"); @@ -520,25 +561,26 @@ describe("RateLimitService - other quota paths", () => { }); expect( - redisClientRef.eval.mock.calls.some( - (call: unknown[]) => String(call[2]) === "key:1:cost_5h_fixed" + pipelineCalls.some( + (call: unknown[]) => call[0] === "eval" && String(call[3]) === "key:1:cost_5h_fixed" ) ).toBe(true); expect( - redisClientRef.eval.mock.calls.some( - (call: unknown[]) => String(call[2]) === "provider:9:cost_5h_fixed" + pipelineCalls.some( + (call: unknown[]) => call[0] === "eval" && String(call[3]) === "provider:9:cost_5h_fixed" ) ).toBe(true); expect( - redisClientRef.eval.mock.calls.some( - (call: unknown[]) => String(call[2]) === "user:7:cost_5h_fixed" + pipelineCalls.some( + (call: unknown[]) => call[0] === "eval" && String(call[3]) === "user:7:cost_5h_fixed" ) ).toBe(true); expect( - redisClientRef.eval.mock.calls.some((call: unknown[]) => - String(call[2]).includes("cost_5h_rolling") + pipelineCalls.some( + (call: unknown[]) => call[0] === "eval" && String(call[3]).includes("cost_5h_rolling") ) ).toBe(false); + expect(redisClientRef.eval).not.toHaveBeenCalled(); }); it("trackCost:fixed 模式应写入 key/provider 的 daily+weekly+monthly(STRING)", async () => { @@ -553,8 +595,9 @@ describe("RateLimitService - other quota paths", () => { createdAtMs: nowMs, }); - // 5h 的 Lua 脚本至少会执行两次(key/provider) - expect(redisClientRef.eval).toHaveBeenCalled(); + // 5h Lua 与固定窗口命令均进入同一个 pipeline。 + expect(pipelineCalls.filter((c) => c[0] === "eval").length).toBeGreaterThanOrEqual(2); + expect(redisClientRef.eval).not.toHaveBeenCalled(); expect(pipelineCalls.filter((c) => c[0] === "incrbyfloat").length).toBeGreaterThanOrEqual(4); expect(pipelineCalls.filter((c) => c[0] === "expire").length).toBeGreaterThanOrEqual(4); }); @@ -569,9 +612,99 @@ describe("RateLimitService - other quota paths", () => { createdAtMs: nowMs, }); - const evalArgs = redisClientRef.eval.mock.calls.map((c: unknown[]) => String(c[2])); + const evalArgs = pipelineCalls + .filter((c) => c[0] === "eval") + .map((c: unknown[]) => String(c[3])); expect(evalArgs.some((k) => k === "key:1:cost_daily_rolling")).toBe(true); expect(evalArgs.some((k) => k === "provider:9:cost_daily_rolling")).toBe(true); + expect(redisClientRef.eval).not.toHaveBeenCalled(); + }); + + it("trackCost:应把 Key、Provider 与 User 的真实消费合并为一个 pipeline", async () => { + const { RateLimitService } = await import("@/lib/rate-limit"); + + const options = { + userId: 7, + key5hResetMode: "rolling" as const, + keyResetMode: "fixed" as const, + keyResetTime: "00:00", + provider5hResetMode: "fixed" as const, + providerResetMode: "rolling" as const, + providerResetTime: "03:30", + user5hResetMode: "rolling" as const, + userResetMode: "rolling" as const, + userResetTime: "06:45", + requestId: 123, + createdAtMs: nowMs, + }; + + await RateLimitService.trackCost(1, 9, "sess", 1.25, options); + + expect(redisClientRef.pipeline).toHaveBeenCalledTimes(1); + expect(redisClientRef.eval).not.toHaveBeenCalled(); + expect(pipelineCalls.filter((call) => call[0] === "exec")).toHaveLength(1); + + const evalKeys = pipelineCalls + .filter((call) => call[0] === "eval") + .map((call) => String(call[3])); + expect(evalKeys).toEqual([ + "key:1:cost_5h_rolling", + "provider:9:cost_5h_fixed", + "user:7:cost_5h_rolling", + "provider:9:cost_daily_rolling", + "user:7:cost_daily_rolling", + ]); + + const fixedCounterKeys = pipelineCalls + .filter((call) => call[0] === "incrbyfloat") + .map((call) => String(call[1])); + expect(fixedCounterKeys).toEqual([ + "key:1:cost_daily_0000", + "key:1:cost_weekly", + "key:1:cost_monthly", + "provider:9:cost_weekly", + "provider:9:cost_monthly", + ]); + expect(fixedCounterKeys.some((key) => key.startsWith("user:7:cost_weekly"))).toBe(false); + expect(fixedCounterKeys.some((key) => key.startsWith("user:7:cost_monthly"))).toBe(false); + }); + + it("trackCost:pipeline 单命令失败时应逐项记录并继续 fail-open", async () => { + const { logger } = await import("@/lib/logger"); + const { RateLimitService } = await import("@/lib/rate-limit"); + const pipeline = makePipeline(); + pipeline.exec.mockResolvedValueOnce([[new Error("pipeline boom"), null]]); + redisClientRef.pipeline.mockReturnValueOnce(pipeline); + + await expect( + RateLimitService.trackCost(1, 9, "sess", 1.25, { + key5hResetMode: "fixed", + provider5hResetMode: "fixed", + keyResetMode: "fixed", + providerResetMode: "fixed", + requestId: 123, + createdAtMs: nowMs, + }) + ).resolves.toBeUndefined(); + + expect(logger.error).toHaveBeenCalledWith( + "[RateLimit] Cost pipeline command failed", + expect.objectContaining({ commandIndex: 0, error: "pipeline boom" }) + ); + }); + + it("trackCost:Redis 非 ready 时不应创建 pipeline 或发送命令", async () => { + const { RateLimitService } = await import("@/lib/rate-limit"); + redisClientRef.status = "connecting"; + + await RateLimitService.trackCost(1, 9, "sess", 1.25, { + userId: 7, + requestId: 123, + createdAtMs: nowMs, + }); + + expect(redisClientRef.pipeline).not.toHaveBeenCalled(); + expect(redisClientRef.eval).not.toHaveBeenCalled(); }); it("getCurrentCostBatch:pipeline.exec 返回 null 时应返回默认值", async () => { diff --git a/tests/unit/lib/redis/client.test.ts b/tests/unit/lib/redis/client.test.ts index eb32520ec..c6dfa6c49 100644 --- a/tests/unit/lib/redis/client.test.ts +++ b/tests/unit/lib/redis/client.test.ts @@ -34,22 +34,47 @@ vi.mock("server-only", () => ({})); import { buildRedisOptionsForUrl, closeRedis, getRedisClient } from "@/lib/redis/client"; describe("buildRedisOptionsForUrl", () => { + afterEach(() => { + delete process.env.REDIS_COMMAND_TIMEOUT_MS; + }); + it("detects TLS from rediss:// protocol", () => { const result = buildRedisOptionsForUrl("rediss://localhost:6380"); expect(result.isTLS).toBe(true); expect(result.options.tls).toBeDefined(); + expect(result.options.commandTimeout).toBe(10_000); + expect(result.options.socketTimeout).toBe(15_000); + expect(result.options.autoResendUnfulfilledCommands).toBe(false); }); it("does not enable TLS for redis:// protocol", () => { const result = buildRedisOptionsForUrl("redis://localhost:6379"); expect(result.isTLS).toBe(false); expect(result.options.tls).toBeUndefined(); + expect(result.options.commandTimeout).toBe(10_000); + expect(result.options.socketTimeout).toBe(15_000); + expect(result.options.autoResendUnfulfilledCommands).toBe(false); }); it("falls back to string-prefix detection for malformed URLs", () => { const result = buildRedisOptionsForUrl("rediss://not a valid url"); expect(result.isTLS).toBe(true); }); + + it.each(["redis://localhost:6379", "rediss://localhost:6380"])( + "supports REDIS_COMMAND_TIMEOUT_MS override for %s", + async (redisUrl) => { + process.env.REDIS_COMMAND_TIMEOUT_MS = "2500"; + vi.resetModules(); + const { buildRedisOptionsForUrl: buildFreshOptions } = await import("@/lib/redis/client"); + + const result = buildFreshOptions(redisUrl); + + expect(result.options.commandTimeout).toBe(2_500); + expect(result.options.socketTimeout).toBe(7_500); + expect(result.options.autoResendUnfulfilledCommands).toBe(false); + } + ); }); describe("getRedisClient", () => { diff --git a/tests/unit/lib/session-manager-helpers.test.ts b/tests/unit/lib/session-manager-helpers.test.ts index 5a02f8dbf..e637f1be0 100644 --- a/tests/unit/lib/session-manager-helpers.test.ts +++ b/tests/unit/lib/session-manager-helpers.test.ts @@ -60,6 +60,17 @@ describe("SessionManager 辅助函数", () => { expect(getParseHeaderRecordWarnCalls()).toHaveLength(0); }); + test("parseHeaderRecord:移除历史快照中的内部 x-cch header", async () => { + vi.clearAllMocks(); + const { parseHeaderRecord } = await loadHelpers(); + + expect( + parseHeaderRecord( + '{"x-cch-internal-secret":"secret-canary","x-cch-future-marker":"1","x-safe":"ok"}' + ) + ).toEqual({ "x-safe": "ok" }); + }); + test("parseHeaderRecord:无效 JSON 应返回 null 并记录 warn", async () => { vi.clearAllMocks(); const { parseHeaderRecord } = await loadHelpers(); diff --git a/tests/unit/lib/session-manager-terminate-session.test.ts b/tests/unit/lib/session-manager-terminate-session.test.ts index f61889538..76145bd34 100644 --- a/tests/unit/lib/session-manager-terminate-session.test.ts +++ b/tests/unit/lib/session-manager-terminate-session.test.ts @@ -35,6 +35,8 @@ describe("SessionManager.terminateSession", () => { status: "ready", get: vi.fn(async () => null), hget: vi.fn(async () => null), + del: vi.fn(async () => 1), + eval: vi.fn(async () => 1), pipeline: vi.fn(() => pipelineRef), }; }); @@ -83,4 +85,26 @@ describe("SessionManager.terminateSession", () => { expect(pipelineRef.zrem).not.toHaveBeenCalledWith(getUserActiveSessionsKey(123), sessionId); }); + + it("迟到 cleanup 仅删除仍绑定到预期 provider 的 session", async () => { + const { SessionManager } = await import("@/lib/session-manager"); + + await expect(SessionManager.clearSessionProvider("sess_compare", 42)).resolves.toBe(true); + + expect(redisClientRef.eval).toHaveBeenCalledWith( + expect.stringContaining('redis.call("GET", KEYS[1]) == ARGV[1]'), + 1, + "session:sess_compare:provider", + "42" + ); + expect(redisClientRef.del).not.toHaveBeenCalled(); + }); + + it("迟到 cleanup 不删除已切换到新 provider 的 session", async () => { + redisClientRef.eval.mockResolvedValueOnce(0); + const { SessionManager } = await import("@/lib/session-manager"); + + await expect(SessionManager.clearSessionProvider("sess_compare", 42)).resolves.toBe(false); + expect(redisClientRef.del).not.toHaveBeenCalled(); + }); }); diff --git a/tests/unit/lib/shutdown.test.ts b/tests/unit/lib/shutdown.test.ts index 3d3a39ae3..264a1b682 100644 --- a/tests/unit/lib/shutdown.test.ts +++ b/tests/unit/lib/shutdown.test.ts @@ -14,15 +14,21 @@ describe.sequential("lifecycle/shutdown", () => { vi.resetModules(); vi.useRealTimers(); delete (globalThis as unknown as { __CCH_LIFECYCLE__?: unknown }).__CCH_LIFECYCLE__; + delete (globalThis as unknown as { __ASYNC_TASK_MANAGER__?: unknown }).__ASYNC_TASK_MANAGER__; delete (globalThis as unknown as { __CCH_CLOUD_PRICE_SYNC_INTERVAL_ID__?: unknown }) .__CCH_CLOUD_PRICE_SYNC_INTERVAL_ID__; delete (globalThis as unknown as { __CCH_API_KEY_VF_SYNC_CLEANUP__?: unknown }) .__CCH_API_KEY_VF_SYNC_CLEANUP__; + delete (globalThis as unknown as { __CCH_STOP_BACKGROUND_QUEUES__?: unknown }) + .__CCH_STOP_BACKGROUND_QUEUES__; }); afterEach(() => { vi.useRealTimers(); vi.unstubAllEnvs(); + delete (globalThis as unknown as { __ASYNC_TASK_MANAGER__?: unknown }).__ASYNC_TASK_MANAGER__; + delete (globalThis as unknown as { __CCH_STOP_BACKGROUND_QUEUES__?: unknown }) + .__CCH_STOP_BACKGROUND_QUEUES__; }); it("markShuttingDown flips isShuttingDown idempotently", async () => { @@ -50,14 +56,14 @@ describe.sequential("lifecycle/shutdown", () => { expect(second).toBe(first); }); - it("runApplicationCleanup invokes the staged modules and survives one step throwing", async () => { + it("runApplicationCleanup invokes staged modules and survives one non-critical step throwing", async () => { const stopCache = vi.fn(); - const stopProbe = vi.fn(); + const stopProbe = vi.fn(() => { + throw new Error("simulated probe scheduler shutdown failure"); + }); const stopPublicStatus = vi.fn(async () => {}); const stopProbeLog = vi.fn(); - const shutdownTasks = vi.fn(() => { - throw new Error("simulated tasks shutdown failure"); - }); + const shutdownTasks = vi.fn(async () => {}); const stopWriteBuffer = vi.fn(async () => {}); const shutdownLf = vi.fn(async () => {}); const closeRedis = vi.fn(async () => {}); @@ -108,6 +114,49 @@ describe.sequential("lifecycle/shutdown", () => { expect(clearIntervalSpy).toHaveBeenCalledWith(intervalId); }); + it("waits for scheduler quiescence after the warning threshold before closing resources", async () => { + let releaseScheduler!: () => void; + const schedulerStopped = new Promise((resolve) => { + releaseScheduler = resolve; + }); + const closeDbPools = vi.fn(async () => {}); + const closeRedis = vi.fn(async () => {}); + + vi.doMock("@/lib/cache/session-cache", () => ({ stopCacheCleanup: () => {} })); + vi.doMock("@/lib/provider-endpoints/probe-scheduler", () => ({ + stopEndpointProbeScheduler: () => schedulerStopped, + })); + vi.doMock("@/lib/public-status/scheduler", () => ({ + stopPublicStatusRebuildScheduler: async () => {}, + })); + vi.doMock("@/lib/provider-endpoints/probe-log-cleanup", () => ({ + stopEndpointProbeLogCleanup: async () => {}, + })); + vi.doMock("@/lib/async-task-manager", () => ({ shutdownAllAsyncTasks: async () => {} })); + vi.doMock("@/repository/message-write-buffer", () => ({ + stopMessageRequestWriteBuffer: async () => {}, + })); + vi.doMock("@/drizzle/db", () => ({ closeDbPools })); + vi.doMock("@/lib/langfuse", () => ({ shutdownLangfuse: async () => {} })); + vi.doMock("@/lib/redis", () => ({ closeRedis })); + + const { runApplicationCleanup } = await import("@/lib/lifecycle/shutdown"); + const cleanup = runApplicationCleanup("SIGTERM", { + totalTimeoutMs: 5_000, + perStepTimeoutMs: 20, + }); + + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(closeDbPools).not.toHaveBeenCalled(); + expect(closeRedis).not.toHaveBeenCalled(); + + releaseScheduler(); + await cleanup; + + expect(closeDbPools).toHaveBeenCalledTimes(1); + expect(closeRedis).toHaveBeenCalledTimes(1); + }); + it("runApplicationCleanup returns within totalTimeoutMs even if one step hangs", async () => { let releaseHang: () => void = () => {}; const hang = new Promise((resolve) => { @@ -140,4 +189,122 @@ describe.sequential("lifecycle/shutdown", () => { expect(elapsed).toBeLessThan(2_000); releaseHang(); }); + + it("runApplicationCleanup abort 后等待所有 async task settled 再启动 writer", async () => { + vi.stubEnv("CI", "true"); + vi.stubEnv("NEXT_RUNTIME", "nodejs"); + vi.doUnmock("@/lib/async-task-manager"); + + vi.doMock("@/lib/cache/session-cache", () => ({ stopCacheCleanup: () => {} })); + vi.doMock("@/lib/provider-endpoints/probe-scheduler", () => ({ + stopEndpointProbeScheduler: () => {}, + })); + vi.doMock("@/lib/public-status/scheduler", () => ({ + stopPublicStatusRebuildScheduler: async () => {}, + })); + vi.doMock("@/lib/provider-endpoints/probe-log-cleanup", () => ({ + stopEndpointProbeLogCleanup: () => {}, + })); + + const writerStarted = vi.fn(); + vi.doMock("@/repository/message-write-buffer", () => ({ + stopMessageRequestWriteBuffer: async () => { + writerStarted(); + }, + })); + vi.doMock("@/drizzle/db", () => ({ closeDbPools: async () => {} })); + vi.doMock("@/lib/langfuse", () => ({ shutdownLangfuse: async () => {} })); + vi.doMock("@/lib/redis", () => ({ closeRedis: async () => {} })); + + const { AsyncTaskManager } = await import("@/lib/async-task-manager"); + let resolveFirst!: () => void; + let resolveSecond!: () => void; + const first = new Promise((resolve) => { + resolveFirst = resolve; + }); + const second = new Promise((resolve) => { + resolveSecond = resolve; + }); + const firstController = new AbortController(); + const secondController = new AbortController(); + const firstAborted = new Promise((resolve) => { + firstController.signal.addEventListener("abort", () => resolve(), { once: true }); + }); + const secondAborted = new Promise((resolve) => { + secondController.signal.addEventListener("abort", () => resolve(), { once: true }); + }); + AsyncTaskManager.register("shutdown-first", () => first, { + abortController: firstController, + }); + AsyncTaskManager.register("shutdown-second", () => second, { + abortController: secondController, + }); + + const { runApplicationCleanup } = await import("@/lib/lifecycle/shutdown"); + const cleanup = runApplicationCleanup("SIGTERM", { + totalTimeoutMs: 5_000, + perStepTimeoutMs: 20, + }); + await Promise.all([firstAborted, secondAborted]); + + expect(firstController.signal.aborted).toBe(true); + expect(secondController.signal.aborted).toBe(true); + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(writerStarted).not.toHaveBeenCalled(); + + resolveFirst(); + await new Promise((resolve) => queueMicrotask(resolve)); + expect(writerStarted).not.toHaveBeenCalled(); + + resolveSecond(); + await cleanup; + + expect(writerStarted).toHaveBeenCalledTimes(1); + }); + + it("continues critical cleanup after background queue shutdown fails", async () => { + const queueError = new Error("queue stop failed"); + const shutdownTasks = vi.fn(async () => {}); + const stopWriteBuffer = vi.fn(async () => {}); + const closeDbPools = vi.fn(async () => {}); + const shutdownLangfuse = vi.fn(async () => {}); + const closeRedis = vi.fn(async () => {}); + + vi.doMock("@/lib/cache/session-cache", () => ({ stopCacheCleanup: () => {} })); + vi.doMock("@/lib/provider-endpoints/probe-scheduler", () => ({ + stopEndpointProbeScheduler: () => {}, + })); + vi.doMock("@/lib/public-status/scheduler", () => ({ + stopPublicStatusRebuildScheduler: async () => {}, + })); + vi.doMock("@/lib/provider-endpoints/probe-log-cleanup", () => ({ + stopEndpointProbeLogCleanup: () => {}, + })); + vi.doMock("@/lib/async-task-manager", () => ({ shutdownAllAsyncTasks: shutdownTasks })); + vi.doMock("@/repository/message-write-buffer", () => ({ + stopMessageRequestWriteBuffer: stopWriteBuffer, + })); + vi.doMock("@/drizzle/db", () => ({ closeDbPools })); + vi.doMock("@/lib/langfuse", () => ({ shutdownLangfuse })); + vi.doMock("@/lib/redis", () => ({ closeRedis })); + ( + globalThis as unknown as { __CCH_STOP_BACKGROUND_QUEUES__?: () => Promise } + ).__CCH_STOP_BACKGROUND_QUEUES__ = vi.fn().mockRejectedValue(queueError); + + const { runApplicationCleanup } = await import("@/lib/lifecycle/shutdown"); + let thrown: unknown; + try { + await runApplicationCleanup("SIGTERM", { totalTimeoutMs: 5_000, perStepTimeoutMs: 100 }); + } catch (error) { + thrown = error; + } + + expect(thrown).toBeInstanceOf(AggregateError); + expect((thrown as AggregateError).errors).toContain(queueError); + expect(shutdownTasks).toHaveBeenCalledOnce(); + expect(stopWriteBuffer).toHaveBeenCalledOnce(); + expect(closeDbPools).toHaveBeenCalledOnce(); + expect(shutdownLangfuse).toHaveBeenCalledOnce(); + expect(closeRedis).toHaveBeenCalledOnce(); + }); }); diff --git a/tests/unit/lib/upstream-error-detection-status.test.ts b/tests/unit/lib/upstream-error-detection-status.test.ts index e34cf75a5..2bac09c5c 100644 --- a/tests/unit/lib/upstream-error-detection-status.test.ts +++ b/tests/unit/lib/upstream-error-detection-status.test.ts @@ -28,93 +28,87 @@ const cloudflareErrorCases = [ ] as const; describe("inferUpstreamErrorStatusCodeFromText numeric boundaries", () => { - it.each(httpStatusCases)("keeps matching a standalone HTTP $statusCode status token", ({ - statusCode, - matcherId, - }) => { - expect(inferUpstreamErrorStatusCodeFromText(`HTTP/1.1 ${statusCode}`)).toEqual({ - statusCode, - matcherId, - }); - }); - - it.each( - httpStatusCases - )("does not treat HTTP $statusCode followed by a decimal fraction as a status token", ({ - statusCode, - }) => { - expect(inferUpstreamErrorStatusCodeFromText(`HTTP/1.1 ${statusCode}.12`)).toBeNull(); - }); - - it.each( - httpStatusCases - )("does not treat HTTP $statusCode embedded in a longer number as a status token", ({ - statusCode, - }) => { - expect(inferUpstreamErrorStatusCodeFromText(`HTTP/1.1 ${statusCode}12`)).toBeNull(); - }); - - it.each( - httpStatusCases - )("does not treat HTTP $statusCode followed by a letter as a status token", ({ statusCode }) => { - expect(inferUpstreamErrorStatusCodeFromText(`HTTP/1.1 ${statusCode}abc`)).toBeNull(); - }); - - it.each(httpStatusCases)("keeps matching HTTP $statusCode followed by sentence punctuation", ({ - statusCode, - matcherId, - }) => { - expect(inferUpstreamErrorStatusCodeFromText(`HTTP/1.1 ${statusCode}.`)).toEqual({ - statusCode, - matcherId, - }); - }); - - it.each(cloudflareErrorCases)("keeps matching a standalone Cloudflare Error $code token", ({ - code, - statusCode, - matcherId, - }) => { - expect(inferUpstreamErrorStatusCodeFromText(`Error ${code}`)).toEqual({ - statusCode, - matcherId, - }); - }); - - it.each( - cloudflareErrorCases - )("does not treat Cloudflare Error $code followed by a decimal fraction as a code token", ({ - code, - }) => { - expect(inferUpstreamErrorStatusCodeFromText(`Error ${code}.7`)).toBeNull(); - }); - - it.each( - cloudflareErrorCases - )("does not treat Cloudflare Error $code embedded in a longer number as a code token", ({ - code, - }) => { - expect(inferUpstreamErrorStatusCodeFromText(`Error ${code}7`)).toBeNull(); - }); - - it.each( - cloudflareErrorCases - )("does not treat Cloudflare Error $code followed by a letter as a code token", ({ code }) => { - expect(inferUpstreamErrorStatusCodeFromText(`Error ${code}x`)).toBeNull(); - }); - - it.each( - cloudflareErrorCases - )("keeps matching Cloudflare Error $code followed by sentence punctuation", ({ - code, - statusCode, - matcherId, - }) => { - expect(inferUpstreamErrorStatusCodeFromText(`Error ${code}.`)).toEqual({ - statusCode, - matcherId, - }); - }); + it.each(httpStatusCases)( + "keeps matching a standalone HTTP $statusCode status token", + ({ statusCode, matcherId }) => { + expect(inferUpstreamErrorStatusCodeFromText(`HTTP/1.1 ${statusCode}`)).toEqual({ + statusCode, + matcherId, + }); + } + ); + + it.each(httpStatusCases)( + "does not treat HTTP $statusCode followed by a decimal fraction as a status token", + ({ statusCode }) => { + expect(inferUpstreamErrorStatusCodeFromText(`HTTP/1.1 ${statusCode}.12`)).toBeNull(); + } + ); + + it.each(httpStatusCases)( + "does not treat HTTP $statusCode embedded in a longer number as a status token", + ({ statusCode }) => { + expect(inferUpstreamErrorStatusCodeFromText(`HTTP/1.1 ${statusCode}12`)).toBeNull(); + } + ); + + it.each(httpStatusCases)( + "does not treat HTTP $statusCode followed by a letter as a status token", + ({ statusCode }) => { + expect(inferUpstreamErrorStatusCodeFromText(`HTTP/1.1 ${statusCode}abc`)).toBeNull(); + } + ); + + it.each(httpStatusCases)( + "keeps matching HTTP $statusCode followed by sentence punctuation", + ({ statusCode, matcherId }) => { + expect(inferUpstreamErrorStatusCodeFromText(`HTTP/1.1 ${statusCode}.`)).toEqual({ + statusCode, + matcherId, + }); + } + ); + + it.each(cloudflareErrorCases)( + "keeps matching a standalone Cloudflare Error $code token", + ({ code, statusCode, matcherId }) => { + expect(inferUpstreamErrorStatusCodeFromText(`Error ${code}`)).toEqual({ + statusCode, + matcherId, + }); + } + ); + + it.each(cloudflareErrorCases)( + "does not treat Cloudflare Error $code followed by a decimal fraction as a code token", + ({ code }) => { + expect(inferUpstreamErrorStatusCodeFromText(`Error ${code}.7`)).toBeNull(); + } + ); + + it.each(cloudflareErrorCases)( + "does not treat Cloudflare Error $code embedded in a longer number as a code token", + ({ code }) => { + expect(inferUpstreamErrorStatusCodeFromText(`Error ${code}7`)).toBeNull(); + } + ); + + it.each(cloudflareErrorCases)( + "does not treat Cloudflare Error $code followed by a letter as a code token", + ({ code }) => { + expect(inferUpstreamErrorStatusCodeFromText(`Error ${code}x`)).toBeNull(); + } + ); + + it.each(cloudflareErrorCases)( + "keeps matching Cloudflare Error $code followed by sentence punctuation", + ({ code, statusCode, matcherId }) => { + expect(inferUpstreamErrorStatusCodeFromText(`Error ${code}.`)).toEqual({ + statusCode, + matcherId, + }); + } + ); it("does not infer service_unavailable from an AWS request id containing 503", () => { const text = "request id: 202604250550399959"; diff --git a/tests/unit/price-sync/cloud-price-updater.test.ts b/tests/unit/price-sync/cloud-price-updater.test.ts index e1ddcfe55..fd17f9fbf 100644 --- a/tests/unit/price-sync/cloud-price-updater.test.ts +++ b/tests/unit/price-sync/cloud-price-updater.test.ts @@ -20,10 +20,26 @@ vi.mock("@/lib/async-task-manager", () => { return { AsyncTaskManager: { getActiveTasks: vi.fn(() => []), - register: vi.fn((_taskId: string, promise: Promise) => { - asyncTasks.push(promise); - return new AbortController(); - }), + register: vi.fn( + ( + _taskId: string, + factory: (signal: AbortSignal) => Promise, + options?: string | { abortController?: AbortController } + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + let promise: Promise; + try { + promise = Promise.resolve(factory(controller.signal)); + } catch (error) { + promise = Promise.reject(error); + } + asyncTasks.push(promise); + return controller; + } + ), }, }; }); diff --git a/tests/unit/proxy/build-request-details-redaction.test.ts b/tests/unit/proxy/build-request-details-redaction.test.ts index 03d91fa7f..3f28ba758 100644 --- a/tests/unit/proxy/build-request-details-redaction.test.ts +++ b/tests/unit/proxy/build-request-details-redaction.test.ts @@ -214,5 +214,20 @@ describe("buildRequestDetails - Redaction based on STORE_SESSION_MESSAGES", () = expect(result.headers).toBe("authorization: Bearer [REDACTED]"); expect(result.headers).not.toContain("stale-token"); }); + + it("should remove reserved internal transport headers", () => { + const headers = new Headers([ + ["x-cch-internal-secret", "ws-secret-canary"], + ["x-cch-responses-ws-forward", "1"], + ["x-benign", "visible"], + ]); + const session = createMockSessionWithHeaders("{}", headers, ""); + + const result = buildRequestDetails(session); + + expect(result.headers).toBe("x-benign: visible"); + expect(result.headers).not.toContain("x-cch-"); + expect(result.headers).not.toContain("ws-secret-canary"); + }); }); }); diff --git a/tests/unit/proxy/client-abort-vs-upstream-499.test.ts b/tests/unit/proxy/client-abort-vs-upstream-499.test.ts index 326ecd0fa..501516f92 100644 --- a/tests/unit/proxy/client-abort-vs-upstream-499.test.ts +++ b/tests/unit/proxy/client-abort-vs-upstream-499.test.ts @@ -13,6 +13,7 @@ import { categorizeErrorAsync, isClientAbortError, } from "@/app/v1/_lib/proxy/errors"; +import { DbPoolAdmissionError } from "@/drizzle/admitted-client"; describe("isClientAbortError - 499 source awareness", () => { // Scenario 1: Local abort (isLocalAbort=true) -> CLIENT_ABORT @@ -102,6 +103,17 @@ describe("categorizeErrorAsync - 499 source awareness", () => { }); }); +describe("categorizeErrorAsync - local database overload", () => { + it("should classify a wrapped DB admission rejection separately from network errors", async () => { + const error = new Error("Failed query", { + cause: new DbPoolAdmissionError("data", 32), + }); + + const category = await categorizeErrorAsync(error); + expect(ErrorCategory[category]).toBe("LOCAL_OVERLOAD"); + }); +}); + describe("ProxyError.fromUpstreamResponse - isLocalAbort default", () => { // Scenario 9: fromUpstreamResponse should produce isLocalAbort=false it("should create ProxyError with isLocalAbort=false from upstream 499 response", async () => { diff --git a/tests/unit/proxy/client-detector.test.ts b/tests/unit/proxy/client-detector.test.ts index d14dabc6f..ee43c028e 100644 --- a/tests/unit/proxy/client-detector.test.ts +++ b/tests/unit/proxy/client-detector.test.ts @@ -85,13 +85,12 @@ describe("client-detector", () => { expect(isBuiltinKeyword(pattern)).toBe(true); }); - test.each([ - "gemini-cli", - "codex-cli", - "custom-pattern", - ])("should return false for non-builtin keyword: %s", (pattern) => { - expect(isBuiltinKeyword(pattern)).toBe(false); - }); + test.each(["gemini-cli", "codex-cli", "custom-pattern"])( + "should return false for non-builtin keyword: %s", + (pattern) => { + expect(isBuiltinKeyword(pattern)).toBe(false); + } + ); }); describe("confirmClaudeCodeSignals via detectClientFull", () => { diff --git a/tests/unit/proxy/codex-provider-overrides.test.ts b/tests/unit/proxy/codex-provider-overrides.test.ts index f015a7dde..bae89adaf 100644 --- a/tests/unit/proxy/codex-provider-overrides.test.ts +++ b/tests/unit/proxy/codex-provider-overrides.test.ts @@ -259,35 +259,38 @@ describe("Codex 供应商级参数覆写", () => { ], }, ], - ])("当强制 image_generation=true 且%s已声明 namespace 时,allowed_tools 应使用同形引用", (_, request) => { - const provider = { - providerType: "codex", - codexImageGenerationPreference: "true", - }; - const input: Record = { - model: "gpt-5.5", - ...request, - tool_choice: { - type: "allowed_tools", - mode: "auto", - tools: [{ type: "function", name: "lookup_weather" }], - }, - }; + ])( + "当强制 image_generation=true 且%s已声明 namespace 时,allowed_tools 应使用同形引用", + (_, request) => { + const provider = { + providerType: "codex", + codexImageGenerationPreference: "true", + }; + const input: Record = { + model: "gpt-5.5", + ...request, + tool_choice: { + type: "allowed_tools", + mode: "auto", + tools: [{ type: "function", name: "lookup_weather" }], + }, + }; - const output = applyCodexProviderOverrides(provider as any, input); + const output = applyCodexProviderOverrides(provider as any, input); - expect(output.tool_choice).toEqual({ - type: "allowed_tools", - mode: "auto", - tools: [ - { type: "function", name: "lookup_weather" }, - { type: "namespace", name: "image_gen" }, - ], - }); - expect(output.tools).not.toEqual( - expect.arrayContaining([expect.objectContaining({ type: "image_generation" })]) - ); - }); + expect(output.tool_choice).toEqual({ + type: "allowed_tools", + mode: "auto", + tools: [ + { type: "function", name: "lookup_weather" }, + { type: "namespace", name: "image_gen" }, + ], + }); + expect(output.tools).not.toEqual( + expect.arrayContaining([expect.objectContaining({ type: "image_generation" })]) + ); + } + ); it("当强制 image_generation=false 时,应从 tools 中移除对应工具", () => { const provider = { @@ -425,27 +428,30 @@ describe("Codex 供应商级参数覆写", () => { ["字符串", "image_generation", "image_generation"], ["namespace 字段", { type: "namespace", namespace: "image_gen" }, "namespace:image_gen"], ["嵌套 tool", { tool: { type: "namespace", name: "image_gen" } }, "tool:image_generation"], - ])("当强制 image_generation=false 时,应移除%s形式的 tool_choice 并记录审计", (_, toolChoice, auditValue) => { - const provider = { - providerType: "codex", - codexImageGenerationPreference: "false", - }; - const input: Record = { - model: "gpt-5.5", - input: [], - tool_choice: toolChoice, - }; - - const result = applyCodexProviderOverridesWithAudit(provider as any, input); - - expect(result.request.tool_choice).toBeUndefined(); - expect(result.audit?.changes.find((change) => change.path === "tool_choice")).toEqual({ - path: "tool_choice", - before: auditValue, - after: null, - changed: true, - }); - }); + ])( + "当强制 image_generation=false 时,应移除%s形式的 tool_choice 并记录审计", + (_, toolChoice, auditValue) => { + const provider = { + providerType: "codex", + codexImageGenerationPreference: "false", + }; + const input: Record = { + model: "gpt-5.5", + input: [], + tool_choice: toolChoice, + }; + + const result = applyCodexProviderOverridesWithAudit(provider as any, input); + + expect(result.request.tool_choice).toBeUndefined(); + expect(result.audit?.changes.find((change) => change.path === "tool_choice")).toEqual({ + path: "tool_choice", + before: auditValue, + after: null, + changed: true, + }); + } + ); it("不应把名为 image_generation 的普通函数选择误判为内置图片工具", () => { const provider = { diff --git a/tests/unit/proxy/connected-non-reader-lifetime.test.ts b/tests/unit/proxy/connected-non-reader-lifetime.test.ts new file mode 100644 index 000000000..f41e2a7c6 --- /dev/null +++ b/tests/unit/proxy/connected-non-reader-lifetime.test.ts @@ -0,0 +1,258 @@ +import { Context } from "hono"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { Readable } from "node:stream"; +import { + createDemandDrivenResponsePump, + type DemandDrivenResponsePump, + type DemandDrivenResponsePumpCompletion, +} from "@/app/v1/_lib/proxy/demand-driven-response-pump"; +import { ProxyForwarder } from "@/app/v1/_lib/proxy/forwarder"; +import { ProxyProviderResolver } from "@/app/v1/_lib/proxy/provider-selector"; +import { ProxyResponseHandler } from "@/app/v1/_lib/proxy/response-handler"; +import { ProxySession } from "@/app/v1/_lib/proxy/session"; + +const encoder = new TextEncoder(); +const processingTasks: Promise[] = []; +const transportMocks = vi.hoisted(() => ({ request: vi.fn() })); + +vi.mock("undici", async (importOriginal) => ({ + ...(await importOriginal()), + request: transportMocks.request, +})); + +vi.mock("@/app/v1/_lib/proxy/response-fixer", () => ({ + ResponseFixer: { process: async (_session: ProxySession, response: Response) => response }, +})); +vi.mock("@/lib/async-task-manager", () => ({ + AsyncTaskManager: { + register: vi.fn((_id: string, factory: () => Promise) => { + const task = factory(); + processingTasks.push(task); + return new AbortController(); + }), + touch: vi.fn(), + }, +})); +vi.mock("@/lib/logger", () => ({ + logger: { debug: vi.fn(), error: vi.fn(), info: vi.fn(), trace: vi.fn(), warn: vi.fn() }, +})); +vi.mock("@/lib/circuit-breaker", () => ({ + getCircuitState: vi.fn(() => "closed"), + getProviderHealthInfo: vi.fn(async () => ({ + health: { failureCount: 0 }, + config: { failureThreshold: 3 }, + })), + recordFailure: vi.fn(), + recordSuccess: vi.fn(), +})); +vi.mock("@/lib/endpoint-circuit-breaker", () => ({ + recordEndpointFailure: vi.fn(), +})); +vi.mock("@/repository/message", () => ({ + updateMessageRequestCostWithBreakdown: vi.fn(), + updateMessageRequestDetails: vi.fn(), + updateMessageRequestDetailsDurably: vi.fn(), + updateMessageRequestDetailsIfUnfinalized: vi.fn(), + updateMessageRequestDuration: vi.fn(), +})); +vi.mock("@/lib/session-manager", () => ({ + SessionManager: { + clearSessionProvider: vi.fn(), + updateSessionBindingSmart: vi.fn(async () => ({ reason: "test", updated: false })), + updateSessionUsage: vi.fn(), + }, +})); + +async function createGeminiSession(signal: AbortSignal | null): Promise { + const request = new Request("https://example.com/v1/chat/completions", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ model: "gemini-2.0-flash", stream: true }), + signal, + }); + const session = await ProxySession.fromContext(new Context(request)); + Object.assign(session, { + authState: { apiKey: null, key: null, success: true, user: null }, + provider: { + id: 1, + name: "gemini", + url: "https://example.com", + key: "test-key", + providerType: "gemini", + firstByteTimeoutStreamingMs: 0, + streamingIdleTimeoutMs: 0, + }, + messageContext: { createdAt: new Date(), id: 1, user: { id: 1, name: "test" } }, + originalFormat: "gemini", + }); + return session; +} + +afterEach(() => { + vi.useRealTimers(); + processingTasks.length = 0; +}); + +describe("connected non-reader response lifetime", () => { + it("cancels one unconsumed lookahead chunk at the 60 second deadline", async () => { + vi.useFakeTimers(); + const cancel = vi.fn(); + const source = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode("pending")); + }, + cancel, + }); + const completions: DemandDrivenResponsePumpCompletion[] = []; + const deadlineError = expect.objectContaining({ name: "AbortError" }); + + const pump = createDemandDrivenResponsePump({ source, onChunk: vi.fn() }); + void pump.completion.then((completion) => completions.push(completion)); + await vi.advanceTimersByTimeAsync(59_999); + + expect(cancel).not.toHaveBeenCalled(); + expect(completions).toEqual([]); + expect(pump.getState()).toBe("client-active"); + + await vi.advanceTimersByTimeAsync(1); + const completion = await pump.completion; + pump.startDrain(new Error("late drain")); + pump.cancelSource(new Error("late cancel")); + await vi.runAllTimersAsync(); + + expect(cancel).toHaveBeenCalledOnce(); + expect(cancel).toHaveBeenCalledWith(deadlineError); + expect(completion).toMatchObject({ streamEndedNormally: false, clientAborted: false }); + expect(completion.error).toEqual(deadlineError); + expect(completions).toEqual([completion]); + expect(pump.getState()).toBe("closed"); + }); + + it("transfers cancellation even when onClientCancel throws", async () => { + let sourceController: ReadableStreamDefaultController | null = null; + const callbackError = new Error("cancel observer failed"); + const source = new ReadableStream({ + start(controller) { + sourceController = controller; + }, + }); + const pump = createDemandDrivenResponsePump({ + source, + onChunk: vi.fn(), + onClientCancel() { + throw callbackError; + }, + }); + + const cancelOutcome = await pump.stream.cancel("client disconnected").then( + () => ({ kind: "resolved" as const }), + (error: unknown) => ({ kind: "rejected" as const, error }) + ); + const stateAfterCancel = pump.getState(); + sourceController?.close(); + const completion = await pump.completion; + + expect(cancelOutcome).toEqual({ kind: "rejected", error: callbackError }); + expect(stateAfterCancel).toBe("draining"); + expect(completion).toMatchObject({ streamEndedNormally: true, clientAborted: true }); + expect(completion.error).toBeNull(); + }); + + it("preserves the first hard-cancel owner during synchronous source reentry", async () => { + const firstError = new DOMException("pending response was not consumed", "AbortError"); + const reentrantError = new Error("reentrant cancel"); + const cancelFailure = new Error("source cancel failed"); + let pump: DemandDrivenResponsePump | null = null; + const cancel = vi.fn(() => { + pump?.cancelSource(reentrantError); + return Promise.reject(cancelFailure); + }); + const source = new ReadableStream({ cancel }); + pump = createDemandDrivenResponsePump({ source, onChunk: vi.fn() }); + const completions: DemandDrivenResponsePumpCompletion[] = []; + void pump.completion.then((completion) => completions.push(completion)); + + pump.cancelSource(firstError); + const completion = await pump.completion; + await new Promise((resolve) => setImmediate(resolve)); + + expect(cancel).toHaveBeenCalledOnce(); + expect(completion.error).toBe(firstError); + expect(firstError.cause).toBe(cancelFailure); + expect(completions).toEqual([completion]); + expect(pump.getState()).toBe("closed"); + }); + + it("keeps Gemini passthrough demand-driven while preserving exact chunks", async () => { + const chunks = [ + '{"candidates":[{"content":{"parts":[{"text":"one"}]}}]}\n', + '{"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":3}}\n', + ]; + let pullCount = 0; + const source = new ReadableStream( + { + pull(controller) { + pullCount += 1; + const chunk = chunks[pullCount - 1]; + if (chunk) controller.enqueue(encoder.encode(chunk)); + else controller.close(); + }, + }, + { highWaterMark: 0 } + ); + const response = new Response(source, { headers: { "content-type": "text/event-stream" } }); + const session = await createGeminiSession(null); + + const returned = await ProxyResponseHandler.dispatch(session, response); + await Promise.all([Promise.resolve(), Promise.resolve()]); + + expect(pullCount).toBe(1); + await expect(returned.text()).resolves.toBe(chunks.join("")); + const settlements = await Promise.allSettled(processingTasks); + expect(settlements.every((settlement) => settlement.status === "fulfilled")).toBe(true); + }); + + it.each([true, false])( + "detaches client cancellation after headers with signal=%s", + async (hasClientSignal) => { + const clientController = new AbortController(); + const session = await createGeminiSession(hasClientSignal ? clientController.signal : null); + let transportSignal: AbortSignal | undefined; + transportMocks.request.mockImplementation(async (_url, options) => { + transportSignal = options.signal; + return { + statusCode: 200, + headers: { "content-type": "text/event-stream" }, + body: Readable.from(["data: {}\n\n"]), + }; + }); + + const response = await ProxyForwarder.send(session); + clientController.abort(new Error("client disconnected after headers")); + expect(transportSignal?.aborted).toBe(false); + await response.body?.cancel(); + } + ); + + it("detaches transport signals after an upstream error response", async () => { + const clientController = new AbortController(); + const clientError = new Error("client abort before headers"); + const session = await createGeminiSession(clientController.signal); + vi.spyOn(ProxyProviderResolver, "pickRandomProviderWithExclusion").mockResolvedValue(null); + let transportSignal: AbortSignal | undefined; + transportMocks.request.mockImplementation(async (_url, options) => { + transportSignal = options.signal; + clientController.abort(clientError); + return { + statusCode: 499, + headers: { "content-type": "application/json" }, + body: Readable.from(["{}"]), + }; + }); + + await expect(ProxyForwarder.send(session)).rejects.toMatchObject({ statusCode: 503 }); + + expect(transportSignal?.aborted).toBe(true); + expect(transportSignal?.reason).toBe(clientError); + }); +}); diff --git a/tests/unit/proxy/endpoint-family-catalog.test.ts b/tests/unit/proxy/endpoint-family-catalog.test.ts index 428fe937c..efd8ea69c 100644 --- a/tests/unit/proxy/endpoint-family-catalog.test.ts +++ b/tests/unit/proxy/endpoint-family-catalog.test.ts @@ -352,11 +352,12 @@ describe("endpoint family catalog", () => { expect(resolveEndpointFamilyByPath(path)?.modelRequired).toBe(true); }); - test.each(FAMILY_SAMPLES.filter((entry) => !entry.modelRequired))("%s 不应要求模型", ({ - path, - }) => { - expect(resolveEndpointFamilyByPath(path)?.modelRequired).toBe(false); - }); + test.each(FAMILY_SAMPLES.filter((entry) => !entry.modelRequired))( + "%s 不应要求模型", + ({ path }) => { + expect(resolveEndpointFamilyByPath(path)?.modelRequired).toBe(false); + } + ); test("Gemini batch body fallback 应识别为 gemini", () => { expect( diff --git a/tests/unit/proxy/endpoint-family-provider-routing.test.ts b/tests/unit/proxy/endpoint-family-provider-routing.test.ts index 14d121fb1..4722353fd 100644 --- a/tests/unit/proxy/endpoint-family-provider-routing.test.ts +++ b/tests/unit/proxy/endpoint-family-provider-routing.test.ts @@ -385,32 +385,31 @@ describe("endpoint family -> provider routing matrix", () => { ); }); - test.each(ENDPOINT_PROVIDER_CASES)("$id should route $path to $expectedProviderType", async ({ - path, - expectedProviderType, - requestedModel, - }) => { - const ProxyProviderResolver = await setupResolverMocks(); + test.each(ENDPOINT_PROVIDER_CASES)( + "$id should route $path to $expectedProviderType", + async ({ path, expectedProviderType, requestedModel }) => { + const ProxyProviderResolver = await setupResolverMocks(); - const providers: Provider[] = [ - createTestProvider(1, "claude"), - createTestProvider(2, "claude-auth"), - createTestProvider(3, "codex"), - createTestProvider(4, "openai-compatible"), - createTestProvider(5, "gemini"), - createTestProvider(6, "gemini-cli"), - ]; - const session = createSessionStub(path, requestedModel); - session.getProvidersSnapshot = async () => providers; + const providers: Provider[] = [ + createTestProvider(1, "claude"), + createTestProvider(2, "claude-auth"), + createTestProvider(3, "codex"), + createTestProvider(4, "openai-compatible"), + createTestProvider(5, "gemini"), + createTestProvider(6, "gemini-cli"), + ]; + const session = createSessionStub(path, requestedModel); + session.getProvidersSnapshot = async () => providers; - const { provider, context } = await (ProxyProviderResolver as any).pickRandomProvider( - session, - [] - ); + const { provider, context } = await (ProxyProviderResolver as any).pickRandomProvider( + session, + [] + ); - expect(provider?.providerType).toBe(expectedProviderType); - expect(context.requestedModel).toBe(requestedModel); - }); + expect(provider?.providerType).toBe(expectedProviderType); + expect(context.requestedModel).toBe(requestedModel); + } + ); test("/v1/chat/completions should never select codex when openai-compatible is available", async () => { const ProxyProviderResolver = await setupResolverMocks(); diff --git a/tests/unit/proxy/endpoint-path-normalization.test.ts b/tests/unit/proxy/endpoint-path-normalization.test.ts index 8b4662e04..183585a7b 100644 --- a/tests/unit/proxy/endpoint-path-normalization.test.ts +++ b/tests/unit/proxy/endpoint-path-normalization.test.ts @@ -38,17 +38,15 @@ describe("endpoint path normalization", () => { expect(isRawPassthroughEndpointPath(pathname)).toBe(true); }); - test.each([ - "/v1/messages", - "/v1/responses", - "/v1/messages/count", - "/v1/responses/mini", - ])("non-target path is not misclassified for %s", (pathname) => { - expect(isCountTokensEndpointPath(pathname)).toBe(false); - expect(isResponseCompactEndpointPath(pathname)).toBe(false); - expect(isRawPassthroughEndpointPath(pathname)).toBe(false); - expect(isCountTokensRequestWithEndpoint(pathname)).toBe(false); - }); + test.each(["/v1/messages", "/v1/responses", "/v1/messages/count", "/v1/responses/mini"])( + "non-target path is not misclassified for %s", + (pathname) => { + expect(isCountTokensEndpointPath(pathname)).toBe(false); + expect(isResponseCompactEndpointPath(pathname)).toBe(false); + expect(isRawPassthroughEndpointPath(pathname)).toBe(false); + expect(isCountTokensRequestWithEndpoint(pathname)).toBe(false); + } + ); test("session count_tokens detection handles null endpoint", () => { expect(isCountTokensRequestWithEndpoint(null)).toBe(false); diff --git a/tests/unit/proxy/error-handler-client-message.test.ts b/tests/unit/proxy/error-handler-client-message.test.ts new file mode 100644 index 000000000..3c5ba79b9 --- /dev/null +++ b/tests/unit/proxy/error-handler-client-message.test.ts @@ -0,0 +1,237 @@ +import { Context } from "hono"; +import { DrizzleQueryError } from "drizzle-orm"; +import { beforeEach, describe, expect, test, vi } from "vitest"; +import { + ProxyErrorHandler, + resolveFinalClientErrorMessage, +} from "@/app/v1/_lib/proxy/error-handler"; +import { ProxyError } from "@/app/v1/_lib/proxy/errors"; +import { ProxySession } from "@/app/v1/_lib/proxy/session"; +import { DbPoolAdmissionError } from "@/drizzle/admitted-client"; +import type { ErrorDetectionResult } from "@/lib/error-rule-detector"; + +const mocks = vi.hoisted(() => ({ + detectAsync: vi.fn<(content: string) => Promise>(), + emitProxyLangfuseTrace: vi.fn(), + getCachedSystemSettings: vi.fn(async () => ({ + verboseProviderError: false, + passThroughUpstreamErrorMessage: true, + })), +})); + +vi.mock("@/lib/error-rule-detector", () => ({ + errorRuleDetector: { detectAsync: mocks.detectAsync }, +})); + +vi.mock("@/lib/config/system-settings-cache", () => ({ + getCachedSystemSettings: mocks.getCachedSystemSettings, +})); + +vi.mock("@/lib/langfuse/emit-proxy-trace", () => ({ + emitProxyLangfuseTrace: mocks.emitProxyLangfuseTrace, +})); + +vi.mock("@/lib/logger", () => ({ + logger: { + debug: vi.fn(), + error: vi.fn(), + fatal: vi.fn(), + info: vi.fn(), + trace: vi.fn(), + warn: vi.fn(), + }, +})); + +async function createSession(): Promise { + const request = new Request("https://hub.test/v1/messages", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + model: "claude-sonnet-4-20250514", + messages: [{ role: "user", content: "hello" }], + }), + }); + return ProxySession.fromContext(new Context(request)); +} + +describe("resolveFinalClientErrorMessage", () => { + test("sanitizes a pass-through message extracted from an upstream body", () => { + const error = new ProxyError("Quota exceeded", 429, { + body: "Quota exceeded for api_key=[REDACTED_KEY]", + rawBody: JSON.stringify({ + error: { + message: + "Quota exceeded for api_key=sk-secret-12345678 at https://api.vendor.example/v1 request_id=req_abc123", + }, + }), + providerName: "provider-a", + }); + + const message = resolveFinalClientErrorMessage({ + error, + currentFallbackMessage: "Quota exceeded Upstream detail: provider-a", + settings: { passThroughUpstreamErrorMessage: true }, + override: null, + }); + + expect(message).toContain("Quota exceeded"); + expect(message).not.toContain("sk-secret"); + expect(message).not.toContain("https://"); + expect(message).not.toContain("api.vendor.example"); + expect(message).not.toContain("req_abc123"); + expect(message).not.toContain("provider-a"); + }); + + test("uses a safe candidate when the raw body has no extractable error", () => { + const error = new ProxyError("Rate limit", 429, { + body: "Rate limit exceeded for this endpoint", + rawBody: "temporary plain-text response", + safeClientMessageCandidate: "Rate limit exceeded for this endpoint", + providerName: "relay-a", + }); + + const message = resolveFinalClientErrorMessage({ + error, + currentFallbackMessage: "Rate limit", + settings: { passThroughUpstreamErrorMessage: true }, + override: null, + }); + + expect(message).toBe("Rate limit exceeded for this endpoint"); + }); + + test.each([ + [400, "上游请求参数无效,请检查后重试"], + [401, "上游鉴权失败,请稍后重试"], + [429, "上游服务当前限流,请稍后重试"], + [503, "上游服务暂时不可用,请稍后重试"], + ])("maps status %i to a generic message when pass-through is disabled", (status, expected) => { + const error = new ProxyError("sensitive upstream failure", status, { + body: "sensitive upstream failure", + }); + + const message = resolveFinalClientErrorMessage({ + error, + currentFallbackMessage: "sensitive upstream failure Upstream detail: relay-a", + settings: { passThroughUpstreamErrorMessage: false }, + override: null, + }); + + expect(message).toBe(expected); + }); + + test("falls back when every upstream candidate exposes a provider label", () => { + const error = new ProxyError("Provider relay-a returned: overload", 503, { + body: "Provider relay-a returned: overload", + safeClientMessageCandidate: "Provider relay-a returned: overload", + providerName: "relay-a", + }); + + const message = resolveFinalClientErrorMessage({ + error, + currentFallbackMessage: "Provider relay-a returned: overload", + settings: { passThroughUpstreamErrorMessage: true }, + override: null, + }); + + expect(message).toBe("上游服务暂时不可用,请稍后重试"); + }); + + test("preserves an explicit override message", () => { + const message = resolveFinalClientErrorMessage({ + error: new ProxyError("Upstream failed", 502, { body: "Upstream failed" }), + currentFallbackMessage: "custom override", + settings: { passThroughUpstreamErrorMessage: false }, + override: { + statusCode: 451, + response: { error: { type: "invalid_request_error", message: "custom override" } }, + }, + }); + + expect(message).toBe("custom override"); + }); +}); + +describe("ProxyErrorHandler.handle client message", () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.detectAsync.mockResolvedValue({ matched: false }); + mocks.getCachedSystemSettings.mockResolvedValue({ + verboseProviderError: false, + passThroughUpstreamErrorMessage: true, + }); + }); + + test("returns only the sanitized upstream message to the client", async () => { + const session = await createSession(); + session.setSessionId("s_client_message"); + const error = new ProxyError("Quota exceeded", 429, { + body: "Quota exceeded for api_key=[REDACTED_KEY]", + rawBody: JSON.stringify({ + error: { + message: + "Quota exceeded for api_key=sk-secret-12345678 at https://api.vendor.example/v1 request_id=req_abc123", + }, + }), + providerName: "provider-a", + }); + + const response = await ProxyErrorHandler.handle(session, error); + const responseText = await response.text(); + + expect(response.status).toBe(429); + expect(responseText).toContain("Quota exceeded"); + expect(responseText).toContain("cch_session_id: s_client_message"); + expect(responseText).not.toContain("sk-secret"); + expect(responseText).not.toContain("https://"); + expect(responseText).not.toContain("req_abc123"); + expect(responseText).not.toContain("provider-a"); + }); + + test("returns a fixed 503 without exposing an admission query or parameters", async () => { + const session = await createSession(); + const canary = "sk-admission-secret-canary"; + const error = new DrizzleQueryError( + "select * from keys where key = $1", + [canary], + new DbPoolAdmissionError("control", 32) + ); + + const response = await ProxyErrorHandler.handle(session, error); + const responseText = await response.text(); + + expect(response.status).toBe(503); + expect(responseText).not.toContain(canary); + expect(responseText).not.toContain("select * from keys"); + expect(responseText).not.toContain("params:"); + expect(mocks.getCachedSystemSettings).not.toHaveBeenCalled(); + expect(mocks.emitProxyLangfuseTrace).toHaveBeenCalledWith( + session, + expect.objectContaining({ + errorMessage: expect.not.stringContaining(canary), + statusCode: 503, + }) + ); + }); + + test("sanitizes ordinary Drizzle query failures before HTTP and Langfuse", async () => { + const session = await createSession(); + const canary = "sk-query-secret-canary"; + const cause = Object.assign(new Error("canceling statement due to lock timeout"), { + code: "55P03", + }); + const error = new DrizzleQueryError("update keys set key = $1", [canary], cause); + + const response = await ProxyErrorHandler.handle(session, error); + const responseText = await response.text(); + const traceArguments = JSON.stringify(mocks.emitProxyLangfuseTrace.mock.calls); + + expect(response.status).toBe(500); + expect(responseText).not.toContain(canary); + expect(responseText).not.toContain("update keys"); + expect(responseText).not.toContain("params:"); + expect(traceArguments).not.toContain(canary); + expect(traceArguments).not.toContain("update keys"); + expect(mocks.getCachedSystemSettings).not.toHaveBeenCalled(); + }); +}); diff --git a/tests/unit/proxy/error-handler-durable-persistence.test.ts b/tests/unit/proxy/error-handler-durable-persistence.test.ts new file mode 100644 index 000000000..803a294d8 --- /dev/null +++ b/tests/unit/proxy/error-handler-durable-persistence.test.ts @@ -0,0 +1,233 @@ +import { Context } from "hono"; +import { DrizzleQueryError } from "drizzle-orm"; +import { beforeEach, describe, expect, test, vi } from "vitest"; +import { ProxyErrorHandler } from "@/app/v1/_lib/proxy/error-handler"; +import { ProxyError } from "@/app/v1/_lib/proxy/errors"; +import { ProxySession } from "@/app/v1/_lib/proxy/session"; +import { DbPoolAdmissionError } from "@/drizzle/admitted-client"; +import type { ErrorDetectionResult } from "@/lib/error-rule-detector"; +import type { emitProxyLangfuseTrace } from "@/lib/langfuse/emit-proxy-trace"; +import type { updateMessageRequestDetailsDurably } from "@/repository/message"; +import type { Key } from "@/types/key"; +import type { User } from "@/types/user"; + +const mocks = vi.hoisted(() => ({ + detectAsync: vi.fn<(content: string) => Promise>(), + emitProxyLangfuseTrace: vi.fn(), + endRequest: vi.fn(), + getCachedSystemSettings: vi.fn(async () => ({ + verboseProviderError: false, + passThroughUpstreamErrorMessage: false, + })), + updateMessageRequestDetailsDurably: vi.fn(), +})); + +vi.mock("@/lib/error-rule-detector", () => ({ + errorRuleDetector: { detectAsync: mocks.detectAsync }, +})); + +vi.mock("@/lib/config/system-settings-cache", () => ({ + getCachedSystemSettings: mocks.getCachedSystemSettings, +})); + +vi.mock("@/lib/langfuse/emit-proxy-trace", () => ({ + emitProxyLangfuseTrace: mocks.emitProxyLangfuseTrace, +})); + +vi.mock("@/repository/message", () => ({ + updateMessageRequestDetailsDurably: mocks.updateMessageRequestDetailsDurably, +})); + +vi.mock("@/lib/proxy-status-tracker", () => ({ + ProxyStatusTracker: { + getInstance: () => ({ endRequest: mocks.endRequest }), + }, +})); + +vi.mock("@/lib/logger", () => ({ + logger: { + debug: vi.fn(), + error: vi.fn(), + fatal: vi.fn(), + info: vi.fn(), + trace: vi.fn(), + warn: vi.fn(), + }, +})); + +const FIXTURE_DATE = new Date("2026-01-01T00:00:00.000Z"); + +const USER = { + id: 42, + name: "test-user", + description: "error-handler fixture", + role: "user", + rpm: null, + dailyQuota: null, + providerGroup: null, + createdAt: FIXTURE_DATE, + updatedAt: FIXTURE_DATE, + limit5hResetMode: "fixed", + dailyResetMode: "fixed", + dailyResetTime: "00:00", + isEnabled: true, +} satisfies User; + +const KEY = { + id: 8, + userId: USER.id, + name: "test-key", + key: "sk-test-key", + isEnabled: true, + canLoginWebUi: false, + limit5hUsd: null, + limit5hResetMode: "fixed", + limitDailyUsd: null, + dailyResetMode: "fixed", + dailyResetTime: "00:00", + limitWeeklyUsd: null, + limitMonthlyUsd: null, + limitConcurrentSessions: 0, + providerGroup: null, + cacheTtlPreference: null, + createdAt: FIXTURE_DATE, + updatedAt: FIXTURE_DATE, +} satisfies Key; + +async function createSession(): Promise { + const request = new Request("https://hub.test/v1/messages", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + model: "claude-sonnet-4-20250514", + messages: [{ role: "user", content: "hello" }], + }), + }); + const session = await ProxySession.fromContext(new Context(request)); + session.setSessionId("s_durable"); + return session; +} + +function attachMessageContext(session: ProxySession): void { + session.setMessageContext({ + id: 901, + createdAt: FIXTURE_DATE, + user: USER, + key: KEY, + apiKey: KEY.key, + }); +} + +describe("ProxyErrorHandler.handle durable persistence", () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.detectAsync.mockResolvedValue({ matched: false }); + mocks.getCachedSystemSettings.mockResolvedValue({ + verboseProviderError: false, + passThroughUpstreamErrorMessage: false, + }); + mocks.updateMessageRequestDetailsDurably.mockResolvedValue(undefined); + }); + + test("emits the trace, awaits persistence, then ends status tracking", async () => { + const commit = Promise.withResolvers(); + mocks.updateMessageRequestDetailsDurably.mockReturnValueOnce(commit.promise); + const session = await createSession(); + attachMessageContext(session); + + const handlePromise = ProxyErrorHandler.handle(session, new Error("fetch failed")); + await vi.waitFor(() => expect(mocks.updateMessageRequestDetailsDurably).toHaveBeenCalledOnce()); + + expect(mocks.emitProxyLangfuseTrace).toHaveBeenCalledOnce(); + expect(mocks.endRequest).not.toHaveBeenCalled(); + const traceOrder = mocks.emitProxyLangfuseTrace.mock.invocationCallOrder[0]; + const persistOrder = mocks.updateMessageRequestDetailsDurably.mock.invocationCallOrder[0]; + expect(traceOrder ?? Number.MAX_SAFE_INTEGER).toBeLessThan(persistOrder ?? -1); + + commit.resolve(); + const response = await handlePromise; + + expect(response.status).toBe(500); + expect(mocks.updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 901, + expect.objectContaining({ + durationMs: expect.any(Number), + errorMessage: "fetch failed", + providerChain: [], + statusCode: 500, + model: "claude-sonnet-4-20250514", + context1mApplied: false, + swapCacheTtlApplied: false, + }) + ); + expect(mocks.endRequest).toHaveBeenCalledWith(USER.id, 901); + const endOrder = mocks.endRequest.mock.invocationCallOrder[0]; + expect(persistOrder ?? Number.MAX_SAFE_INTEGER).toBeLessThan(endOrder ?? -1); + }); + + test("keeps the trace when durable persistence rejects", async () => { + mocks.updateMessageRequestDetailsDurably.mockRejectedValueOnce(new Error("db down")); + const session = await createSession(); + attachMessageContext(session); + + await expect(ProxyErrorHandler.handle(session, new Error("fetch failed"))).rejects.toThrow( + "db down" + ); + + expect(mocks.emitProxyLangfuseTrace).toHaveBeenCalledWith( + session, + expect.objectContaining({ statusCode: 500, errorMessage: "fetch failed" }) + ); + expect(mocks.endRequest).not.toHaveBeenCalled(); + }); + + test("persists the final overridden status", async () => { + mocks.detectAsync.mockResolvedValue({ matched: true, overrideStatusCode: 429 }); + const session = await createSession(); + attachMessageContext(session); + const error = new ProxyError("Upstream failed", 502, { + body: "Upstream failed", + providerId: 7, + providerName: "provider-a", + }); + + const response = await ProxyErrorHandler.handle(session, error); + + expect(response.status).toBe(429); + expect(mocks.updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 901, + expect.objectContaining({ + statusCode: 429, + errorMessage: expect.stringContaining("Upstream"), + }) + ); + expect(mocks.endRequest).toHaveBeenCalledWith(USER.id, 901); + }); + + test("skips persistence and tracking when no message context exists", async () => { + const session = await createSession(); + + const response = await ProxyErrorHandler.handle(session, new Error("fetch failed")); + + expect(response.status).toBe(500); + expect(mocks.emitProxyLangfuseTrace).toHaveBeenCalledOnce(); + expect(mocks.updateMessageRequestDetailsDurably).not.toHaveBeenCalled(); + expect(mocks.endRequest).not.toHaveBeenCalled(); + }); + + test("returns admission 503 without recursively waiting on durable persistence", async () => { + const session = await createSession(); + attachMessageContext(session); + const error = new DrizzleQueryError( + "select * from keys where key = $1", + ["sk-admission-canary"], + new DbPoolAdmissionError("control", 32) + ); + + const response = await ProxyErrorHandler.handle(session, error); + + expect(response.status).toBe(503); + expect(mocks.updateMessageRequestDetailsDurably).not.toHaveBeenCalled(); + expect(mocks.endRequest).toHaveBeenCalledWith(USER.id, 901); + }); +}); diff --git a/tests/unit/proxy/error-handler-langfuse-trace.test.ts b/tests/unit/proxy/error-handler-langfuse-trace.test.ts index 80994a904..2f37a3471 100644 --- a/tests/unit/proxy/error-handler-langfuse-trace.test.ts +++ b/tests/unit/proxy/error-handler-langfuse-trace.test.ts @@ -1,15 +1,22 @@ +import { Context } from "hono"; import { beforeEach, describe, expect, test, vi } from "vitest"; +import { ProxyErrorHandler } from "@/app/v1/_lib/proxy/error-handler"; +import { ProxyError } from "@/app/v1/_lib/proxy/errors"; +import { ProxySession } from "@/app/v1/_lib/proxy/session"; +import type { ErrorDetectionResult } from "@/lib/error-rule-detector"; +import type { emitProxyLangfuseTrace } from "@/lib/langfuse/emit-proxy-trace"; const mocks = vi.hoisted(() => ({ - emitProxyLangfuseTrace: vi.fn(), + detectAsync: vi.fn<(content: string) => Promise>(), + emitProxyLangfuseTrace: vi.fn(), getCachedSystemSettings: vi.fn(async () => ({ verboseProviderError: false, passThroughUpstreamErrorMessage: false, })), - getErrorOverrideAsync: vi.fn(async () => undefined), - updateMessageRequestDetails: vi.fn(async () => undefined), - updateMessageRequestDuration: vi.fn(async () => undefined), - endRequest: vi.fn(), +})); + +vi.mock("@/lib/error-rule-detector", () => ({ + errorRuleDetector: { detectAsync: mocks.detectAsync }, })); vi.mock("@/lib/langfuse/emit-proxy-trace", () => ({ @@ -20,92 +27,47 @@ vi.mock("@/lib/config/system-settings-cache", () => ({ getCachedSystemSettings: mocks.getCachedSystemSettings, })); -vi.mock("@/repository/message", () => ({ - updateMessageRequestDetails: mocks.updateMessageRequestDetails, - updateMessageRequestDuration: mocks.updateMessageRequestDuration, -})); - -vi.mock("@/lib/proxy-status-tracker", () => ({ - ProxyStatusTracker: { - getInstance: () => ({ - endRequest: mocks.endRequest, - }), - }, -})); - vi.mock("@/lib/logger", () => ({ logger: { debug: vi.fn(), - info: vi.fn(), - warn: vi.fn(), - trace: vi.fn(), error: vi.fn(), fatal: vi.fn(), + info: vi.fn(), + trace: vi.fn(), + warn: vi.fn(), }, })); -vi.mock("@/app/v1/_lib/proxy/errors", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - getErrorOverrideAsync: mocks.getErrorOverrideAsync, - }; -}); - -import { ProxyErrorHandler } from "@/app/v1/_lib/proxy/error-handler"; -import { ProxyError, RateLimitError } from "@/app/v1/_lib/proxy/errors"; +type SessionInput = { + readonly stream?: boolean; + readonly url?: string; +}; -function createSession(overrides: Record = {}): any { - const requestMessage = { - model: "claude-sonnet-4-20250514", - messages: [{ role: "user", content: "hello" }], - }; - - return { - sessionId: "s_langfuse_error", - messageContext: { - id: "msg_langfuse_error", - user: { id: 42, name: "test-user" }, - key: { name: "test-key" }, - }, - startTime: Date.now() - 25, +async function createSession(input: SessionInput = {}): Promise { + const request = new Request(input.url ?? "https://hub.test/v1/messages", { method: "POST", - originalFormat: "claude", - request: { - message: requestMessage, - model: requestMessage.model, - log: JSON.stringify(requestMessage), - }, - headers: new Headers({ "user-agent": "vitest" }), - provider: { - id: 7, - name: "provider-a", - providerType: "claude", - swapCacheTtlBilling: false, - }, - getProviderChain: () => [], - getCurrentModel: () => requestMessage.model, - getContext1mApplied: () => false, - getGroupCostMultiplier: () => 1, - getSpecialSettings: () => null, - getEndpoint: () => "/v1/messages", - getRequestSequence: () => 1, - ...overrides, - }; + headers: { "content-type": "application/json", "user-agent": "vitest" }, + body: JSON.stringify({ + model: "claude-sonnet-4-20250514", + messages: [{ role: "user", content: "hello" }], + stream: input.stream ?? false, + }), + }); + return ProxySession.fromContext(new Context(request)); } -describe("ProxyErrorHandler.handle - Langfuse error traces", () => { +describe("ProxyErrorHandler.handle Langfuse traces", () => { beforeEach(() => { vi.clearAllMocks(); + mocks.detectAsync.mockResolvedValue({ matched: false }); mocks.getCachedSystemSettings.mockResolvedValue({ verboseProviderError: false, passThroughUpstreamErrorMessage: false, }); - mocks.getErrorOverrideAsync.mockResolvedValue(undefined); }); - test("emits trace for local request errors without upstream output", async () => { - const session = createSession(); + test("emits an empty-output trace for a local request error", async () => { + const session = await createSession(); await ProxyErrorHandler.handle(session, new ProxyError("Invalid request: missing model", 400)); @@ -121,11 +83,12 @@ describe("ProxyErrorHandler.handle - Langfuse error traces", () => { errorMessage: "Invalid request: missing model", }) ); - expect(mocks.emitProxyLangfuseTrace.mock.calls[0][1].durationMs).toBeGreaterThanOrEqual(0); + const trace = mocks.emitProxyLangfuseTrace.mock.calls[0]?.[1]; + expect(trace?.durationMs).toBeGreaterThanOrEqual(0); }); - test("emits trace for thrown network errors without upstream output", async () => { - const session = createSession(); + test("emits an empty-output trace for a thrown network error", async () => { + const session = await createSession(); await ProxyErrorHandler.handle(session, new Error("fetch failed")); @@ -140,26 +103,8 @@ describe("ProxyErrorHandler.handle - Langfuse error traces", () => { ); }); - test("emits trace before database persistence failures can abort handling", async () => { - const session = createSession(); - mocks.updateMessageRequestDuration.mockRejectedValueOnce(new Error("db down")); - - await expect(ProxyErrorHandler.handle(session, new Error("fetch failed"))).rejects.toThrow( - "db down" - ); - - expect(mocks.emitProxyLangfuseTrace).toHaveBeenCalledWith( - session, - expect.objectContaining({ - responseText: "", - statusCode: 500, - errorMessage: "fetch failed", - }) - ); - }); - - test("uses upstream raw body as trace output when available", async () => { - const session = createSession(); + test("prefers the raw upstream body as trace output", async () => { + const session = await createSession(); const error = new ProxyError("Upstream failed", 502, { body: "sanitized upstream body", rawBody: '{"error":{"message":"raw upstream failure"}}', @@ -180,49 +125,8 @@ describe("ProxyErrorHandler.handle - Langfuse error traces", () => { ); }); - test("emits final override response and status after error override is applied", async () => { - const session = createSession(); - mocks.getErrorOverrideAsync.mockResolvedValueOnce({ - statusCode: 429, - response: { - error: { - message: "masked quota message", - type: "rate_limit_error", - }, - }, - }); - - const response = await ProxyErrorHandler.handle( - session, - new ProxyError("Upstream failed", 502, { - rawBody: '{"error":{"message":"raw upstream failure"}}', - providerId: 7, - providerName: "provider-a", - }) - ); - - expect(response.status).toBe(429); - expect(mocks.emitProxyLangfuseTrace).toHaveBeenCalledWith( - session, - expect.objectContaining({ - responseText: expect.stringContaining("masked quota message"), - statusCode: 429, - errorMessage: "masked quota message", - }) - ); - expect(mocks.emitProxyLangfuseTrace.mock.calls[0][1].responseText).not.toContain( - "raw upstream failure" - ); - expect(mocks.updateMessageRequestDetails).toHaveBeenCalledWith( - session.messageContext.id, - expect.objectContaining({ - statusCode: 429, - }) - ); - }); - - test("falls back to upstream body when raw body is missing", async () => { - const session = createSession(); + test("falls back to the sanitized upstream body when no raw body exists", async () => { + const session = await createSession(); const error = new ProxyError("Upstream failed", 502, { body: "sanitized upstream body", rawBodyTruncated: false, @@ -242,19 +146,8 @@ describe("ProxyErrorHandler.handle - Langfuse error traces", () => { ); }); - test("preserves streaming request context for early error traces", async () => { - const requestMessage = { - model: "claude-sonnet-4-20250514", - messages: [{ role: "user", content: "hello" }], - stream: true, - }; - const session = createSession({ - request: { - message: requestMessage, - model: requestMessage.model, - log: JSON.stringify(requestMessage), - }, - }); + test("preserves body-declared streaming context on an early error", async () => { + const session = await createSession({ stream: true }); await ProxyErrorHandler.handle(session, new Error("fetch failed")); @@ -270,11 +163,9 @@ describe("ProxyErrorHandler.handle - Langfuse error traces", () => { ); }); - test("detects Gemini SSE URLs as streaming for early error traces", async () => { - const session = createSession({ - requestUrl: new URL( - "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:streamGenerateContent?alt=sse" - ), + test("detects a Gemini SSE URL as streaming on an early error", async () => { + const session = await createSession({ + url: "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:streamGenerateContent?alt=sse", }); await ProxyErrorHandler.handle(session, new Error("fetch failed")); @@ -290,23 +181,4 @@ describe("ProxyErrorHandler.handle - Langfuse error traces", () => { }) ); }); - - test("emits trace for rate limit early returns", async () => { - const session = createSession(); - - await ProxyErrorHandler.handle( - session, - new RateLimitError("rate_limit_error", "limit exceeded", "daily_quota", 12, 20, null) - ); - - expect(mocks.emitProxyLangfuseTrace).toHaveBeenCalledWith( - session, - expect.objectContaining({ - responseText: "", - statusCode: 402, - isStreaming: false, - errorMessage: "limit exceeded", - }) - ); - }); }); diff --git a/tests/unit/proxy/error-handler-overrides.test.ts b/tests/unit/proxy/error-handler-overrides.test.ts new file mode 100644 index 000000000..b3d9bd269 --- /dev/null +++ b/tests/unit/proxy/error-handler-overrides.test.ts @@ -0,0 +1,176 @@ +import { Context } from "hono"; +import { beforeEach, describe, expect, test, vi } from "vitest"; +import { ProxyErrorHandler } from "@/app/v1/_lib/proxy/error-handler"; +import { ProxyError } from "@/app/v1/_lib/proxy/errors"; +import { ProxySession } from "@/app/v1/_lib/proxy/session"; +import type { ErrorDetectionResult } from "@/lib/error-rule-detector"; +import type { emitProxyLangfuseTrace } from "@/lib/langfuse/emit-proxy-trace"; + +const mocks = vi.hoisted(() => ({ + detectAsync: vi.fn<(content: string) => Promise>(), + emitProxyLangfuseTrace: vi.fn(), + getCachedSystemSettings: vi.fn(async () => ({ + verboseProviderError: false, + passThroughUpstreamErrorMessage: true, + })), +})); + +vi.mock("@/lib/error-rule-detector", () => ({ + errorRuleDetector: { detectAsync: mocks.detectAsync }, +})); + +vi.mock("@/lib/config/system-settings-cache", () => ({ + getCachedSystemSettings: mocks.getCachedSystemSettings, +})); + +vi.mock("@/lib/langfuse/emit-proxy-trace", () => ({ + emitProxyLangfuseTrace: mocks.emitProxyLangfuseTrace, +})); + +vi.mock("@/lib/logger", () => ({ + logger: { + debug: vi.fn(), + error: vi.fn(), + fatal: vi.fn(), + info: vi.fn(), + trace: vi.fn(), + warn: vi.fn(), + }, +})); + +async function createSession(): Promise { + const request = new Request("https://hub.test/v1/messages", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + model: "claude-sonnet-4-20250514", + messages: [{ role: "user", content: "hello" }], + }), + }); + const session = await ProxySession.fromContext(new Context(request)); + session.setSessionId("s_override"); + return session; +} + +function createUpstreamError(): ProxyError { + return new ProxyError("Upstream failed", 502, { + body: "Quota exceeded", + rawBody: '{"error":{"message":"raw upstream failure"}}', + providerId: 7, + providerName: "provider-a", + requestId: "req_upstream", + safeClientMessageCandidate: "Quota exceeded", + }); +} + +describe("ProxyErrorHandler.handle overrides", () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.detectAsync.mockResolvedValue({ matched: false }); + mocks.getCachedSystemSettings.mockResolvedValue({ + verboseProviderError: false, + passThroughUpstreamErrorMessage: true, + }); + }); + + test("applies an explicit response and status before upstream content", async () => { + mocks.detectAsync.mockResolvedValue({ + matched: true, + overrideStatusCode: 429, + overrideResponse: { + error: { + type: "rate_limit_error", + message: "masked quota message", + code: "provider_unavailable", + }, + }, + }); + const session = await createSession(); + + const response = await ProxyErrorHandler.handle(session, createUpstreamError()); + const responseText = await response.text(); + + expect(response.status).toBe(429); + expect(responseText).toContain("masked quota message (cch_session_id: s_override)"); + expect(responseText).toContain("provider_unavailable"); + expect(responseText).not.toContain("raw upstream failure"); + expect(responseText).not.toContain("req_upstream"); + expect(mocks.emitProxyLangfuseTrace).toHaveBeenCalledWith( + session, + expect.objectContaining({ + responseText: expect.stringContaining("masked quota message"), + statusCode: 429, + errorMessage: "masked quota message", + }) + ); + const trace = mocks.emitProxyLangfuseTrace.mock.calls[0]?.[1]; + expect(trace?.responseText).not.toContain("raw upstream failure"); + }); + + test("keeps the upstream status for a response-only override", async () => { + mocks.detectAsync.mockResolvedValue({ + matched: true, + overrideResponse: { + type: "error", + error: { type: "invalid_request_error", message: "custom response" }, + }, + }); + const session = await createSession(); + + const response = await ProxyErrorHandler.handle(session, createUpstreamError()); + const responseText = await response.text(); + + expect(response.status).toBe(502); + expect(responseText).toContain("custom response (cch_session_id: s_override)"); + expect(responseText).not.toContain("req_upstream"); + }); + + test("applies a status-only override while retaining the resolved client message", async () => { + mocks.detectAsync.mockResolvedValue({ matched: true, overrideStatusCode: 418 }); + const session = await createSession(); + + const response = await ProxyErrorHandler.handle(session, createUpstreamError()); + const responseText = await response.text(); + + expect(response.status).toBe(418); + expect(responseText).toContain("raw upstream failure (cch_session_id: s_override)"); + expect(responseText).toContain('"request_id":"req_upstream"'); + expect(responseText).not.toContain("provider-a"); + }); + + test("falls back to the upstream status when an override status is invalid", async () => { + mocks.detectAsync.mockResolvedValue({ matched: true, overrideStatusCode: 200 }); + const session = await createSession(); + + const response = await ProxyErrorHandler.handle(session, createUpstreamError()); + + expect(response.status).toBe(502); + expect(mocks.emitProxyLangfuseTrace).toHaveBeenCalledWith( + session, + expect.objectContaining({ statusCode: 502 }) + ); + }); + + test("resolves a blank override message through the client-safe fallback", async () => { + mocks.detectAsync.mockResolvedValue({ + matched: true, + overrideStatusCode: 451, + overrideResponse: { + error: { type: "invalid_request_error", message: "", code: "provider_unavailable" }, + }, + }); + mocks.getCachedSystemSettings.mockResolvedValue({ + verboseProviderError: false, + passThroughUpstreamErrorMessage: false, + }); + const session = await createSession(); + + const response = await ProxyErrorHandler.handle(session, createUpstreamError()); + const responseText = await response.text(); + + expect(response.status).toBe(451); + expect(responseText).toContain("上游服务暂时不可用,请稍后重试"); + expect(responseText).not.toContain("Quota exceeded"); + expect(responseText).not.toContain("raw upstream failure"); + }); +}); diff --git a/tests/unit/proxy/error-handler-terminal-status.test.ts b/tests/unit/proxy/error-handler-terminal-status.test.ts new file mode 100644 index 000000000..98b2d5a60 --- /dev/null +++ b/tests/unit/proxy/error-handler-terminal-status.test.ts @@ -0,0 +1,219 @@ +import { Context } from "hono"; +import { beforeEach, describe, expect, test, vi } from "vitest"; +import { ProxyErrorHandler } from "@/app/v1/_lib/proxy/error-handler"; +import { ProxyError, RateLimitError } from "@/app/v1/_lib/proxy/errors"; +import { ProxySession } from "@/app/v1/_lib/proxy/session"; +import type { ErrorDetectionResult } from "@/lib/error-rule-detector"; +import type { Provider } from "@/types/provider"; + +const mocks = vi.hoisted(() => ({ + detectAsync: vi.fn<(content: string) => Promise>(), + emitProxyLangfuseTrace: vi.fn(), + getCachedSystemSettings: vi.fn(async () => ({ + verboseProviderError: false, + passThroughUpstreamErrorMessage: false, + })), +})); + +vi.mock("@/lib/error-rule-detector", () => ({ + errorRuleDetector: { detectAsync: mocks.detectAsync }, +})); + +vi.mock("@/lib/config/system-settings-cache", () => ({ + getCachedSystemSettings: mocks.getCachedSystemSettings, +})); + +vi.mock("@/lib/langfuse/emit-proxy-trace", () => ({ + emitProxyLangfuseTrace: mocks.emitProxyLangfuseTrace, +})); + +vi.mock("@/lib/logger", () => ({ + logger: { + debug: vi.fn(), + error: vi.fn(), + fatal: vi.fn(), + info: vi.fn(), + trace: vi.fn(), + warn: vi.fn(), + }, +})); + +const PROVIDER = { + id: 7, + name: "provider-a", + url: "https://provider-a.example.com", + key: "provider-key", + providerVendorId: 70, + isEnabled: true, + weight: 1, + priority: 0, + groupPriorities: null, + costMultiplier: 1, + groupTag: null, + providerType: "claude", + preserveClientIp: false, + disableSessionReuse: false, + modelRedirects: null, + activeTimeStart: null, + activeTimeEnd: null, + allowedModels: null, + allowedClients: [], + blockedClients: [], + mcpPassthroughType: "none", + mcpPassthroughUrl: null, + limit5hUsd: null, + limit5hResetMode: "fixed", + limitDailyUsd: null, + dailyResetMode: "fixed", + dailyResetTime: "00:00", + limitWeeklyUsd: null, + limitMonthlyUsd: null, + limitTotalUsd: null, + totalCostResetAt: null, + limitConcurrentSessions: 0, + maxRetryAttempts: 3, + circuitBreakerFailureThreshold: 5, + circuitBreakerOpenDuration: 1_800_000, + circuitBreakerHalfOpenSuccessThreshold: 2, + proxyUrl: null, + proxyFallbackToDirect: false, + customHeaders: null, + firstByteTimeoutStreamingMs: 30_000, + streamingIdleTimeoutMs: 10_000, + requestTimeoutNonStreamingMs: 600_000, + websiteUrl: null, + faviconUrl: null, + cacheTtlPreference: null, + swapCacheTtlBilling: false, + context1mPreference: null, + codexReasoningEffortPreference: null, + codexReasoningSummaryPreference: null, + codexTextVerbosityPreference: null, + codexParallelToolCallsPreference: null, + codexImageGenerationPreference: null, + codexServiceTierPreference: null, + anthropicMaxTokensPreference: null, + anthropicThinkingBudgetPreference: null, + anthropicAdaptiveThinking: null, + geminiGoogleSearchPreference: null, + tpm: 0, + rpm: 0, + rpd: 0, + cc: 0, + createdAt: new Date("2026-01-01T00:00:00.000Z"), + updatedAt: new Date("2026-01-01T00:00:00.000Z"), +} satisfies Provider; + +const RATE_LIMIT_CASES = [ + { limitType: "rpm", expectedStatus: 429 }, + { limitType: "concurrent_sessions", expectedStatus: 429 }, + { limitType: "daily_quota", expectedStatus: 402 }, + { limitType: "usd_5h", expectedStatus: 402 }, +] as const; + +async function createSession(): Promise { + const request = new Request("https://hub.test/v1/messages", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + model: "claude-sonnet-4-20250514", + messages: [{ role: "user", content: "hello" }], + }), + }); + return ProxySession.fromContext(new Context(request)); +} + +describe("ProxyErrorHandler.handle terminal status", () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.detectAsync.mockResolvedValue({ matched: false }); + mocks.getCachedSystemSettings.mockResolvedValue({ + verboseProviderError: false, + passThroughUpstreamErrorMessage: false, + }); + }); + + test.each([400, 404, 429, 524])("preserves ProxyError status %i", async (status) => { + const session = await createSession(); + const error = new ProxyError("Upstream failed", status, { body: "Upstream failed" }); + + const response = await ProxyErrorHandler.handle(session, error); + + expect(response.status).toBe(status); + expect(mocks.emitProxyLangfuseTrace).toHaveBeenCalledWith( + session, + expect.objectContaining({ statusCode: status }) + ); + }); + + test("uses the last failed provider-chain status for a generic error", async () => { + const session = await createSession(); + session.addProviderToChain(PROVIDER, { reason: "retry_failed", statusCode: 503 }); + + const response = await ProxyErrorHandler.handle(session, new Error("fetch failed")); + + expect(response.status).toBe(503); + expect(mocks.emitProxyLangfuseTrace).toHaveBeenCalledWith( + session, + expect.objectContaining({ statusCode: 503, errorMessage: "fetch failed" }) + ); + }); + + test.each(RATE_LIMIT_CASES)( + "maps $limitType limits to HTTP $expectedStatus", + async ({ limitType, expectedStatus }) => { + const session = await createSession(); + const error = new RateLimitError( + "rate_limit_error", + "limit exceeded", + limitType, + 12, + 20, + null + ); + + const response = await ProxyErrorHandler.handle(session, error); + + expect(response.status).toBe(expectedStatus); + expect(await response.json()).toEqual({ + error: { + type: "rate_limit_error", + message: "limit exceeded", + code: "rate_limit_exceeded", + limit_type: limitType, + current: 12, + limit: 20, + reset_time: null, + }, + }); + expect(mocks.emitProxyLangfuseTrace).toHaveBeenCalledWith( + session, + expect.objectContaining({ + responseText: "", + statusCode: expectedStatus, + errorMessage: "limit exceeded", + }) + ); + } + ); + + test("keeps fixed-window rate-limit headers", async () => { + const session = await createSession(); + const error = new RateLimitError( + "rate_limit_error", + "daily limit exceeded", + "daily_quota", + 12, + 20, + "2026-04-22T13:30:00.000Z" + ); + + const response = await ProxyErrorHandler.handle(session, error); + + expect(response.status).toBe(402); + expect(response.headers.get("X-RateLimit-Limit")).toBe("20"); + expect(response.headers.get("X-RateLimit-Remaining")).toBe("8"); + expect(response.headers.get("X-RateLimit-Reset")).toBe("1776864600"); + expect(response.headers.get("Retry-After")).toBe("0"); + }); +}); diff --git a/tests/unit/proxy/fake-streaming-response-validator.test.ts b/tests/unit/proxy/fake-streaming-response-validator.test.ts index 83b33deba..310f45b0c 100644 --- a/tests/unit/proxy/fake-streaming-response-validator.test.ts +++ b/tests/unit/proxy/fake-streaming-response-validator.test.ts @@ -15,64 +15,54 @@ function failure(family: ProtocolFamily, body: string, isStream: boolean, status describe("validateUpstreamResponse", () => { describe("status code handling", () => { - test.each([ - "anthropic", - "openai-chat", - "openai-responses", - "gemini", - ])("%s: non-2xx is failure regardless of body", (family) => { - const valid = `{"id":"ok","model":"m","content":[{"type":"text","text":"hi"}]}`; - expect(failure(family, valid, false, 500).ok).toBe(false); - expect(failure(family, valid, false, 502).ok).toBe(false); - expect(failure(family, valid, false, 429).ok).toBe(false); - expect(failure(family, valid, false, 401).ok).toBe(false); - }); + test.each(["anthropic", "openai-chat", "openai-responses", "gemini"])( + "%s: non-2xx is failure regardless of body", + (family) => { + const valid = `{"id":"ok","model":"m","content":[{"type":"text","text":"hi"}]}`; + expect(failure(family, valid, false, 500).ok).toBe(false); + expect(failure(family, valid, false, 502).ok).toBe(false); + expect(failure(family, valid, false, 429).ok).toBe(false); + expect(failure(family, valid, false, 401).ok).toBe(false); + } + ); }); describe("empty / whitespace bodies", () => { - test.each([ - "anthropic", - "openai-chat", - "openai-responses", - "gemini", - ])("%s: empty body fails (non-stream)", (family) => { - expect(failure(family, "", false).ok).toBe(false); - expect(failure(family, " ", false).ok).toBe(false); - expect(failure(family, "\n\n \t\n", false).ok).toBe(false); - }); + test.each(["anthropic", "openai-chat", "openai-responses", "gemini"])( + "%s: empty body fails (non-stream)", + (family) => { + expect(failure(family, "", false).ok).toBe(false); + expect(failure(family, " ", false).ok).toBe(false); + expect(failure(family, "\n\n \t\n", false).ok).toBe(false); + } + ); - test.each([ - "anthropic", - "openai-chat", - "openai-responses", - "gemini", - ])("%s: empty body fails (stream)", (family) => { - expect(failure(family, "", true).ok).toBe(false); - expect(failure(family, " ", true).ok).toBe(false); - }); + test.each(["anthropic", "openai-chat", "openai-responses", "gemini"])( + "%s: empty body fails (stream)", + (family) => { + expect(failure(family, "", true).ok).toBe(false); + expect(failure(family, " ", true).ok).toBe(false); + } + ); }); describe("invalid JSON for non-stream", () => { - test.each([ - "anthropic", - "openai-chat", - "openai-responses", - "gemini", - ])("%s: invalid JSON fails non-stream", (family) => { - expect(failure(family, "not-json", false).ok).toBe(false); - expect(failure(family, "{ truncated", false).ok).toBe(false); - }); + test.each(["anthropic", "openai-chat", "openai-responses", "gemini"])( + "%s: invalid JSON fails non-stream", + (family) => { + expect(failure(family, "not-json", false).ok).toBe(false); + expect(failure(family, "{ truncated", false).ok).toBe(false); + } + ); }); describe("SSE failure cases", () => { - test.each([ - "anthropic", - "openai-chat", - "openai-responses", - "gemini", - ])("%s: comment-only SSE fails", (family) => { - expect(failure(family, ": ping\n\n: ping\n\n", true).ok).toBe(false); - }); + test.each(["anthropic", "openai-chat", "openai-responses", "gemini"])( + "%s: comment-only SSE fails", + (family) => { + expect(failure(family, ": ping\n\n: ping\n\n", true).ok).toBe(false); + } + ); test("openai-chat: [DONE]-only SSE fails", () => { expect(failure("openai-chat", "data: [DONE]\n\n", true).ok).toBe(false); diff --git a/tests/unit/proxy/fake-streaming-response.test.ts b/tests/unit/proxy/fake-streaming-response.test.ts index 88cad4aa6..bc6350f76 100644 --- a/tests/unit/proxy/fake-streaming-response.test.ts +++ b/tests/unit/proxy/fake-streaming-response.test.ts @@ -40,15 +40,13 @@ function parseSseEvents(body: string): Array<{ event: string | null; data: strin } describe("emitFinalNonStream", () => { - test.each([ - "anthropic", - "openai-chat", - "openai-responses", - "gemini", - ])("%s: returns the validated final body verbatim", (family) => { - const body = JSON.stringify({ id: "x", model: "m", content: [{ type: "text", text: "hi" }] }); - expect(emitFinalNonStream({ family, finalBody: body })).toBe(body); - }); + test.each(["anthropic", "openai-chat", "openai-responses", "gemini"])( + "%s: returns the validated final body verbatim", + (family) => { + const body = JSON.stringify({ id: "x", model: "m", content: [{ type: "text", text: "hi" }] }); + expect(emitFinalNonStream({ family, finalBody: body })).toBe(body); + } + ); }); describe("emitFinalStream — anthropic", () => { diff --git a/tests/unit/proxy/fake-streaming-stream-intent.test.ts b/tests/unit/proxy/fake-streaming-stream-intent.test.ts index 75f65067e..f4b623048 100644 --- a/tests/unit/proxy/fake-streaming-stream-intent.test.ts +++ b/tests/unit/proxy/fake-streaming-stream-intent.test.ts @@ -26,33 +26,31 @@ function inputs({ describe("detectClientStreamIntent", () => { describe("standard formats (claude / openai / response)", () => { - test.each([ - "claude", - "openai", - "response", - ])("%s: body.stream === true => stream", (format) => { - expect( - detectClientStreamIntent( - inputs({ format, pathname: "/v1/messages", body: { stream: true } }) - ) - ).toBe(true); - }); + test.each(["claude", "openai", "response"])( + "%s: body.stream === true => stream", + (format) => { + expect( + detectClientStreamIntent( + inputs({ format, pathname: "/v1/messages", body: { stream: true } }) + ) + ).toBe(true); + } + ); - test.each([ - "claude", - "openai", - "response", - ])("%s: body.stream missing or false => non-stream", (format) => { - expect( - detectClientStreamIntent( - inputs({ format, pathname: "/v1/messages", body: { stream: false } }) - ) - ).toBe(false); - expect(detectClientStreamIntent(inputs({ format, pathname: "/v1/messages", body: {} }))).toBe( - false - ); - expect(detectClientStreamIntent(inputs({ format, pathname: "/v1/messages" }))).toBe(false); - }); + test.each(["claude", "openai", "response"])( + "%s: body.stream missing or false => non-stream", + (format) => { + expect( + detectClientStreamIntent( + inputs({ format, pathname: "/v1/messages", body: { stream: false } }) + ) + ).toBe(false); + expect( + detectClientStreamIntent(inputs({ format, pathname: "/v1/messages", body: {} })) + ).toBe(false); + expect(detectClientStreamIntent(inputs({ format, pathname: "/v1/messages" }))).toBe(false); + } + ); test("standard formats ignore path / query for stream intent", () => { expect( @@ -69,20 +67,20 @@ describe("detectClientStreamIntent", () => { }); describe("gemini family", () => { - test.each([ - "gemini", - "gemini-cli", - ])("%s: streamGenerateContent in path => stream", (format) => { - expect( - detectClientStreamIntent( - inputs({ - format, - pathname: "/v1beta/models/gemini-1.5-pro:streamGenerateContent", - body: {}, - }) - ) - ).toBe(true); - }); + test.each(["gemini", "gemini-cli"])( + "%s: streamGenerateContent in path => stream", + (format) => { + expect( + detectClientStreamIntent( + inputs({ + format, + pathname: "/v1beta/models/gemini-1.5-pro:streamGenerateContent", + body: {}, + }) + ) + ).toBe(true); + } + ); test.each(["gemini", "gemini-cli"])("%s: alt=sse query => stream", (format) => { expect( @@ -97,43 +95,43 @@ describe("detectClientStreamIntent", () => { ).toBe(true); }); - test.each([ - "gemini", - "gemini-cli", - ])("%s: body.stream === true => stream", (format) => { - expect( - detectClientStreamIntent( - inputs({ - format, - pathname: "/v1beta/models/gemini-1.5-pro:generateContent", - body: { stream: true }, - }) - ) - ).toBe(true); - }); + test.each(["gemini", "gemini-cli"])( + "%s: body.stream === true => stream", + (format) => { + expect( + detectClientStreamIntent( + inputs({ + format, + pathname: "/v1beta/models/gemini-1.5-pro:generateContent", + body: { stream: true }, + }) + ) + ).toBe(true); + } + ); - test.each([ - "gemini", - "gemini-cli", - ])("%s: no streaming signal => non-stream", (format) => { - expect( - detectClientStreamIntent( - inputs({ - format, - pathname: "/v1beta/models/gemini-1.5-pro:generateContent", - body: { stream: false }, - }) - ) - ).toBe(false); - expect( - detectClientStreamIntent( - inputs({ - format, - pathname: "/v1beta/models/gemini-1.5-pro:generateContent", - }) - ) - ).toBe(false); - }); + test.each(["gemini", "gemini-cli"])( + "%s: no streaming signal => non-stream", + (format) => { + expect( + detectClientStreamIntent( + inputs({ + format, + pathname: "/v1beta/models/gemini-1.5-pro:generateContent", + body: { stream: false }, + }) + ) + ).toBe(false); + expect( + detectClientStreamIntent( + inputs({ + format, + pathname: "/v1beta/models/gemini-1.5-pro:generateContent", + }) + ) + ).toBe(false); + } + ); test("gemini search supports object form", () => { expect( diff --git a/tests/unit/proxy/pricing-no-price.test.ts b/tests/unit/proxy/pricing-no-price.test.ts index ef15fc14e..78e2cd4f7 100644 --- a/tests/unit/proxy/pricing-no-price.test.ts +++ b/tests/unit/proxy/pricing-no-price.test.ts @@ -37,6 +37,7 @@ vi.mock("@/repository/system-config", () => ({ vi.mock("@/repository/message", () => ({ updateMessageRequestCost: vi.fn(), updateMessageRequestDetails: vi.fn(), + updateMessageRequestDetailsDurably: vi.fn(), updateMessageRequestDuration: vi.fn(), })); @@ -67,7 +68,11 @@ vi.mock("@/lib/proxy-status-tracker", () => ({ import { finalizeRequestStats } from "@/app/v1/_lib/proxy/response-handler"; import { ProxySession } from "@/app/v1/_lib/proxy/session"; import { RateLimitService } from "@/lib/rate-limit"; -import { updateMessageRequestCost, updateMessageRequestDetails } from "@/repository/message"; +import { + updateMessageRequestCost, + updateMessageRequestDetails, + updateMessageRequestDetailsDurably, +} from "@/repository/message"; import { findLatestPriceByModel } from "@/repository/model-price"; import { getSystemSettings } from "@/repository/system-config"; @@ -212,7 +217,7 @@ describe("价格表缺失/查询失败:请求不计费且不报错", () => { "bad upstream" ); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( 2000, expect.objectContaining({ statusCode: 502, diff --git a/tests/unit/proxy/provider-selector-cross-type-model.test.ts b/tests/unit/proxy/provider-selector-cross-type-model.test.ts index 60ff7b9bf..f171f6ffe 100644 --- a/tests/unit/proxy/provider-selector-cross-type-model.test.ts +++ b/tests/unit/proxy/provider-selector-cross-type-model.test.ts @@ -198,21 +198,18 @@ describe("providerSupportsModel - direct unit tests (#832)", () => { }, ]; - test.each(cases)("$name", async ({ - providerType, - allowedModels, - modelRedirects, - requestedModel, - expected, - }) => { - const { providerSupportsModel } = await import("@/app/v1/_lib/proxy/provider-selector"); - const provider = createProvider({ - providerType, - allowedModels, - ...(modelRedirects && { modelRedirects }), - }); - expect(providerSupportsModel(provider, requestedModel)).toBe(expected); - }); + test.each(cases)( + "$name", + async ({ providerType, allowedModels, modelRedirects, requestedModel, expected }) => { + const { providerSupportsModel } = await import("@/app/v1/_lib/proxy/provider-selector"); + const provider = createProvider({ + providerType, + allowedModels, + ...(modelRedirects && { modelRedirects }), + }); + expect(providerSupportsModel(provider, requestedModel)).toBe(expected); + } + ); }); // ══════════════════════════════════════════════════════════════════ @@ -305,7 +302,8 @@ describe("findReusable - cross-type model routing (#832)", () => { expect(result).toBeNull(); expect(sessionManagerMocks.SessionManager.clearSessionProvider).toHaveBeenCalledWith( - "cross-type-3" + "cross-type-3", + 12 ); }); @@ -337,7 +335,8 @@ describe("findReusable - cross-type model routing (#832)", () => { expect(result).toBeNull(); expect(sessionManagerMocks.SessionManager.clearSessionProvider).toHaveBeenCalledWith( - "cross-type-6" + "cross-type-6", + 15 ); }); }); diff --git a/tests/unit/proxy/provider-selector-model-mismatch-binding.test.ts b/tests/unit/proxy/provider-selector-model-mismatch-binding.test.ts index 2de0429b2..d21486433 100644 --- a/tests/unit/proxy/provider-selector-model-mismatch-binding.test.ts +++ b/tests/unit/proxy/provider-selector-model-mismatch-binding.test.ts @@ -115,7 +115,8 @@ describe("findReusable - model mismatch clears stale binding", () => { expect(result).toBeNull(); expect(sessionManagerMocks.SessionManager.clearSessionProvider).toHaveBeenCalledWith( - "sess_disable_reuse" + "sess_disable_reuse", + 78 ); }); @@ -139,7 +140,8 @@ describe("findReusable - model mismatch clears stale binding", () => { expect(result).toBeNull(); // Key assertion: clearSessionProvider should have been called expect(sessionManagerMocks.SessionManager.clearSessionProvider).toHaveBeenCalledWith( - "4c25cf92" + "4c25cf92", + 78 ); }); @@ -162,7 +164,8 @@ describe("findReusable - model mismatch clears stale binding", () => { expect(result).toBeNull(); expect(sessionManagerMocks.SessionManager.clearSessionProvider).toHaveBeenCalledWith( - "sess_response_format_mismatch" + "sess_response_format_mismatch", + 94 ); }); @@ -236,7 +239,8 @@ describe("findReusable - model mismatch clears stale binding", () => { expect(result).toBeNull(); expect(sessionManagerMocks.SessionManager.clearSessionProvider).toHaveBeenCalledWith( - "sess_variant" + "sess_variant", + 78 ); }); }); diff --git a/tests/unit/proxy/proxy-forwarder-endpoint-audit.test.ts b/tests/unit/proxy/proxy-forwarder-endpoint-audit.test.ts index 51610de69..a02dc05a2 100644 --- a/tests/unit/proxy/proxy-forwarder-endpoint-audit.test.ts +++ b/tests/unit/proxy/proxy-forwarder-endpoint-audit.test.ts @@ -418,68 +418,69 @@ describe("ProxyForwarder - endpoint audit", () => { test.each([ { requestPath: "/v1/messages/count_tokens", providerType: "claude" as const }, { requestPath: "/v1/responses/compact", providerType: "codex" as const }, - ])("raw 端点 $requestPath: endpoint 选择失败时不应静默回退到 provider.url", async ({ - requestPath, - providerType, - }) => { - const session = createSession(new URL(`https://example.com${requestPath}`)); - const provider = createProvider({ - providerType, - providerVendorId: 123, - url: `https://provider.example.com${requestPath}?key=SECRET`, - }); - session.setProvider(provider); + ])( + "raw 端点 $requestPath: endpoint 选择失败时不应静默回退到 provider.url", + async ({ requestPath, providerType }) => { + const session = createSession(new URL(`https://example.com${requestPath}`)); + const provider = createProvider({ + providerType, + providerVendorId: 123, + url: `https://provider.example.com${requestPath}?key=SECRET`, + }); + session.setProvider(provider); - mocks.getPreferredProviderEndpoints.mockRejectedValueOnce(new Error("boom")); + mocks.getPreferredProviderEndpoints.mockRejectedValueOnce(new Error("boom")); - const doForward = vi.spyOn( - ProxyForwarder as unknown as { doForward: (...args: unknown[]) => unknown }, - "doForward" - ); - doForward.mockResolvedValueOnce( - new Response("{}", { - status: 200, - headers: { - "content-type": "application/json", - "content-length": "2", - }, - }) - ); + const doForward = vi.spyOn( + ProxyForwarder as unknown as { doForward: (...args: unknown[]) => unknown }, + "doForward" + ); + doForward.mockResolvedValueOnce( + new Response("{}", { + status: 200, + headers: { + "content-type": "application/json", + "content-length": "2", + }, + }) + ); - const rejected = await ProxyForwarder.send(session) - .then(() => false) - .catch(() => true); + const rejected = await ProxyForwarder.send(session) + .then(() => false) + .catch(() => true); - expect(rejected, `raw 端点 ${requestPath} endpoint 选择失败后不允许静默回退 provider.url`).toBe( - true - ); - expect(doForward).not.toHaveBeenCalled(); + expect( + rejected, + `raw 端点 ${requestPath} endpoint 选择失败后不允许静默回退 provider.url` + ).toBe(true); + expect(doForward).not.toHaveBeenCalled(); - expect(logger.warn).toHaveBeenCalledWith( - "[ProxyForwarder] Failed to load provider endpoints", - expect.objectContaining({ - providerId: provider.id, - vendorId: 123, - providerType, - strictEndpointPolicy: true, - reason: "selector_error", - error: "boom", - }) - ); + expect(logger.warn).toHaveBeenCalledWith( + "[ProxyForwarder] Failed to load provider endpoints", + expect.objectContaining({ + providerId: provider.id, + vendorId: 123, + providerType, + strictEndpointPolicy: true, + reason: "selector_error", + error: "boom", + }) + ); - expect(logger.warn).toHaveBeenCalledWith( - "ProxyForwarder: Strict endpoint policy blocked legacy provider.url fallback", - expect.objectContaining({ - providerId: provider.id, - vendorId: 123, - providerType, - requestPath, - reason: "strict_blocked_legacy_fallback", - strictBlockCause: "selector_error", - selectorError: "boom", - }) - ); - }); + expect(logger.warn).toHaveBeenCalledWith( + "ProxyForwarder: Strict endpoint policy blocked legacy provider.url fallback", + expect.objectContaining({ + providerId: provider.id, + vendorId: 123, + providerType, + requestPath, + reason: "strict_blocked_legacy_fallback", + strictBlockCause: "selector_error", + selectorError: "boom", + }) + ); + } + ); test("raw 端点空候选应记录 no_endpoint_candidates 且不混淆为 selector_error", async () => { const requestPath = "/v1/messages/count_tokens"; diff --git a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts index fe8d2464d..014219a24 100644 --- a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts +++ b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts @@ -1454,7 +1454,7 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { await rejection; expect(controller1.signal.aborted).toBe(true); expect(controller2.signal.aborted).toBe(true); - expect(mocks.clearSessionProvider).toHaveBeenCalledWith("sess-hedge"); + expect(mocks.clearSessionProvider).toHaveBeenCalledWith("sess-hedge", 1); expect(mocks.recordFailure).not.toHaveBeenCalled(); expect(mocks.recordSuccess).not.toHaveBeenCalled(); @@ -1584,74 +1584,74 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { category: ProxyErrorCategory.SYSTEM_ERROR, errorFactory: () => new Error("fetch failed"), }, - ])("when a real hedge race ends with only $name, terminal error should be generic fallback", async ({ - category, - errorFactory, - }) => { - vi.useFakeTimers(); - - try { - const provider1 = createProvider({ id: 1, name: "p1", firstByteTimeoutStreamingMs: 100 }); - const provider2 = createProvider({ id: 2, name: "p2", firstByteTimeoutStreamingMs: 100 }); - const session = createSession(); - session.setProvider(provider1); - - mocks.pickRandomProviderWithExclusion - .mockResolvedValueOnce(provider2) - .mockResolvedValueOnce(null); - mocks.categorizeErrorAsync.mockResolvedValueOnce(category).mockResolvedValueOnce(category); - - const doForward = vi.spyOn( - ProxyForwarder as unknown as { - doForward: (...args: unknown[]) => Promise; - }, - "doForward" - ); - - const controller1 = new AbortController(); - const controller2 = new AbortController(); - - doForward.mockImplementationOnce(async (attemptSession) => { - const runtime = attemptSession as ProxySession & AttemptRuntime; - runtime.responseController = controller1; - runtime.clearResponseTimeout = vi.fn(); - return createDelayedFailure({ - delayMs: 150, - error: errorFactory(provider1), - controller: controller1, + ])( + "when a real hedge race ends with only $name, terminal error should be generic fallback", + async ({ category, errorFactory }) => { + vi.useFakeTimers(); + + try { + const provider1 = createProvider({ id: 1, name: "p1", firstByteTimeoutStreamingMs: 100 }); + const provider2 = createProvider({ id: 2, name: "p2", firstByteTimeoutStreamingMs: 100 }); + const session = createSession(); + session.setProvider(provider1); + + mocks.pickRandomProviderWithExclusion + .mockResolvedValueOnce(provider2) + .mockResolvedValueOnce(null); + mocks.categorizeErrorAsync.mockResolvedValueOnce(category).mockResolvedValueOnce(category); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + + const controller1 = new AbortController(); + const controller2 = new AbortController(); + + doForward.mockImplementationOnce(async (attemptSession) => { + const runtime = attemptSession as ProxySession & AttemptRuntime; + runtime.responseController = controller1; + runtime.clearResponseTimeout = vi.fn(); + return createDelayedFailure({ + delayMs: 150, + error: errorFactory(provider1), + controller: controller1, + }); }); - }); - doForward.mockImplementationOnce(async (attemptSession) => { - const runtime = attemptSession as ProxySession & AttemptRuntime; - runtime.responseController = controller2; - runtime.clearResponseTimeout = vi.fn(); - return createDelayedFailure({ - delayMs: 160, - error: errorFactory(provider2), - controller: controller2, + doForward.mockImplementationOnce(async (attemptSession) => { + const runtime = attemptSession as ProxySession & AttemptRuntime; + runtime.responseController = controller2; + runtime.clearResponseTimeout = vi.fn(); + return createDelayedFailure({ + delayMs: 160, + error: errorFactory(provider2), + controller: controller2, + }); }); - }); - const responsePromise = ProxyForwarder.send(session); - const errorPromise = responsePromise.catch((rejection) => rejection as UpstreamProxyError); + const responsePromise = ProxyForwarder.send(session); + const errorPromise = responsePromise.catch((rejection) => rejection as UpstreamProxyError); - await vi.advanceTimersByTimeAsync(100); - expect(doForward).toHaveBeenCalledTimes(2); + await vi.advanceTimersByTimeAsync(100); + expect(doForward).toHaveBeenCalledTimes(2); - await vi.runAllTimersAsync(); - const error = await errorPromise; + await vi.runAllTimersAsync(); + const error = await errorPromise; - expect(error).toBeInstanceOf(UpstreamProxyError); - expect(error.statusCode).toBe(503); - expect(error.message).toBe("所有供应商暂时不可用,请稍后重试"); - expect(error.message).not.toContain("invalid key"); - expect(error.message).not.toContain("model not found"); - expect(mocks.clearSessionProvider).toHaveBeenCalledWith("sess-hedge"); - } finally { - vi.useRealTimers(); + expect(error).toBeInstanceOf(UpstreamProxyError); + expect(error.statusCode).toBe(503); + expect(error.message).toBe("所有供应商暂时不可用,请稍后重试"); + expect(error.message).not.toContain("invalid key"); + expect(error.message).not.toContain("model not found"); + expect(mocks.clearSessionProvider).toHaveBeenCalledWith("sess-hedge", 1); + } finally { + vi.useRealTimers(); + } } - }); + ); test("non-retryable client errors should stop hedge immediately and preserve original error", async () => { const provider1 = createProvider({ id: 1, name: "p1", firstByteTimeoutStreamingMs: 100 }); @@ -1694,7 +1694,7 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { expect(error.message).toBe("prompt too long"); expect(doForward).toHaveBeenCalledTimes(1); expect(mocks.pickRandomProviderWithExclusion).not.toHaveBeenCalled(); - expect(mocks.clearSessionProvider).toHaveBeenCalledWith("sess-hedge"); + expect(mocks.clearSessionProvider).toHaveBeenCalledWith("sess-hedge", 1); expect(session.getProviderChain()).toEqual( expect.arrayContaining([ expect.objectContaining({ @@ -1722,6 +1722,49 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { ); }); + test("local DB admission overload should stop hedge without circuit mutation or failover", async () => { + const provider = createProvider({ id: 1, name: "p1", firstByteTimeoutStreamingMs: 100 }); + const session = createSession(); + session.setProvider(provider); + + const admissionCause = Object.assign(new Error("Database pool data is full"), { + name: "DbPoolAdmissionError", + code: "DB_POOL_ADMISSION_EXCEEDED", + pool: "data", + maxOutstanding: 32, + }); + const wrappedError = new Error("Failed query", { cause: admissionCause }); + mocks.categorizeErrorAsync.mockResolvedValueOnce(ProxyErrorCategory.LOCAL_OVERLOAD); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockRejectedValueOnce(wrappedError); + + const error = await ProxyForwarder.send(session).catch((rejection) => rejection as Error); + + expect(error).toBe(wrappedError); + expect(doForward).toHaveBeenCalledTimes(1); + expect(mocks.pickRandomProviderWithExclusion).not.toHaveBeenCalled(); + expect(mocks.recordEndpointFailure).not.toHaveBeenCalled(); + expect(mocks.recordFailure).not.toHaveBeenCalled(); + expect(mocks.clearSessionProvider).toHaveBeenCalledWith("sess-hedge", 1); + expect(session.getProviderChain()).toEqual([ + expect.objectContaining({ + id: provider.id, + reason: "system_error", + errorDetails: expect.objectContaining({ + system: expect.objectContaining({ + errorCode: "DB_POOL_ADMISSION_EXCEEDED", + }), + }), + }), + ]); + }); + test("hedge 备选供应商命中 thinking signature 错误时,应整流后在同供应商重试并保留审计", async () => { vi.useFakeTimers(); diff --git a/tests/unit/proxy/proxy-forwarder-retry-limit.test.ts b/tests/unit/proxy/proxy-forwarder-retry-limit.test.ts index ef193a722..454b3a7d6 100644 --- a/tests/unit/proxy/proxy-forwarder-retry-limit.test.ts +++ b/tests/unit/proxy/proxy-forwarder-retry-limit.test.ts @@ -243,64 +243,64 @@ describe("ProxyForwarder - raw passthrough fallback parity", () => { vi.mocked(categorizeErrorAsync).mockResolvedValue(ErrorCategory.PROVIDER_ERROR); }); - test.each([ - V1_ENDPOINT_PATHS.MESSAGES_COUNT_TOKENS, - V1_ENDPOINT_PATHS.RESPONSES_COMPACT, - ])("%s 失败时应允许跨 provider fallback,但仍保持 no-circuit", async (pathname) => { - vi.useFakeTimers(); - - try { - const session = createSession(new URL(`https://example.com${pathname}`)); - const provider = createProvider({ - providerType: "claude", - providerVendorId: 123, - maxRetryAttempts: 3, - }); - session.setProvider(provider); - - mocks.getPreferredProviderEndpoints.mockResolvedValue([ - makeEndpoint({ - id: 1, - vendorId: 123, + test.each([V1_ENDPOINT_PATHS.MESSAGES_COUNT_TOKENS, V1_ENDPOINT_PATHS.RESPONSES_COMPACT])( + "%s 失败时应允许跨 provider fallback,但仍保持 no-circuit", + async (pathname) => { + vi.useFakeTimers(); + + try { + const session = createSession(new URL(`https://example.com${pathname}`)); + const provider = createProvider({ providerType: "claude", - url: "https://ep1.example.com", - }), - makeEndpoint({ - id: 2, - vendorId: 123, - providerType: "claude", - url: "https://ep2.example.com", - }), - ]); - - const doForward = vi.spyOn( - ProxyForwarder as unknown as { doForward: (...args: unknown[]) => unknown }, - "doForward" - ); - const selectAlternative = vi.spyOn( - ProxyForwarder as unknown as { selectAlternative: (...args: unknown[]) => unknown }, - "selectAlternative" - ); - - doForward.mockImplementation(async () => { - throw new ProxyError("upstream failed", 500); - }); - - const sendPromise = ProxyForwarder.send(session); - let caughtError: Error | null = null; - sendPromise.catch((error) => { - caughtError = error as Error; - }); - await vi.runAllTimersAsync(); + providerVendorId: 123, + maxRetryAttempts: 3, + }); + session.setProvider(provider); + + mocks.getPreferredProviderEndpoints.mockResolvedValue([ + makeEndpoint({ + id: 1, + vendorId: 123, + providerType: "claude", + url: "https://ep1.example.com", + }), + makeEndpoint({ + id: 2, + vendorId: 123, + providerType: "claude", + url: "https://ep2.example.com", + }), + ]); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { doForward: (...args: unknown[]) => unknown }, + "doForward" + ); + const selectAlternative = vi.spyOn( + ProxyForwarder as unknown as { selectAlternative: (...args: unknown[]) => unknown }, + "selectAlternative" + ); + + doForward.mockImplementation(async () => { + throw new ProxyError("upstream failed", 500); + }); - expect(caughtError).toBeInstanceOf(ProxyError); - expect(doForward).toHaveBeenCalledTimes(1); - expect(selectAlternative).toHaveBeenCalledTimes(1); - expect(mocks.recordFailure).not.toHaveBeenCalled(); - } finally { - vi.useRealTimers(); + const sendPromise = ProxyForwarder.send(session); + let caughtError: Error | null = null; + sendPromise.catch((error) => { + caughtError = error as Error; + }); + await vi.runAllTimersAsync(); + + expect(caughtError).toBeInstanceOf(ProxyError); + expect(doForward).toHaveBeenCalledTimes(1); + expect(selectAlternative).toHaveBeenCalledTimes(1); + expect(mocks.recordFailure).not.toHaveBeenCalled(); + } finally { + vi.useRealTimers(); + } } - }); + ); }); describe("ProxyForwarder - retry limit enforcement", () => { @@ -308,6 +308,52 @@ describe("ProxyForwarder - retry limit enforcement", () => { vi.clearAllMocks(); }); + test("local DB admission overload should not retry or penalize Provider and endpoint circuits", async () => { + const session = createSession(); + const provider = createProvider({ + providerType: "claude", + providerVendorId: 123, + maxRetryAttempts: 3, + }); + session.setProvider(provider); + mocks.getPreferredProviderEndpoints.mockResolvedValue([ + makeEndpoint({ + id: 1, + vendorId: 123, + providerType: "claude", + url: "https://ep1.example.com", + }), + ]); + + vi.mocked(categorizeErrorAsync).mockResolvedValue(5 as ErrorCategory); + const admissionCause = Object.assign(new Error("Database pool data is full"), { + name: "DbPoolAdmissionError", + code: "DB_POOL_ADMISSION_EXCEEDED", + pool: "data", + maxOutstanding: 32, + }); + const wrappedError = new Error("Failed query", { cause: admissionCause }); + const doForward = vi.spyOn( + ProxyForwarder as unknown as { doForward: (...args: unknown[]) => unknown }, + "doForward" + ); + doForward.mockRejectedValue(wrappedError); + + await expect(ProxyForwarder.send(session)).rejects.toBeDefined(); + + expect(doForward).toHaveBeenCalledTimes(1); + expect(mocks.recordEndpointFailure).not.toHaveBeenCalled(); + expect(mocks.recordFailure).not.toHaveBeenCalled(); + expect(session.getProviderChain()).toEqual([ + expect.objectContaining({ + id: provider.id, + endpointId: 1, + reason: "system_error", + attemptNumber: 1, + }), + ]); + }); + test("endpoints > maxRetry: should only use top N lowest-latency endpoints", async () => { vi.useFakeTimers(); diff --git a/tests/unit/proxy/proxy-handler-concurrency-ownership.test.ts b/tests/unit/proxy/proxy-handler-concurrency-ownership.test.ts new file mode 100644 index 000000000..9e3b5650e --- /dev/null +++ b/tests/unit/proxy/proxy-handler-concurrency-ownership.test.ts @@ -0,0 +1,103 @@ +import { Context } from "hono"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { ProxySession } from "@/app/v1/_lib/proxy/session"; + +type ProxySettingsFixture = { + readonly enableHighConcurrencyMode: boolean; + readonly allowNonConversationEndpointProviderFallback: boolean; +}; + +const boundary = vi.hoisted(() => ({ + decrementConcurrentCount: vi.fn<(sessionId: string) => Promise>(), + incrementConcurrentCount: vi.fn<(sessionId: string) => Promise>(), + loadSettings: vi.fn<() => Promise>(), + runGuards: vi.fn<(session: ProxySession) => Promise>(), + send: vi.fn<(session: ProxySession) => Promise>(), +})); + +vi.mock("@/lib/config", async (importOriginal) => ({ + ...(await importOriginal()), + getCachedSystemSettings: boundary.loadSettings, +})); + +vi.mock("@/app/v1/_lib/proxy/guard-pipeline", () => ({ + GuardPipelineBuilder: { + fromSession: () => ({ run: boundary.runGuards }), + }, +})); + +vi.mock("@/app/v1/_lib/proxy/forwarder", () => ({ + ProxyForwarder: { send: boundary.send }, +})); + +vi.mock("@/lib/session-tracker", () => ({ + SessionTracker: { + decrementConcurrentCount: boundary.decrementConcurrentCount, + incrementConcurrentCount: boundary.incrementConcurrentCount, + }, +})); + +vi.mock("@/lib/proxy-status-tracker", () => ({ + ProxyStatusTracker: { + getInstance: () => ({ endRequest: vi.fn(), startRequest: vi.fn() }), + }, +})); + +import { handleProxyRequest } from "@/app/v1/_lib/proxy-handler"; + +function createContext(): Context { + const request = new Request("http://localhost/v1/messages", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ model: "claude-test", messages: [] }), + }); + return new Context(request); +} + +describe("handleProxyRequest concurrency ownership", () => { + beforeEach(() => { + boundary.runGuards.mockReset(); + boundary.send.mockReset(); + boundary.incrementConcurrentCount.mockReset(); + boundary.decrementConcurrentCount.mockReset(); + boundary.loadSettings.mockResolvedValue({ + enableHighConcurrencyMode: false, + allowNonConversationEndpointProviderFallback: true, + }); + boundary.incrementConcurrentCount.mockResolvedValue(undefined); + boundary.decrementConcurrentCount.mockResolvedValue(undefined); + boundary.send.mockResolvedValue(new Response("unused", { status: 200 })); + }); + + it("does not release concurrency for an early guard response before acquisition", async () => { + boundary.runGuards.mockImplementation(async (session) => { + session.setSessionId("session-early"); + return new Response("guard rejected", { status: 429 }); + }); + + const response = await handleProxyRequest(createContext()); + + expect(response.status).toBe(429); + expect(await response.text()).toBe("guard rejected"); + expect(boundary.incrementConcurrentCount).not.toHaveBeenCalled(); + expect(boundary.decrementConcurrentCount).not.toHaveBeenCalled(); + expect(boundary.send).not.toHaveBeenCalled(); + }); + + it("releases exactly one concurrency count after acquiring it", async () => { + boundary.runGuards.mockImplementation(async (session) => { + session.setSessionId("session-forwarded"); + return null; + }); + boundary.send.mockResolvedValue(new Response("forwarded", { status: 201 })); + + const response = await handleProxyRequest(createContext()); + + expect(response.status).toBe(201); + expect(await response.text()).toBe("forwarded"); + expect(boundary.incrementConcurrentCount).toHaveBeenCalledOnce(); + expect(boundary.incrementConcurrentCount).toHaveBeenCalledWith("session-forwarded"); + expect(boundary.decrementConcurrentCount).toHaveBeenCalledOnce(); + expect(boundary.decrementConcurrentCount).toHaveBeenCalledWith("session-forwarded"); + }); +}); diff --git a/tests/unit/proxy/proxy-handler-public-errors.test.ts b/tests/unit/proxy/proxy-handler-public-errors.test.ts new file mode 100644 index 000000000..a0a504328 --- /dev/null +++ b/tests/unit/proxy/proxy-handler-public-errors.test.ts @@ -0,0 +1,177 @@ +import { Context } from "hono"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { ProxySession } from "@/app/v1/_lib/proxy/session"; +import type { FakeStreamingWhitelistEntry } from "@/types/system-config"; + +type ProxySettingsFixture = { + readonly enableHighConcurrencyMode: boolean; + readonly allowNonConversationEndpointProviderFallback: boolean; + readonly fakeStreamingWhitelist: FakeStreamingWhitelistEntry[]; + readonly passThroughUpstreamErrorMessage: boolean; + readonly verboseProviderError: boolean; +}; + +const boundary = vi.hoisted(() => ({ + decrementConcurrentCount: vi.fn<(sessionId: string) => Promise>(), + emitProxyLangfuseTrace: vi.fn(), + getErrorOverride: vi.fn<(error: Error) => Promise>(), + incrementConcurrentCount: vi.fn<(sessionId: string) => Promise>(), + loadSettings: vi.fn<() => Promise>(), + runGuards: vi.fn<(session: ProxySession) => Promise>(), + send: vi.fn<(session: ProxySession) => Promise>(), + updateMessageRequestDetailsDurably: vi.fn(), +})); + +vi.mock("@/lib/config", async (importOriginal) => ({ + ...(await importOriginal()), + getCachedSystemSettings: boundary.loadSettings, +})); + +vi.mock("@/lib/config/system-settings-cache", async (importOriginal) => ({ + ...(await importOriginal()), + getCachedSystemSettings: boundary.loadSettings, +})); + +vi.mock("@/app/v1/_lib/proxy/guard-pipeline", () => ({ + GuardPipelineBuilder: { + fromSession: () => ({ run: boundary.runGuards }), + }, +})); + +vi.mock("@/app/v1/_lib/proxy/forwarder", () => ({ + ProxyForwarder: { send: boundary.send }, +})); + +vi.mock("@/app/v1/_lib/proxy/errors", async (importOriginal) => ({ + ...(await importOriginal()), + getErrorOverrideAsync: boundary.getErrorOverride, +})); + +vi.mock("@/lib/langfuse/emit-proxy-trace", () => ({ + emitProxyLangfuseTrace: boundary.emitProxyLangfuseTrace, +})); + +vi.mock("@/repository/message", () => ({ + updateMessageRequestDetailsDurably: boundary.updateMessageRequestDetailsDurably, +})); + +vi.mock("@/lib/session-tracker", () => ({ + SessionTracker: { + decrementConcurrentCount: boundary.decrementConcurrentCount, + incrementConcurrentCount: boundary.incrementConcurrentCount, + refreshSession: vi.fn(), + }, +})); + +vi.mock("@/lib/proxy-status-tracker", () => ({ + ProxyStatusTracker: { + getInstance: () => ({ endRequest: vi.fn(), startRequest: vi.fn() }), + }, +})); + +vi.mock("@/lib/logger", () => ({ + logger: { + debug: vi.fn(), + error: vi.fn(), + info: vi.fn(), + trace: vi.fn(), + warn: vi.fn(), + }, +})); + +import { ProxyError } from "@/app/v1/_lib/proxy/errors"; +import { handleProxyRequest } from "@/app/v1/_lib/proxy-handler"; + +const settings: ProxySettingsFixture = { + enableHighConcurrencyMode: false, + allowNonConversationEndpointProviderFallback: true, + fakeStreamingWhitelist: [], + passThroughUpstreamErrorMessage: false, + verboseProviderError: false, +}; + +describe("handleProxyRequest public error behavior", () => { + beforeEach(() => { + boundary.runGuards.mockReset(); + boundary.send.mockReset(); + boundary.incrementConcurrentCount.mockReset(); + boundary.decrementConcurrentCount.mockReset(); + boundary.loadSettings.mockReset(); + boundary.getErrorOverride.mockReset(); + boundary.loadSettings.mockResolvedValue(settings); + boundary.getErrorOverride.mockResolvedValue(null); + boundary.incrementConcurrentCount.mockResolvedValue(undefined); + boundary.decrementConcurrentCount.mockResolvedValue(undefined); + }); + + it("translates a post-session forwarding error through the real error handler", async () => { + boundary.runGuards.mockImplementation(async (session) => { + session.setSessionId("session-forward-error"); + return null; + }); + boundary.send.mockRejectedValue(new ProxyError("upstream unavailable", 503)); + const request = new Request("http://localhost/v1/messages", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ model: "claude-test", messages: [] }), + }); + + const response = await handleProxyRequest(new Context(request)); + + expect(response.status).toBe(503); + expect(await response.json()).toEqual({ + error: { + message: "上游服务暂时不可用,请稍后重试 (cch_session_id: session-forward-error)", + type: "service_unavailable_error", + code: "service_unavailable_error", + }, + }); + expect(boundary.incrementConcurrentCount).toHaveBeenCalledWith("session-forward-error"); + expect(boundary.decrementConcurrentCount).toHaveBeenCalledWith("session-forward-error"); + }); + + it("returns a public ProxyError response when request decoding fails before session creation", async () => { + const request = new Request("http://localhost/v1/messages", { + method: "POST", + headers: { + "content-encoding": "gzip", + "content-type": "application/json", + }, + body: "not-a-gzip-stream", + }); + + const response = await handleProxyRequest(new Context(request)); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.type).toBe("invalid_request_error"); + expect(body.error.message).toContain("Failed to decode 'gzip' request body"); + expect(boundary.runGuards).not.toHaveBeenCalled(); + expect(boundary.decrementConcurrentCount).not.toHaveBeenCalled(); + }); + + it("hides an unknown failure that occurs before session creation", async () => { + const request = new (class extends Request { + override clone(): Request { + throw new Error("request clone failed"); + } + })("http://localhost/v1/messages", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ model: "claude-test", messages: [] }), + }); + + const response = await handleProxyRequest(new Context(request)); + + expect(response.status).toBe(500); + expect(await response.json()).toEqual({ + error: { + message: "代理请求发生未知错误", + type: "internal_server_error", + code: "internal_server_error", + }, + }); + expect(boundary.runGuards).not.toHaveBeenCalled(); + expect(boundary.decrementConcurrentCount).not.toHaveBeenCalled(); + }); +}); diff --git a/tests/unit/proxy/proxy-handler-public-success.test.ts b/tests/unit/proxy/proxy-handler-public-success.test.ts new file mode 100644 index 000000000..99a9a0675 --- /dev/null +++ b/tests/unit/proxy/proxy-handler-public-success.test.ts @@ -0,0 +1,177 @@ +import { Context } from "hono"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { ProxySession } from "@/app/v1/_lib/proxy/session"; +import type { FakeStreamingWhitelistEntry } from "@/types/system-config"; + +type ProxySettingsFixture = { + readonly enableHighConcurrencyMode: boolean; + readonly allowNonConversationEndpointProviderFallback: boolean; + readonly enableResponseFixer: boolean; + readonly enableResponseInputRectifier: boolean; + readonly fakeStreamingWhitelist: FakeStreamingWhitelistEntry[]; +}; + +const boundary = vi.hoisted(() => ({ + decrementConcurrentCount: vi.fn<(sessionId: string) => Promise>(), + incrementConcurrentCount: vi.fn<(sessionId: string) => Promise>(), + loadSettings: vi.fn<() => Promise>(), + runGuards: vi.fn<(session: ProxySession) => Promise>(), + send: vi.fn<(session: ProxySession) => Promise>(), +})); + +vi.mock("@/lib/config", async (importOriginal) => ({ + ...(await importOriginal()), + getCachedSystemSettings: boundary.loadSettings, +})); + +vi.mock("@/lib/config/system-settings-cache", async (importOriginal) => ({ + ...(await importOriginal()), + getCachedSystemSettings: boundary.loadSettings, +})); + +vi.mock("@/app/v1/_lib/proxy/guard-pipeline", () => ({ + GuardPipelineBuilder: { + fromSession: () => ({ run: boundary.runGuards }), + }, +})); + +vi.mock("@/app/v1/_lib/proxy/forwarder", () => ({ + ProxyForwarder: { send: boundary.send }, +})); + +vi.mock("@/lib/session-tracker", () => ({ + SessionTracker: { + decrementConcurrentCount: boundary.decrementConcurrentCount, + incrementConcurrentCount: boundary.incrementConcurrentCount, + refreshSession: vi.fn(), + }, +})); + +vi.mock("@/lib/proxy-status-tracker", () => ({ + ProxyStatusTracker: { + getInstance: () => ({ endRequest: vi.fn(), startRequest: vi.fn() }), + }, +})); + +vi.mock("@/lib/logger", () => ({ + logger: { + debug: vi.fn(), + error: vi.fn(), + info: vi.fn(), + trace: vi.fn(), + warn: vi.fn(), + }, +})); + +import { handleProxyRequest } from "@/app/v1/_lib/proxy-handler"; + +const defaultSettings: ProxySettingsFixture = { + enableHighConcurrencyMode: false, + allowNonConversationEndpointProviderFallback: true, + enableResponseFixer: true, + enableResponseInputRectifier: true, + fakeStreamingWhitelist: [], +}; + +function createContext(pathname: string, body: Record): Context { + const request = new Request(`http://localhost${pathname}`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(body), + }); + return new Context(request); +} + +describe("handleProxyRequest public success behavior", () => { + beforeEach(() => { + boundary.runGuards.mockReset(); + boundary.send.mockReset(); + boundary.incrementConcurrentCount.mockReset(); + boundary.decrementConcurrentCount.mockReset(); + boundary.loadSettings.mockReset(); + boundary.loadSettings.mockResolvedValue(defaultSettings); + boundary.runGuards.mockResolvedValue(null); + boundary.incrementConcurrentCount.mockResolvedValue(undefined); + boundary.decrementConcurrentCount.mockResolvedValue(undefined); + }); + + it("returns a successful upstream response through the real dispatcher", async () => { + boundary.send.mockResolvedValue( + new Response(JSON.stringify({ id: "msg_1", type: "message", content: [] }), { + status: 200, + headers: { "content-type": "application/json" }, + }) + ); + + const response = await handleProxyRequest( + createContext("/v1/messages", { model: "claude-test", messages: [] }) + ); + + expect(response.status).toBe(200); + expect(response.headers.get("content-type")).toContain("application/json"); + expect(await response.json()).toEqual({ id: "msg_1", type: "message", content: [] }); + expect(boundary.send).toHaveBeenCalledOnce(); + }); + + it("returns synthesized SSE when the request is fake-stream eligible", async () => { + boundary.loadSettings.mockResolvedValue({ + ...defaultSettings, + fakeStreamingWhitelist: [{ model: "gpt-image-2", groupTags: [] }], + }); + boundary.send.mockResolvedValue( + new Response( + JSON.stringify({ + id: "msg_fake", + type: "message", + role: "assistant", + model: "claude-test", + content: [{ type: "text", text: "generated" }], + stop_reason: "end_turn", + usage: { input_tokens: 3, output_tokens: 1 }, + }), + { status: 200, headers: { "content-type": "application/json" } } + ) + ); + + const response = await handleProxyRequest( + createContext("/v1/messages", { model: "gpt-image-2", messages: [], stream: true }) + ); + const body = await response.text(); + + expect(response.status).toBe(200); + expect(response.headers.get("content-type")).toContain("text/event-stream"); + expect(body).toContain("event: message_start"); + expect(body).toContain('"text":"generated"'); + expect(body).toContain("event: message_stop"); + expect(boundary.send).toHaveBeenCalledOnce(); + }); + + it("normalizes Responses input and output at the public boundary", async () => { + boundary.send.mockImplementation( + async (session) => + new Response( + JSON.stringify({ + id: "resp_1", + object: "response", + echoed_input: session.request.message.input, + output: [{ type: "message", content: null }], + tools: null, + }), + { status: 200, headers: { "content-type": "application/json" } } + ) + ); + + const response = await handleProxyRequest( + createContext("/v1/responses", { model: "gpt-5", input: "hello" }) + ); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ + id: "resp_1", + object: "response", + echoed_input: [{ role: "user", content: [{ type: "input_text", text: "hello" }] }], + output: [{ type: "message", content: [] }], + tools: [], + }); + }); +}); diff --git a/tests/unit/proxy/response-handler-abort-listener-cleanup.test.ts b/tests/unit/proxy/response-handler-abort-listener-cleanup.test.ts index 73f6933b9..c18814926 100644 --- a/tests/unit/proxy/response-handler-abort-listener-cleanup.test.ts +++ b/tests/unit/proxy/response-handler-abort-listener-cleanup.test.ts @@ -18,9 +18,23 @@ vi.mock("@/app/v1/_lib/proxy/response-fixer", () => ({ vi.mock("@/lib/async-task-manager", () => ({ AsyncTaskManager: { - register: (_taskId: string, promise: Promise) => { + register: ( + _taskId: string, + factory: (signal: AbortSignal) => Promise, + options?: string | { abortController?: AbortController } + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + let promise: Promise; + try { + promise = Promise.resolve(factory(controller.signal)); + } catch (error) { + promise = Promise.reject(error); + } testState.asyncTasks.push(promise); - return new AbortController(); + return controller; }, touch: () => true, cleanup: testState.cleanupTask, @@ -96,13 +110,21 @@ vi.mock("@/lib/endpoint-circuit-breaker", () => ({ vi.mock("@/repository/message", () => ({ updateMessageRequestCostWithBreakdown: vi.fn(), updateMessageRequestDetails: vi.fn(), + updateMessageRequestDetailsDurably: vi.fn(async () => {}), + updateMessageRequestDetailsIfUnfinalized: vi.fn(async () => {}), updateMessageRequestDuration: vi.fn(), })); async function drainAsyncTasks(): Promise { while (testState.asyncTasks.length > 0) { const tasks = testState.asyncTasks.splice(0); - await Promise.allSettled(tasks); + const settlements = await Promise.allSettled(tasks); + const failures = settlements + .filter((settlement): settlement is PromiseRejectedResult => settlement.status === "rejected") + .map((settlement) => settlement.reason); + if (failures.length > 0) { + throw new AggregateError(failures, "Async task failed during test drain"); + } await new Promise((resolve) => setTimeout(resolve, 0)); } } diff --git a/tests/unit/proxy/response-handler-bill-non-success.test.ts b/tests/unit/proxy/response-handler-bill-non-success.test.ts index f4d04466e..3b9ae2674 100644 --- a/tests/unit/proxy/response-handler-bill-non-success.test.ts +++ b/tests/unit/proxy/response-handler-bill-non-success.test.ts @@ -16,7 +16,18 @@ vi.mock("@/lib/logger", () => ({ vi.mock("@/lib/async-task-manager", () => ({ AsyncTaskManager: { - register: () => new AbortController(), + register: ( + _taskId: string, + factory: (signal: AbortSignal) => Promise, + options?: string | { abortController?: AbortController } + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + void Promise.resolve(factory(controller.signal)).catch(() => {}); + return controller; + }, touch: () => true, cleanup: () => {}, cancel: () => {}, diff --git a/tests/unit/proxy/response-handler-client-abort-drain.test.ts b/tests/unit/proxy/response-handler-client-abort-drain.test.ts index f1a0bfb20..ee817d348 100644 --- a/tests/unit/proxy/response-handler-client-abort-drain.test.ts +++ b/tests/unit/proxy/response-handler-client-abort-drain.test.ts @@ -6,13 +6,22 @@ import { } from "@/app/v1/_lib/proxy/response-handler"; import { ProxySession } from "@/app/v1/_lib/proxy/session"; import { setDeferredStreamingFinalization } from "@/app/v1/_lib/proxy/stream-finalization"; -import { AsyncTaskManager } from "@/lib/async-task-manager"; +import { AsyncTaskManager, shutdownAllAsyncTasks } from "@/lib/async-task-manager"; +import { recordFailure } from "@/lib/circuit-breaker"; import { emitProxyLangfuseTrace } from "@/lib/langfuse/emit-proxy-trace"; +import { RateLimitService } from "@/lib/rate-limit"; import { SessionManager } from "@/lib/session-manager"; -import { updateMessageRequestDetails, updateMessageRequestDuration } from "@/repository/message"; +import { + updateMessageRequestCostWithBreakdown, + updateMessageRequestDetails, + updateMessageRequestDetailsDurably, + updateMessageRequestDetailsIfUnfinalized, + updateMessageRequestDuration, +} from "@/repository/message"; import type { Provider } from "@/types/provider"; const asyncTasks: Promise[] = []; +const registeredTasks: Array<{ taskType: string; promise: Promise }> = []; const STREAM_STATS_HEAD_BYTES_FOR_TEST = 1024 * 1024; vi.mock("@/app/v1/_lib/proxy/response-fixer", () => ({ @@ -23,14 +32,40 @@ vi.mock("@/app/v1/_lib/proxy/response-fixer", () => ({ vi.mock("@/lib/async-task-manager", () => ({ AsyncTaskManager: { - register: vi.fn((_taskId: string, promise: Promise) => { - asyncTasks.push(promise); - return new AbortController(); - }), + register: vi.fn( + ( + _taskId: string, + factory: (signal: AbortSignal) => Promise, + options?: string | { abortController?: AbortController; taskType?: string } + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + let promise: Promise; + try { + promise = Promise.resolve(factory(controller.signal)); + } catch (error) { + promise = Promise.reject(error); + } + asyncTasks.push(promise); + registeredTasks.push({ + taskType: typeof options === "object" ? (options.taskType ?? "unknown") : "unknown", + promise, + }); + return controller; + } + ), touch: vi.fn(() => true), cleanup: vi.fn(), cancel: vi.fn(), }, + shutdownAllAsyncTasks: vi.fn(async () => { + while (asyncTasks.length > 0) { + const tasks = asyncTasks.splice(0, asyncTasks.length); + await Promise.allSettled(tasks); + } + }), })); vi.mock("@/lib/config/system-settings-cache", () => ({ @@ -68,6 +103,7 @@ vi.mock("@/lib/rate-limit", () => ({ trackCost: vi.fn(), trackUserDailyCost: vi.fn(), decrementLeaseBudget: vi.fn(), + settleLeaseBudgets: vi.fn(), }, })); @@ -114,6 +150,8 @@ vi.mock("@/lib/endpoint-circuit-breaker", () => ({ vi.mock("@/repository/message", () => ({ updateMessageRequestCostWithBreakdown: vi.fn(), updateMessageRequestDetails: vi.fn(), + updateMessageRequestDetailsDurably: vi.fn(), + updateMessageRequestDetailsIfUnfinalized: vi.fn(), updateMessageRequestDuration: vi.fn(), })); @@ -282,6 +320,203 @@ function createResponsesSse(): Response { }); } +function createPullTrackedResponsesSse(): { + response: Response; + getPullCount: () => number; +} { + const encoder = new TextEncoder(); + const totalChunks = 32; + let index = 0; + let pullCount = 0; + + const stream = new ReadableStream({ + pull(controller) { + pullCount++; + if (index < totalChunks - 1) { + controller.enqueue( + encoder.encode( + `event: response.output_text.delta\ndata: ${JSON.stringify({ + type: "response.output_text.delta", + delta: `chunk-${index++}`, + })}\n\n` + ) + ); + return; + } + if (index++ === totalChunks - 1) { + controller.enqueue( + encoder.encode( + `event: response.completed\ndata: ${JSON.stringify({ + type: "response.completed", + response: { + id: "resp_pull_tracked", + model: "gpt-5.4-mini-2026-03-17", + usage: { input_tokens: 463, output_tokens: 11 }, + }, + })}\n\n` + ) + ); + return; + } + controller.close(); + }, + }); + + return { + response: new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }), + getPullCount: () => pullCount, + }; +} + +function createControllableTransportErrorResponsesSse(): { + response: Response; + fail: () => void; +} { + let controller: ReadableStreamDefaultController | null = null; + const stream = new ReadableStream({ + start(streamController) { + controller = streamController; + }, + }); + + return { + response: new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }), + fail() { + const error = Object.assign(new Error("socket closed after client cancel"), { + code: "ECONNRESET", + }); + controller?.error(error); + }, + }; +} + +function createControllableEmptyResponsesSse(): { + response: Response; + close: () => void; +} { + let controller: ReadableStreamDefaultController | null = null; + const stream = new ReadableStream({ + start(streamController) { + controller = streamController; + }, + }); + + return { + response: new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }), + close() { + try { + controller?.close(); + } catch { + // The timeout hard-cancel path may already have closed the source. + } + }, + }; +} + +function createControllableIdleTimeoutResponsesSse(): { + response: Response; + failIdle: () => void; +} { + const encoder = new TextEncoder(); + let controller: ReadableStreamDefaultController | null = null; + const stream = new ReadableStream({ + start(streamController) { + controller = streamController; + streamController.enqueue( + encoder.encode( + `event: response.output_text.delta\ndata: ${JSON.stringify({ + type: "response.output_text.delta", + delta: "partial", + })}\n\n` + ) + ); + }, + }); + + return { + response: new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }), + failIdle() { + const error = new Error("streaming_idle"); + error.name = "AbortError"; + controller?.error(error); + }, + }; +} + +function createAbortInsensitiveHangingResponsesSse(): { + response: Response; + close: () => void; +} { + let controller: ReadableStreamDefaultController | null = null; + const stream = new ReadableStream({ + start(streamController) { + controller = streamController; + }, + }); + + return { + response: new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }), + close() { + try { + controller?.close(); + } catch { + // The hard-cap path may already have cancelled and closed the source. + } + }, + }; +} + +function createAbortInsensitivePostChunkHangingResponsesSse(): Response { + const encoder = new TextEncoder(); + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue( + encoder.encode( + `event: response.output_text.delta\ndata: ${JSON.stringify({ + type: "response.output_text.delta", + delta: "first", + })}\n\n` + ) + ); + }, + }), + { + status: 200, + headers: { "content-type": "text/event-stream" }, + } + ); +} + +function createEmptyResponsesSse(): Response { + return new Response( + new ReadableStream({ + start(controller) { + controller.close(); + }, + }), + { + status: 200, + headers: { "content-type": "text/event-stream" }, + } + ); +} + function createResponsesJson(): Response { return new Response( JSON.stringify({ @@ -647,18 +882,266 @@ function createCompletedThenAbortedClaudeSse(): Response { }); } +async function expectAllFulfilled(tasks: readonly Promise[]): Promise { + const settlements = await Promise.allSettled(tasks); + const rejections = settlements + .filter((settlement): settlement is PromiseRejectedResult => settlement.status === "rejected") + .map((settlement) => settlement.reason); + if (rejections.length > 0) { + throw new AggregateError(rejections, "Unexpected async task rejection"); + } +} + async function drainAsyncTasks(): Promise { while (asyncTasks.length > 0) { const tasks = asyncTasks.splice(0, asyncTasks.length); - await Promise.allSettled(tasks); + await expectAllFulfilled(tasks); await new Promise((resolve) => setTimeout(resolve, 0)); } } +function getRegisteredTask(taskType: string): Promise | undefined { + return registeredTasks.filter((task) => task.taskType === taskType).at(-1)?.promise; +} + +function createDeferred() { + let resolve!: (value: T | PromiseLike) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((promiseResolve, promiseReject) => { + resolve = promiseResolve; + reject = promiseReject; + }); + return { promise, resolve, reject }; +} + +function createRecoveryContinuationResponse(contentType: string): Response { + const processingError = new Error("upstream body processing failed"); + let nameReads = 0; + Object.defineProperty(processingError, "name", { + configurable: true, + get() { + nameReads++; + if (nameReads === 1) { + throw new Error("error classification failed"); + } + return "Error"; + }, + }); + + return new Response( + new ReadableStream({ + start(controller) { + controller.error(processingError); + }, + }), + { + status: 200, + headers: { "content-type": contentType }, + } + ); +} + +async function expectPromiseToRemainPending(promise: Promise): Promise { + let outcome: "pending" | "resolved" | "rejected" = "pending"; + void promise.then( + () => { + outcome = "resolved"; + }, + () => { + outcome = "rejected"; + } + ); + await new Promise((resolve) => setImmediate(resolve)); + expect(outcome).toBe("pending"); +} + +async function expectTaskToResolveWithoutWaiting(promise: Promise): Promise { + let outcome: "pending" | "resolved" | "rejected" = "pending"; + void promise.then( + () => { + outcome = "resolved"; + }, + () => { + outcome = "rejected"; + } + ); + await new Promise((resolve) => setImmediate(resolve)); + expect(outcome).toBe("resolved"); +} + +function publishCommitObserver( + details: Parameters[1], + options: Parameters[2] +): void { + try { + const result = options?.onCommitted?.(details); + if (result) void Promise.resolve(result).catch(() => undefined); + } catch { + // Test mock mirrors the repository's commit-observer boundary. + } +} + +function createAbortableNonStreamResponse(signal: AbortSignal): Response { + const stream = new ReadableStream({ + start(controller) { + const abort = () => { + const reason = + signal.reason instanceof Error ? signal.reason : new Error("non-stream response aborted"); + controller.error(reason); + }; + signal.addEventListener("abort", abort, { once: true }); + if (signal.aborted) { + abort(); + } + }, + }); + + return new Response(stream, { + status: 200, + headers: { "content-type": "application/json" }, + }); +} + describe("ProxyResponseHandler stream client abort finalization", () => { beforeEach(() => { asyncTasks.splice(0, asyncTasks.length); + registeredTasks.splice(0, registeredTasks.length); vi.clearAllMocks(); + vi.mocked(updateMessageRequestDetailsDurably).mockImplementation( + async (_id, details, options) => { + try { + const result = options?.onCommitted?.(details); + if (result) void Promise.resolve(result).catch(() => undefined); + } catch { + // Test mock mirrors the repository's commit-observer boundary. + } + return true; + } + ); + vi.mocked(updateMessageRequestDetailsIfUnfinalized).mockImplementation( + async (_id, details, options) => { + try { + const result = options?.onCommitted?.(details); + if (result) void Promise.resolve(result).catch(() => undefined); + } catch { + // Test mock mirrors the repository's commit-observer boundary. + } + return true; + } + ); + }); + + it("propagates unexpected registered task rejections during drain", async () => { + const failure = new Error("factory task failed"); + AsyncTaskManager.register("rejecting-test-task", async () => { + throw failure; + }); + + let rejection: unknown; + try { + await drainAsyncTasks(); + } catch (error) { + rejection = error; + } + + expect(rejection).toBeInstanceOf(AggregateError); + expect((rejection as AggregateError).errors).toEqual([failure]); + }); + + it("keeps shutdown pending until generic non-stream recovery persistence settles", async () => { + const recoveryStarted = createDeferred(); + const releaseRecovery = createDeferred(); + vi.mocked(updateMessageRequestDetailsIfUnfinalized).mockImplementationOnce(async () => { + recoveryStarted.resolve(); + await releaseRecovery.promise; + return true; + }); + let shutdownPromise: Promise | undefined; + + try { + const session = createSession(new AbortController().signal); + const response = await ProxyResponseHandler.dispatch( + session, + createRecoveryContinuationResponse("application/json") + ); + + expect(response.status).toBe(200); + await recoveryStarted.promise; + const processingTask = getRegisteredTask("non-stream-processing"); + expect(processingTask).toBeDefined(); + shutdownPromise = shutdownAllAsyncTasks(); + + await expectPromiseToRemainPending(shutdownPromise); + expect(updateMessageRequestDetails).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); + + releaseRecovery.resolve(); + await shutdownPromise; + await expect(processingTask).rejects.toThrow("error classification failed"); + + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + durationMs: expect.any(Number), + statusCode: 500, + errorMessage: "Error: error classification failed", + }) + ); + } finally { + releaseRecovery.resolve(); + await shutdownPromise?.catch(() => {}); + await drainAsyncTasks(); + } + }); + + it("keeps shutdown pending until generic stream recovery persistence settles", async () => { + const recoveryStarted = createDeferred(); + const releaseRecovery = createDeferred(); + vi.mocked(updateMessageRequestDetailsIfUnfinalized).mockImplementationOnce(async () => { + recoveryStarted.resolve(); + await releaseRecovery.promise; + return true; + }); + let downstreamRead: Promise | undefined; + let shutdownPromise: Promise | undefined; + + try { + const session = createSession(new AbortController().signal); + const response = await ProxyResponseHandler.dispatch( + session, + createRecoveryContinuationResponse("text/event-stream") + ); + + expect(response.status).toBe(200); + downstreamRead = response.text().catch(() => "stream failed"); + await recoveryStarted.promise; + const processingTask = getRegisteredTask("stream-processing"); + expect(processingTask).toBeDefined(); + shutdownPromise = shutdownAllAsyncTasks(); + + await expectPromiseToRemainPending(shutdownPromise); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); + + releaseRecovery.resolve(); + await shutdownPromise; + await downstreamRead; + await expect(processingTask).rejects.toThrow("error classification failed"); + + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + durationMs: expect.any(Number), + statusCode: 500, + errorMessage: "Error: error classification failed", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + } finally { + releaseRecovery.resolve(); + await shutdownPromise?.catch(() => {}); + await downstreamRead?.catch(() => {}); + await drainAsyncTasks(); + } }); it("copies Buffer-backed stream windows before retaining stats snapshots", () => { @@ -681,10 +1164,9 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect(snapshot.text).not.toContain("zzzzzzzzzzzzzzzz"); }); - it("does not apply the default stale cleanup when stream idle timeout is disabled", async () => { - const controller = new AbortController(); - const session = createSession(controller.signal); - session.provider.streamingIdleTimeoutMs = 0; + it("does not pull the upstream stream before downstream demand", async () => { + const clientController = new AbortController(); + const session = createSession(clientController.signal); setDeferredStreamingFinalization(session, { providerId: 1, providerName: "avemujica-responses", @@ -697,48 +1179,22 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); + const tracked = createPullTrackedResponsesSse(); - await ProxyResponseHandler.dispatch(session, createResponsesSse()); - await drainAsyncTasks(); - - const streamRegisterCall = vi.mocked(AsyncTaskManager.register).mock.calls.find((call) => { - const options = call[2] as { taskType?: string } | undefined; - return options?.taskType === "stream-processing"; - }); - - expect(streamRegisterCall).toBeDefined(); - expect(streamRegisterCall?.[2]).toEqual( - expect.objectContaining({ - staleTimeoutMs: Number.POSITIVE_INFINITY, - }) - ); - }); + const downstream = await ProxyResponseHandler.dispatch(session, tracked.response); + await new Promise((resolve) => setImmediate(resolve)); + await new Promise((resolve) => setImmediate(resolve)); - it("does not apply the default stale cleanup when non-stream request timeout is disabled", async () => { - const controller = new AbortController(); - const session = createSession(controller.signal); - session.provider.requestTimeoutNonStreamingMs = 0; + expect(tracked.getPullCount()).toBeLessThanOrEqual(4); - await ProxyResponseHandler.dispatch(session, createResponsesJson()); + clientController.abort(new Error("test cleanup")); + await downstream.body?.cancel().catch(() => {}); await drainAsyncTasks(); - - const nonStreamRegisterCall = vi.mocked(AsyncTaskManager.register).mock.calls.find((call) => { - const options = call[2] as { taskType?: string } | undefined; - return options?.taskType === "non-stream-processing"; - }); - - expect(nonStreamRegisterCall).toBeDefined(); - expect(nonStreamRegisterCall?.[2]).toEqual( - expect.objectContaining({ - staleTimeoutMs: Number.POSITIVE_INFINITY, - }) - ); }); - it("finalizes a complete upstream responses stream as success when the downstream client already closed", async () => { - const controller = new AbortController(); - controller.abort(); - const session = createSession(controller.signal); + it("keeps upstream lookahead bounded while the downstream consumer is paused", async () => { + const clientController = new AbortController(); + const session = createSession(clientController.signal); setDeferredStreamingFinalization(session, { providerId: 1, providerName: "avemujica-responses", @@ -751,70 +1207,106 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); + const tracked = createPullTrackedResponsesSse(); - await ProxyResponseHandler.dispatch(session, createResponsesSse()); - await drainAsyncTasks(); + const downstream = await ProxyResponseHandler.dispatch(session, tracked.response); + const reader = downstream.body?.getReader(); + expect(reader).toBeDefined(); + const first = await reader?.read(); + expect(first?.done).toBe(false); + await new Promise((resolve) => setImmediate(resolve)); + await new Promise((resolve) => setImmediate(resolve)); - expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); - expect(updateMessageRequestDuration).toHaveBeenCalledWith(123, expect.any(Number)); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( + expect(tracked.getPullCount()).toBeLessThanOrEqual(3); + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); + + await reader?.cancel("test cleanup"); + await drainAsyncTasks(); + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( 123, - expect.objectContaining({ - statusCode: 200, - inputTokens: 463, - outputTokens: 11, - }) + expect.objectContaining({ durationMs: expect.any(Number) }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); }); - it("keeps stream accounting bounded for oversized successful streams", async () => { - const controller = new AbortController(); - const session = createSession(controller.signal); - session.sessionId = "session_large"; - Object.assign(session, { - shouldPersistSessionDebugArtifacts: () => true, - }); - setDeferredStreamingFinalization(session, { - providerId: 1, - providerName: "avemujica-responses", - providerPriority: 1, - attemptNumber: 1, - totalProvidersAttempted: 1, - isFirstAttempt: true, - isFailoverSuccess: false, - endpointId: 42, - endpointUrl: "https://api.test.invalid/v1", - upstreamStatusCode: 200, - }); - - await ProxyResponseHandler.dispatch(session, createOversizedResponsesSse()); - await drainAsyncTasks(); - - expect(updateMessageRequestDetails).toHaveBeenCalledWith( - 123, - expect.objectContaining({ - statusCode: 200, - inputTokens: 463, - outputTokens: 11, - }) - ); - expect(SessionManager.storeSessionResponse).not.toHaveBeenCalled(); + it("does not count a pending chunk for an active slow consumer as Provider idle", async () => { + vi.useFakeTimers(); + const responseController = new AbortController(); + let reader: ReadableStreamDefaultReader | undefined; + try { + const clientController = new AbortController(); + const session = createSession(clientController.signal); + session.provider.streamingIdleTimeoutMs = 5_000; + Object.assign(session, { responseController }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + const tracked = createPullTrackedResponsesSse(); - const traceCall = vi.mocked(emitProxyLangfuseTrace).mock.calls.at(-1); - expect(traceCall).toBeDefined(); - const traceData = traceCall?.[1]; - const responseText = traceData?.responseText ?? ""; - expect(responseText).toContain("[cch_truncated]"); - expect(responseText.length).toBeLessThan(10 * 1024 * 1024 + 1024); + const downstream = await ProxyResponseHandler.dispatch(session, tracked.response); + reader = downstream.body?.getReader(); + expect(reader).toBeDefined(); + + const decoder = new TextDecoder(); + const first = await reader?.read(); + expect(first?.done).toBe(false); + let responseText = first?.value ? decoder.decode(first.value, { stream: true }) : ""; + await vi.advanceTimersByTimeAsync(0); + + const pullsBeforePause = tracked.getPullCount(); + expect(pullsBeforePause).toBeGreaterThanOrEqual(2); + + await vi.advanceTimersByTimeAsync(10_000); + + expect(responseController.signal.aborted).toBe(false); + expect(tracked.getPullCount()).toBe(pullsBeforePause); + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); + + while (true) { + const result = await reader?.read(); + if (!result || result.done) break; + responseText += decoder.decode(result.value, { stream: true }); + } + responseText += decoder.decode(); + + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + + expect(responseText).toContain("event: response.completed"); + expect(responseController.signal.aborted).toBe(false); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 200, + inputTokens: 463, + outputTokens: 11, + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect(recordFailure).not.toHaveBeenCalled(); + } finally { + await reader?.cancel("test cleanup").catch(() => {}); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + vi.clearAllTimers(); + vi.useRealTimers(); + } }); - it("decodes an untruncated stream as contiguous UTF-8 across the head/tail split", async () => { + it("does not apply the default stale cleanup when stream idle timeout is disabled", async () => { const controller = new AbortController(); const session = createSession(controller.signal); - session.sessionId = "session_utf8_boundary"; - Object.assign(session, { - shouldPersistSessionDebugArtifacts: () => true, - }); + session.provider.streamingIdleTimeoutMs = 0; setDeferredStreamingFinalization(session, { providerId: 1, providerName: "avemujica-responses", @@ -828,33 +1320,28 @@ describe("ProxyResponseHandler stream client abort finalization", () => { upstreamStatusCode: 200, }); - await ProxyResponseHandler.dispatch(session, createUtf8SplitHeadTailResponsesSse()); + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); await drainAsyncTasks(); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( - 123, + const streamRegisterCall = vi.mocked(AsyncTaskManager.register).mock.calls.find((call) => { + const options = call[2] as { taskType?: string } | undefined; + return options?.taskType === "stream-processing"; + }); + + expect(streamRegisterCall).toBeDefined(); + expect(streamRegisterCall?.[2]).toEqual( expect.objectContaining({ - statusCode: 200, - inputTokens: 463, - outputTokens: 11, + staleTimeoutMs: Number.POSITIVE_INFINITY, }) ); - - const traceCall = vi.mocked(emitProxyLangfuseTrace).mock.calls.at(-1); - expect(traceCall).toBeDefined(); - const responseText = traceCall?.[1].responseText ?? ""; - expect(responseText).toContain("界"); - expect(responseText).not.toContain("\uFFFD"); - expect(responseText).not.toContain("[cch_truncated]"); }); - it("keeps usage when a terminal responses event is split across tail chunk eviction", async () => { + it("does not apply the generic stale watchdog when stream timeouts are enabled", async () => { const controller = new AbortController(); const session = createSession(controller.signal); - session.sessionId = "session_split_tail"; - Object.assign(session, { - shouldPersistSessionDebugArtifacts: () => true, - }); + session.provider.streamingIdleTimeoutMs = 5_000; + session.provider.firstByteTimeoutStreamingMs = 2_000; setDeferredStreamingFinalization(session, { providerId: 1, providerName: "avemujica-responses", @@ -868,23 +1355,22 @@ describe("ProxyResponseHandler stream client abort finalization", () => { upstreamStatusCode: 200, }); - await ProxyResponseHandler.dispatch(session, createSplitTailBoundaryResponsesSse()); + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); await drainAsyncTasks(); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( - 123, - expect.objectContaining({ - statusCode: 200, - inputTokens: 463, - outputTokens: 11, - }) + const streamRegisterCall = vi.mocked(AsyncTaskManager.register).mock.calls.find((call) => { + const options = call[2] as { taskType?: string } | undefined; + return options?.taskType === "stream-processing"; + }); + + expect(streamRegisterCall?.[2]).toEqual( + expect.objectContaining({ staleTimeoutMs: Number.POSITIVE_INFINITY }) ); - expect(SessionManager.storeSessionResponse).not.toHaveBeenCalled(); }); - it("reclassifies a client-aborted stream as success when final usage was already received", async () => { + it("clears the first-byte timeout when an empty upstream stream reaches EOF", async () => { const controller = new AbortController(); - controller.abort(); const session = createSession(controller.signal); setDeferredStreamingFinalization(session, { providerId: 1, @@ -899,29 +1385,16 @@ describe("ProxyResponseHandler stream client abort finalization", () => { upstreamStatusCode: 200, }); - await ProxyResponseHandler.dispatch(session, createCompletedThenErroredResponsesSse()); + const downstream = await ProxyResponseHandler.dispatch(session, createEmptyResponsesSse()); + await downstream.text(); await drainAsyncTasks(); - expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( - 123, - expect.objectContaining({ - statusCode: 200, - inputTokens: 463, - outputTokens: 11, - providerChain: [ - expect.objectContaining({ - reason: "request_success", - statusCode: 200, - }), - ], - }) - ); + expect(session.recordTtfb).not.toHaveBeenCalled(); + expect(session.clearResponseTimeout).toHaveBeenCalledTimes(1); }); - it("keeps a genuinely aborted upstream responses stream as 499", async () => { + it("does not overwrite persisted terminal details when a later side effect fails", async () => { const controller = new AbortController(); - controller.abort(); const session = createSession(controller.signal); setDeferredStreamingFinalization(session, { providerId: 1, @@ -935,28 +1408,27 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); + vi.mocked(emitProxyLangfuseTrace).mockImplementationOnce(() => { + throw new Error("final trace failed"); + }); - await ProxyResponseHandler.dispatch(session, createErroredResponsesSse()); + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); await drainAsyncTasks(); - expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( - 123, - expect.objectContaining({ - statusCode: 499, - errorMessage: "CLIENT_ABORTED", - }) - ); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); + expect(emitProxyLangfuseTrace).toHaveBeenCalledTimes(1); }); - it("keeps a truncated client-aborted Claude stream as 499 despite message_start usage (U01)", async () => { - const controller = new AbortController(); - controller.abort(); - const session = createSession(controller.signal, { - providerType: "anthropic", - originalFormat: "claude", - endpoint: "/v1/messages", - model: "claude-x", + it("settles stream processing before deferred success side effects even when tracing fails", async () => { + const clientController = new AbortController(); + const session = createSession(clientController.signal); + session.sessionId = "deferred-success-side-effects"; + let resolveBinding!: (result: { updated: boolean; reason: string }) => void; + let markBindingStarted!: () => void; + const bindingStarted = new Promise((resolve) => { + markBindingStarted = resolve; }); setDeferredStreamingFinalization(session, { providerId: 1, @@ -970,35 +1442,60 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); + vi.mocked(SessionManager.updateSessionBindingSmart).mockImplementationOnce(() => { + markBindingStarted(); + return new Promise((resolve) => { + resolveBinding = resolve; + }); + }); + vi.mocked(emitProxyLangfuseTrace).mockImplementationOnce(() => { + throw new Error("final trace failed"); + }); - await ProxyResponseHandler.dispatch(session, createTruncatedClaudeSse()); - await drainAsyncTasks(); + try { + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); + await bindingStarted; - expect(updateMessageRequestDetails).toHaveBeenCalledWith( - 123, - expect.objectContaining({ - statusCode: 499, - errorMessage: "CLIENT_ABORTED", - }) - ); - // Must NOT have been recorded as a billed 200 success. - const calls = (updateMessageRequestDetails as unknown as { mock: { calls: unknown[][] } }).mock - .calls; - const recorded = calls.find((c) => (c[0] as number) === 123)?.[1] as - | { statusCode?: number } - | undefined; - expect(recorded?.statusCode).not.toBe(200); + const streamProcessingTask = getRegisteredTask("stream-processing"); + expect(streamProcessingTask).toBeDefined(); + await expectTaskToResolveWithoutWaiting(streamProcessingTask as Promise); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 200, + providerId: 1, + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect(emitProxyLangfuseTrace).toHaveBeenCalledTimes(1); + expect( + vi.mocked(AsyncTaskManager.register).mock.calls.some((call) => { + const options = call[2] as { taskType?: string } | undefined; + return options?.taskType === "post-terminal-side-effects"; + }) + ).toBe(true); + } finally { + resolveBinding?.({ updated: false, reason: "test cleanup" }); + await drainAsyncTasks(); + } }); - it("bills a complete-then-aborted Claude stream as success on the message_stop marker (U01)", async () => { - const controller = new AbortController(); - controller.abort(); - const session = createSession(controller.signal, { - providerType: "anthropic", + it("persists an upstream failure before waiting for Session cleanup side effects", async () => { + const clientController = new AbortController(); + const session = createSession(clientController.signal, { + providerType: "claude", originalFormat: "claude", endpoint: "/v1/messages", model: "claude-x", }); + session.sessionId = "deferred-failure-side-effects"; + let resolveCleanup!: () => void; + let markCleanupStarted!: () => void; + const cleanupStarted = new Promise((resolve) => { + markCleanupStarted = resolve; + }); setDeferredStreamingFinalization(session, { providerId: 1, providerName: "avemujica-responses", @@ -1011,27 +1508,67 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); + vi.mocked(SessionManager.clearSessionProvider).mockImplementationOnce(() => { + markCleanupStarted(); + return new Promise((resolve) => { + resolveCleanup = resolve; + }); + }); - await ProxyResponseHandler.dispatch(session, createCompletedThenAbortedClaudeSse()); - await drainAsyncTasks(); + try { + const downstream = await ProxyResponseHandler.dispatch(session, createTruncatedClaudeSse()); + await downstream.text().catch(() => "client stream closed"); + await cleanupStarted; - expect(updateMessageRequestDetails).toHaveBeenCalledWith( - 123, - expect.objectContaining({ - statusCode: 200, - inputTokens: 463, - }) - ); + const streamProcessingTask = getRegisteredTask("stream-processing"); + expect(streamProcessingTask).toBeDefined(); + await expectTaskToResolveWithoutWaiting(streamProcessingTask as Promise); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 502, + errorMessage: "STREAM_UPSTREAM_ABORTED", + providerId: 1, + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect( + vi.mocked(AsyncTaskManager.register).mock.calls.some((call) => { + const options = call[2] as { taskType?: string } | undefined; + return options?.taskType === "post-terminal-side-effects"; + }) + ).toBe(true); + } finally { + resolveCleanup?.(); + await drainAsyncTasks(); + } }); - it("keeps client-abort drain independent from a small idle timeout while chunks are active", async () => { + it("releases stream listeners and timeouts before deferred persistence finishes", async () => { vi.useFakeTimers(); + const clientController = new AbortController(); + const addSpy = vi.spyOn(clientController.signal, "addEventListener"); + const removeSpy = vi.spyOn(clientController.signal, "removeEventListener"); + const responseController = new AbortController(); + let resolvePersistence!: () => void; + const blockedPersistence = new Promise((resolve) => { + resolvePersistence = () => resolve(true); + }); + let markPersistenceStarted!: () => void; + const persistenceStarted = new Promise((resolve) => { + markPersistenceStarted = resolve; + }); + let timersRestored = false; try { - const clientController = new AbortController(); - const upstreamController = new AbortController(); const session = createSession(clientController.signal); + const releaseAgent = vi.fn(); session.provider.streamingIdleTimeoutMs = 5_000; - Object.assign(session, { responseController: upstreamController }); + Object.assign(session, { responseController, releaseAgent }); + const responseTimeoutId = setTimeout(() => { + responseController.abort(new Error("response timeout was not cleared")); + }, 5_000); + session.clearResponseTimeout = vi.fn(() => clearTimeout(responseTimeoutId)); setDeferredStreamingFinalization(session, { providerId: 1, providerName: "avemujica-responses", @@ -1044,42 +1581,62 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); + vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce(() => { + markPersistenceStarted(); + return blockedPersistence; + }); - await ProxyResponseHandler.dispatch( + const downstream = await ProxyResponseHandler.dispatch( session, - createActiveHangingResponsesSse(upstreamController.signal) + createPullTrackedResponsesSse().response ); - clientController.abort(); + await downstream.text(); + await persistenceStarted; - await vi.advanceTimersByTimeAsync(59_000); - expect(upstreamController.signal.aborted).toBe(false); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDetailsIfUnfinalized).not.toHaveBeenCalled(); + expect(session.clearResponseTimeout).toHaveBeenCalledTimes(1); + expect(releaseAgent).toHaveBeenCalledTimes(1); - await vi.advanceTimersByTimeAsync(1_000); - const tasks = asyncTasks.splice(0, asyncTasks.length); - await Promise.allSettled(tasks); + const abortAddCalls = addSpy.mock.calls.filter(([type]) => type === "abort"); + expect(abortAddCalls).toHaveLength(1); + expect(removeSpy).toHaveBeenCalledWith("abort", abortAddCalls[0][1]); - expect(upstreamController.signal.aborted).toBe(true); - expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( + await vi.advanceTimersByTimeAsync(5_000); + expect(responseController.signal.aborted).toBe(false); + + clientController.abort(new Error("late client disconnect")); + await vi.advanceTimersByTimeAsync(60_000); + expect(responseController.signal.aborted).toBe(false); + expect(updateMessageRequestDetailsIfUnfinalized).not.toHaveBeenCalled(); + + resolvePersistence(); + vi.clearAllTimers(); + vi.useRealTimers(); + timersRestored = true; + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( 123, - expect.objectContaining({ - statusCode: 499, - errorMessage: "CLIENT_ABORTED", - }) + expect.objectContaining({ durationMs: expect.any(Number), statusCode: 200 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); } finally { - vi.useRealTimers(); + resolvePersistence(); + if (!timersRestored) { + vi.clearAllTimers(); + vi.useRealTimers(); + } + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); } }); - it("uses idle timeout for client-aborted streams that hang before the first chunk", async () => { + it("bounds deferred stream finalization when persistence never settles", async () => { vi.useFakeTimers(); try { const clientController = new AbortController(); - const upstreamController = new AbortController(); const session = createSession(clientController.signal); - session.provider.streamingIdleTimeoutMs = 5_000; - Object.assign(session, { responseController: upstreamController }); setDeferredStreamingFinalization(session, { providerId: 1, providerName: "avemujica-responses", @@ -1092,42 +1649,43 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); - - await ProxyResponseHandler.dispatch( - session, - createPreBodyHangingResponsesSse(upstreamController.signal) + vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce( + () => new Promise(() => {}) ); - clientController.abort(); - await vi.advanceTimersByTimeAsync(4_999); - expect(upstreamController.signal.aborted).toBe(false); + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); + await vi.advanceTimersByTimeAsync(120_000); + await vi.advanceTimersByTimeAsync(0); - await vi.advanceTimersByTimeAsync(1); const tasks = asyncTasks.splice(0, asyncTasks.length); - await Promise.allSettled(tasks); - - expect(upstreamController.signal.aborted).toBe(true); - expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( + await expectAllFulfilled(tasks); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledWith( 123, expect.objectContaining({ - statusCode: 499, - errorMessage: "CLIENT_ABORTED", - }) + durationMs: expect.any(Number), + statusCode: 500, + errorMessage: "Error: stream_finalization_timeout", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); } finally { + vi.clearAllTimers(); vi.useRealTimers(); } }); - it("preserves an existing idle deadline when the client aborts after a chunk", async () => { + it("keeps a late fallback commit observable after the failure deadline", async () => { vi.useFakeTimers(); + const flushMicrotasks = (remaining = 32): Promise => + remaining === 0 + ? Promise.resolve() + : Promise.resolve().then(() => flushMicrotasks(remaining - 1)); + let resolveFallback!: () => void; + let committedCallback: (() => void | Promise) | undefined; try { const clientController = new AbortController(); - const upstreamController = new AbortController(); const session = createSession(clientController.signal); - session.provider.streamingIdleTimeoutMs = 5_000; - Object.assign(session, { responseController: upstreamController }); setDeferredStreamingFinalization(session, { providerId: 1, providerName: "avemujica-responses", @@ -1140,42 +1698,976 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); - - await ProxyResponseHandler.dispatch( - session, - createHangingResponsesSse(upstreamController.signal) + vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce( + () => new Promise(() => {}) + ); + vi.mocked(updateMessageRequestDetailsIfUnfinalized).mockImplementationOnce( + async (_id, details, options) => { + committedCallback = options?.onCommitted; + await new Promise((resolve) => { + resolveFallback = resolve; + }); + await options?.onCommitted?.(details); + return true; + } ); - await vi.advanceTimersByTimeAsync(0); - await vi.advanceTimersByTimeAsync(4_999); - expect(upstreamController.signal.aborted).toBe(false); - clientController.abort(); - await vi.advanceTimersByTimeAsync(1); + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); + await vi.advanceTimersByTimeAsync(120_000); + await flushMicrotasks(); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); + + await vi.advanceTimersByTimeAsync(5_000); + await flushMicrotasks(); + expect(asyncTasks.length).toBeGreaterThanOrEqual(2); + + resolveFallback(); + await flushMicrotasks(); + expect(committedCallback).toBeTypeOf("function"); const tasks = asyncTasks.splice(0, asyncTasks.length); - await Promise.allSettled(tasks); + await expectAllFulfilled(tasks); - expect(upstreamController.signal.aborted).toBe(true); - expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( - 123, - expect.objectContaining({ - statusCode: 499, - errorMessage: "CLIENT_ABORTED", - }) - ); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); } finally { + resolveFallback?.(); + await flushMicrotasks(); + vi.clearAllTimers(); vi.useRealTimers(); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); } }); - it("caps client-abort drain at 60s when the upstream stream hangs", async () => { - vi.useFakeTimers(); - try { - const clientController = new AbortController(); - const upstreamController = new AbortController(); - const session = createSession(clientController.signal); + it("persists stream duration in the same durable terminal patch", async () => { + const clientController = new AbortController(); + const session = createSession(clientController.signal); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); + await drainAsyncTasks(); + + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ durationMs: expect.any(Number), statusCode: 200 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + }); + + it("does not let a timeout fallback overwrite a late primary terminal write", async () => { + vi.useFakeTimers(); + const flushMicrotasks = (remaining = 32): Promise => + remaining === 0 + ? Promise.resolve() + : Promise.resolve().then(() => flushMicrotasks(remaining - 1)); + let terminalStatusCode: number | null = null; + let resolvePrimaryDetails!: () => void; + let markPrimaryDetailsStarted!: () => void; + const primaryDetailsStarted = new Promise((resolve) => { + markPrimaryDetailsStarted = resolve; + }); + let resolveFallback!: () => void; + let markFallbackStarted!: () => void; + const fallbackStarted = new Promise((resolve) => { + markFallbackStarted = resolve; + }); + try { + const clientController = new AbortController(); + const session = createSession(clientController.signal); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + vi.mocked(updateMessageRequestDetailsDurably) + .mockImplementationOnce((_id, details) => { + markPrimaryDetailsStarted(); + return new Promise((resolve) => { + resolvePrimaryDetails = () => { + terminalStatusCode = details.statusCode ?? null; + resolve(true); + }; + }); + }) + .mockImplementation(async (_id, details) => { + terminalStatusCode = details.statusCode ?? null; + return true; + }); + vi.mocked(updateMessageRequestDetailsIfUnfinalized).mockImplementation( + async (_id, details) => { + markFallbackStarted(); + if (terminalStatusCode === null) { + terminalStatusCode = details.statusCode ?? null; + } + return new Promise((resolve) => { + resolveFallback = () => resolve(false); + }); + } + ); + + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); + await primaryDetailsStarted; + await vi.advanceTimersByTimeAsync(120_000); + await fallbackStarted; + + resolvePrimaryDetails(); + await flushMicrotasks(); + expect(terminalStatusCode).toBe(200); + + resolveFallback(); + await flushMicrotasks(); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); + expect(terminalStatusCode).toBe(200); + } finally { + resolvePrimaryDetails?.(); + resolveFallback?.(); + await flushMicrotasks(); + vi.clearAllTimers(); + vi.useRealTimers(); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + } + }); + + it("does not apply the default stale cleanup when non-stream request timeout is disabled", async () => { + const controller = new AbortController(); + const session = createSession(controller.signal); + session.provider.requestTimeoutNonStreamingMs = 0; + + await ProxyResponseHandler.dispatch(session, createResponsesJson()); + await drainAsyncTasks(); + + const nonStreamRegisterCall = vi.mocked(AsyncTaskManager.register).mock.calls.find((call) => { + const options = call[2] as { taskType?: string } | undefined; + return options?.taskType === "non-stream-processing"; + }); + + expect(nonStreamRegisterCall).toBeDefined(); + expect(nonStreamRegisterCall?.[2]).toEqual( + expect.objectContaining({ + staleTimeoutMs: Number.POSITIVE_INFINITY, + }) + ); + }); + + it("finalizes a complete upstream responses stream as success when the downstream client already closed", async () => { + const controller = new AbortController(); + controller.abort(); + const session = createSession(controller.signal); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await drainAsyncTasks(); + + expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 200, + inputTokens: 463, + outputTokens: 11, + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + }); + + it("keeps stream accounting bounded for oversized successful streams", async () => { + const controller = new AbortController(); + const session = createSession(controller.signal); + session.sessionId = "session_large"; + Object.assign(session, { + shouldPersistSessionDebugArtifacts: () => true, + }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + const downstream = await ProxyResponseHandler.dispatch(session, createOversizedResponsesSse()); + await downstream.text(); + await drainAsyncTasks(); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 200, + inputTokens: 463, + outputTokens: 11, + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect(SessionManager.storeSessionResponse).not.toHaveBeenCalled(); + + const traceCall = vi.mocked(emitProxyLangfuseTrace).mock.calls.at(-1); + expect(traceCall).toBeDefined(); + const traceData = traceCall?.[1]; + const responseText = traceData?.responseText ?? ""; + expect(responseText).toContain("[cch_truncated]"); + expect(responseText.length).toBeLessThan(10 * 1024 * 1024 + 1024); + }); + + it("decodes an untruncated stream as contiguous UTF-8 across the head/tail split", async () => { + const controller = new AbortController(); + const session = createSession(controller.signal); + session.sessionId = "session_utf8_boundary"; + Object.assign(session, { + shouldPersistSessionDebugArtifacts: () => true, + }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + const downstream = await ProxyResponseHandler.dispatch( + session, + createUtf8SplitHeadTailResponsesSse() + ); + await downstream.text(); + await drainAsyncTasks(); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 200, + inputTokens: 463, + outputTokens: 11, + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + + const traceCall = vi.mocked(emitProxyLangfuseTrace).mock.calls.at(-1); + expect(traceCall).toBeDefined(); + const responseText = traceCall?.[1].responseText ?? ""; + expect(responseText).toContain("界"); + expect(responseText).not.toContain("\uFFFD"); + expect(responseText).not.toContain("[cch_truncated]"); + }); + + it("keeps usage when a terminal responses event is split across tail chunk eviction", async () => { + const controller = new AbortController(); + const session = createSession(controller.signal); + session.sessionId = "session_split_tail"; + Object.assign(session, { + shouldPersistSessionDebugArtifacts: () => true, + }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + const downstream = await ProxyResponseHandler.dispatch( + session, + createSplitTailBoundaryResponsesSse() + ); + await downstream.text(); + await drainAsyncTasks(); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 200, + inputTokens: 463, + outputTokens: 11, + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect(SessionManager.storeSessionResponse).not.toHaveBeenCalled(); + }); + + it("reclassifies a client-aborted stream as success when final usage was already received", async () => { + const controller = new AbortController(); + controller.abort(); + const session = createSession(controller.signal); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + await ProxyResponseHandler.dispatch(session, createCompletedThenErroredResponsesSse()); + await drainAsyncTasks(); + + expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 200, + inputTokens: 463, + outputTokens: 11, + providerChain: [ + expect.objectContaining({ + reason: "request_success", + statusCode: 200, + }), + ], + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + }); + + it("keeps a genuinely aborted upstream responses stream as 499", async () => { + const controller = new AbortController(); + controller.abort(); + const session = createSession(controller.signal); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + await ProxyResponseHandler.dispatch(session, createErroredResponsesSse()); + await drainAsyncTasks(); + + expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 499, + errorMessage: "CLIENT_ABORTED", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + }); + + it("keeps a truncated client-aborted Claude stream as 499 despite message_start usage (U01)", async () => { + const controller = new AbortController(); + controller.abort(); + const session = createSession(controller.signal, { + providerType: "anthropic", + originalFormat: "claude", + endpoint: "/v1/messages", + model: "claude-x", + }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + await ProxyResponseHandler.dispatch(session, createTruncatedClaudeSse()); + await drainAsyncTasks(); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 499, + errorMessage: "CLIENT_ABORTED", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + // Must NOT have been recorded as a billed 200 success. + const calls = ( + updateMessageRequestDetailsDurably as unknown as { mock: { calls: unknown[][] } } + ).mock.calls; + const recorded = calls.find((c) => (c[0] as number) === 123)?.[1] as + | { statusCode?: number } + | undefined; + expect(recorded?.statusCode).not.toBe(200); + }); + + it("bills a complete-then-aborted Claude stream as success on the message_stop marker (U01)", async () => { + const controller = new AbortController(); + controller.abort(); + const session = createSession(controller.signal, { + providerType: "anthropic", + originalFormat: "claude", + endpoint: "/v1/messages", + model: "claude-x", + }); + Object.assign(session, { + sessionId: "session-complete-then-aborted", + getResolvedPricingByBillingSource: vi.fn(async () => ({ + resolvedModelName: "claude-x", + resolvedPricingProviderKey: "anthropic", + source: "local_manual" as const, + priceData: { + input_cost_per_token: 0.000003, + output_cost_per_token: 0.000015, + }, + })), + }); + Object.assign(session.authState?.user ?? {}, { + dailyResetTime: "00:00", + dailyResetMode: "fixed", + limit5hResetMode: "rolling", + }); + Object.assign(session.authState?.key ?? {}, { + dailyResetTime: "00:00", + dailyResetMode: "fixed", + limit5hResetMode: "rolling", + }); + Object.assign(session.provider ?? {}, { + limit5hResetMode: "rolling", + }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + await ProxyResponseHandler.dispatch(session, createCompletedThenAbortedClaudeSse()); + await drainAsyncTasks(); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 200, + inputTokens: 463, + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + + const expectedCost = 0.001554; + const dbCostCall = vi.mocked(updateMessageRequestCostWithBreakdown).mock.calls.at(-1); + expect(dbCostCall).toBeDefined(); + expect(dbCostCall?.[0]).toBe(123); + expect(String(dbCostCall?.[1])).toBe("0.001554"); + expect(dbCostCall?.[2]).toEqual( + expect.objectContaining({ + input: "0.001389", + output: "0.000165", + total: "0.001554", + }) + ); + + expect(RateLimitService.trackCost).toHaveBeenCalledWith( + 2, + 1, + "session-complete-then-aborted", + expectedCost, + expect.objectContaining({ + userId: 1, + userResetTime: "00:00", + userResetMode: "fixed", + requestId: 123, + }) + ); + expect(RateLimitService.trackUserDailyCost).not.toHaveBeenCalled(); + + expect(RateLimitService.settleLeaseBudgets).toHaveBeenCalledTimes(1); + expect(RateLimitService.settleLeaseBudgets).toHaveBeenCalledWith({ + requestId: 123, + cost: expectedCost, + entities: { + key: expect.objectContaining({ id: 2 }), + user: expect.objectContaining({ id: 1 }), + provider: expect.objectContaining({ id: 1 }), + }, + }); + expect(RateLimitService.decrementLeaseBudget).not.toHaveBeenCalled(); + }); + + it("keeps client-abort drain independent from a small idle timeout while chunks are active", async () => { + vi.useFakeTimers(); + try { + const clientController = new AbortController(); + const upstreamController = new AbortController(); + const session = createSession(clientController.signal); + session.provider.streamingIdleTimeoutMs = 5_000; + Object.assign(session, { responseController: upstreamController }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + await ProxyResponseHandler.dispatch( + session, + createActiveHangingResponsesSse(upstreamController.signal) + ); + clientController.abort(); + + await vi.advanceTimersByTimeAsync(59_000); + expect(upstreamController.signal.aborted).toBe(false); + + await vi.advanceTimersByTimeAsync(1_000); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + + expect(upstreamController.signal.aborted).toBe(true); + expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 499, + errorMessage: "CLIENT_ABORTED", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + } finally { + vi.useRealTimers(); + } + }); + + it("uses idle timeout for client-aborted streams that hang before the first chunk", async () => { + vi.useFakeTimers(); + try { + const clientController = new AbortController(); + const upstreamController = new AbortController(); + const session = createSession(clientController.signal); + session.provider.streamingIdleTimeoutMs = 5_000; + Object.assign(session, { responseController: upstreamController }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + await ProxyResponseHandler.dispatch( + session, + createPreBodyHangingResponsesSse(upstreamController.signal) + ); + clientController.abort(); + + await vi.advanceTimersByTimeAsync(4_999); + expect(upstreamController.signal.aborted).toBe(false); + + await vi.advanceTimersByTimeAsync(1); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + + expect(upstreamController.signal.aborted).toBe(true); + expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 499, + errorMessage: "CLIENT_ABORTED", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + } finally { + vi.useRealTimers(); + } + }); + + it("preserves an existing idle deadline when the client aborts after a chunk", async () => { + vi.useFakeTimers(); + try { + const clientController = new AbortController(); + const upstreamController = new AbortController(); + const session = createSession(clientController.signal); + session.provider.streamingIdleTimeoutMs = 5_000; + Object.assign(session, { responseController: upstreamController }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + const downstream = await ProxyResponseHandler.dispatch( + session, + createHangingResponsesSse(upstreamController.signal) + ); + const downstreamReader = downstream.body?.getReader(); + expect(downstreamReader).toBeDefined(); + await downstreamReader?.read(); + await vi.advanceTimersByTimeAsync(0); + await vi.advanceTimersByTimeAsync(4_999); + expect(upstreamController.signal.aborted).toBe(false); + + clientController.abort(); + await vi.advanceTimersByTimeAsync(1); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + + expect(upstreamController.signal.aborted).toBe(true); + expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 499, + errorMessage: "CLIENT_ABORTED", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + } finally { + vi.useRealTimers(); + } + }); + + it("preserves an earlier Provider idle timeout when the client aborts before source settlement", async () => { + vi.useFakeTimers(); + try { + const clientController = new AbortController(); + const upstreamController = new AbortController(); + const session = createSession(clientController.signal); + session.provider.streamingIdleTimeoutMs = 5_000; + Object.assign(session, { responseController: upstreamController }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + const controlled = createControllableIdleTimeoutResponsesSse(); + + const downstream = await ProxyResponseHandler.dispatch(session, controlled.response); + const downstreamReader = downstream.body?.getReader(); + expect(downstreamReader).toBeDefined(); + void downstreamReader?.closed.catch(() => {}); + await downstreamReader?.read(); + await vi.advanceTimersByTimeAsync(0); + + await vi.advanceTimersByTimeAsync(5_000); + expect(upstreamController.signal.aborted).toBe(true); + + clientController.abort(); + controlled.failIdle(); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 502, + errorMessage: "STREAM_IDLE_TIMEOUT", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect(recordFailure).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + + it("caps client-abort drain at 60s when the upstream stream hangs", async () => { + vi.useFakeTimers(); + try { + const clientController = new AbortController(); + const upstreamController = new AbortController(); + const session = createSession(clientController.signal); + session.provider.streamingIdleTimeoutMs = 120_000; + Object.assign(session, { responseController: upstreamController }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + await ProxyResponseHandler.dispatch( + session, + createHangingResponsesSse(upstreamController.signal) + ); + clientController.abort(); + + await vi.advanceTimersByTimeAsync(60_000); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + + expect(upstreamController.signal.aborted).toBe(true); + expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 499, + errorMessage: "CLIENT_ABORTED", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + } finally { + vi.useRealTimers(); + } + }); + + it("caps body-cancel drain even when the session abort signal does not fire", async () => { + vi.useFakeTimers(); + const upstreamController = new AbortController(); + try { + const clientController = new AbortController(); + const session = createSession(clientController.signal); + session.provider.streamingIdleTimeoutMs = 120_000; + Object.assign(session, { responseController: upstreamController }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + const downstream = await ProxyResponseHandler.dispatch( + session, + createHangingResponsesSse(upstreamController.signal) + ); + await downstream.body?.cancel("body_cancel_only"); + expect(clientController.signal.aborted).toBe(false); + + await vi.advanceTimersByTimeAsync(60_000); + expect(upstreamController.signal.aborted).toBe(true); + + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 499, + errorMessage: "CLIENT_ABORTED", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + } finally { + if (!upstreamController.signal.aborted) { + upstreamController.abort(new Error("test cleanup")); + } + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + vi.useRealTimers(); + } + }); + + it("classifies a transport error after body cancel as client-aborted", async () => { + const clientController = new AbortController(); + const session = createSession(clientController.signal); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + const upstream = createControllableTransportErrorResponsesSse(); + + const downstream = await ProxyResponseHandler.dispatch(session, upstream.response); + await downstream.body?.cancel("body_cancel_only"); + upstream.fail(); + await drainAsyncTasks(); + + expect(clientController.signal.aborted).toBe(false); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 499, + errorMessage: "CLIENT_ABORTED", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + }); + + it("keeps an earlier Provider transport error when the client signal aborts later", async () => { + const clientController = new AbortController(); + const session = createSession(clientController.signal); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + const upstream = createControllableTransportErrorResponsesSse(); + + await ProxyResponseHandler.dispatch(session, upstream.response); + upstream.fail(); + queueMicrotask(() => clientController.abort(new Error("late client cleanup"))); + await drainAsyncTasks(); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 502, + errorMessage: "STREAM_UPSTREAM_ABORTED", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + }); + + it("persists a fallback when client-detached transport finalization fails", async () => { + const clientController = new AbortController(); + const session = createSession(clientController.signal); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + vi.mocked(updateMessageRequestDetailsDurably).mockRejectedValueOnce( + new Error("client-detached terminal details failed") + ); + const upstream = createControllableTransportErrorResponsesSse(); + + const downstream = await ProxyResponseHandler.dispatch(session, upstream.response); + await downstream.body?.cancel("body_cancel_only"); + upstream.fail(); + await drainAsyncTasks(); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDetails).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); + expect(vi.mocked(updateMessageRequestDetailsIfUnfinalized).mock.calls[0]).toEqual([ + 123, + expect.objectContaining({ + durationMs: expect.any(Number), + statusCode: 499, + errorMessage: "CLIENT_ABORTED", + providerId: 1, + providerChain: [ + expect.objectContaining({ + id: 1, + name: "avemujica-responses", + statusCode: 499, + errorMessage: "CLIENT_ABORTED", + }), + ], + }), + expect.objectContaining({ onCommitted: expect.any(Function) }), + ]); + }); + + it("hard-caps body-cancel drain when the source ignores AbortSignals", async () => { + vi.useFakeTimers(); + const upstream = createAbortInsensitiveHangingResponsesSse(); + try { + const clientController = new AbortController(); + const responseController = new AbortController(); + const session = createSession(clientController.signal); session.provider.streamingIdleTimeoutMs = 120_000; - Object.assign(session, { responseController: upstreamController }); + Object.assign(session, { responseController }); setDeferredStreamingFinalization(session, { providerId: 1, providerName: "avemujica-responses", @@ -1189,27 +2681,633 @@ describe("ProxyResponseHandler stream client abort finalization", () => { upstreamStatusCode: 200, }); - await ProxyResponseHandler.dispatch( + const downstream = await ProxyResponseHandler.dispatch(session, upstream.response); + await downstream.body?.cancel("body_cancel_only"); + await vi.advanceTimersByTimeAsync(60_000); + + const tasks = asyncTasks.splice(0, asyncTasks.length); + const outcome = await Promise.race([ + expectAllFulfilled(tasks).then(() => "settled" as const), + new Promise<"pending">((resolve) => setImmediate(() => resolve("pending"))), + ]); + expect(outcome).toBe("settled"); + expect(responseController.signal.aborted).toBe(true); + } finally { + upstream.close(); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + vi.useRealTimers(); + } + }); + + it("settles an abort-insensitive source immediately after Provider idle timeout", async () => { + vi.useFakeTimers(); + try { + const clientController = new AbortController(); + const responseController = new AbortController(); + const session = createSession(clientController.signal); + session.provider.streamingIdleTimeoutMs = 5_000; + Object.assign(session, { responseController }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + const downstream = await ProxyResponseHandler.dispatch( session, - createHangingResponsesSse(upstreamController.signal) + createAbortInsensitivePostChunkHangingResponsesSse() ); - clientController.abort(); + const reader = downstream.body?.getReader(); + expect(reader).toBeDefined(); + await reader?.read(); + void reader?.closed.catch(() => {}); - await vi.advanceTimersByTimeAsync(60_000); + await vi.advanceTimersByTimeAsync(5_000); const tasks = asyncTasks.splice(0, asyncTasks.length); - await Promise.allSettled(tasks); + await expectAllFulfilled(tasks); - expect(upstreamController.signal.aborted).toBe(true); - expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( + expect(responseController.signal.aborted).toBe(true); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( 123, expect.objectContaining({ - statusCode: 499, - errorMessage: "CLIENT_ABORTED", - }) + statusCode: 502, + errorMessage: "STREAM_IDLE_TIMEOUT", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); } finally { vi.useRealTimers(); } }); + + it("settles an abort-insensitive source immediately after response timeout", async () => { + const clientController = new AbortController(); + const responseController = new AbortController(); + const session = createSession(clientController.signal); + Object.assign(session, { responseController }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + const upstream = createAbortInsensitiveHangingResponsesSse(); + + const downstream = await ProxyResponseHandler.dispatch(session, upstream.response); + const downstreamOutcome = downstream.text().catch(() => "client stream closed"); + const timeoutError = new Error("response timeout"); + timeoutError.name = "AbortError"; + responseController.abort(timeoutError); + await downstreamOutcome; + await drainAsyncTasks(); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 502, + errorMessage: "STREAM_RESPONSE_TIMEOUT", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + upstream.close(); + }); + + it("preserves a response timeout when the client aborts before source settlement", async () => { + const clientController = new AbortController(); + const responseController = new AbortController(); + const session = createSession(clientController.signal); + Object.assign(session, { responseController }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + const upstream = createControllableEmptyResponsesSse(); + + const downstream = await ProxyResponseHandler.dispatch(session, upstream.response); + const downstreamOutcome = downstream.text().catch(() => "client stream closed"); + const timeoutError = new Error("response timeout"); + timeoutError.name = "AbortError"; + responseController.abort(timeoutError); + clientController.abort(new Error("late client disconnect")); + upstream.close(); + await downstreamOutcome; + await drainAsyncTasks(); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 502, + errorMessage: "STREAM_RESPONSE_TIMEOUT", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect(recordFailure).toHaveBeenCalledTimes(1); + expect(recordFailure).toHaveBeenCalledWith( + 1, + expect.objectContaining({ message: "STREAM_RESPONSE_TIMEOUT" }) + ); + expect(SessionManager.updateSessionBindingSmart).not.toHaveBeenCalled(); + }); + + it("waits for durable non-stream failure details before mutating the Provider circuit", async () => { + const durableAck = createDeferred(); + vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce( + async (_id, details, options) => { + await durableAck.promise; + try { + const result = options?.onCommitted?.(details); + if (result) void Promise.resolve(result).catch(() => undefined); + } catch { + // Test mock mirrors the repository's commit-observer boundary. + } + return true; + } + ); + const session = createSession(new AbortController().signal); + const response = new Response('{"error":{"message":"provider failed"}}', { + status: 500, + headers: { "content-type": "application/json" }, + }); + + await ProxyResponseHandler.dispatch(session, response); + while (vi.mocked(updateMessageRequestDetailsDurably).mock.calls.length === 0) { + await new Promise((resolve) => setImmediate(resolve)); + } + + expect(session.getProviderChain()).toEqual([ + expect.objectContaining({ + id: 1, + reason: "retry_failed", + statusCode: 500, + }), + ]); + expect(vi.mocked(updateMessageRequestDetailsDurably).mock.calls[0]?.[1]).toEqual( + expect.objectContaining({ + statusCode: 500, + errorMessage: "FAKE_200_JSON_ERROR_MESSAGE_NON_EMPTY", + }) + ); + const recordFailureCallsBeforeAck = vi.mocked(recordFailure).mock.calls.length; + + durableAck.resolve(); + await drainAsyncTasks(); + + expect(recordFailureCallsBeforeAck).toBe(0); + expect(recordFailure).toHaveBeenCalledTimes(1); + expect(recordFailure).toHaveBeenCalledWith( + 1, + expect.objectContaining({ message: "FAKE_200_JSON_ERROR_MESSAGE_NON_EMPTY" }) + ); + }); + + it("reuses the original non-stream terminal details in the conditional fallback", async () => { + vi.mocked(updateMessageRequestDetailsDurably).mockRejectedValueOnce( + new Error("durable acknowledgement failed") + ); + const session = createSession(new AbortController().signal); + const response = new Response('{"error":{"message":"provider failed"}}', { + status: 500, + headers: { "content-type": "application/json" }, + }); + + await ProxyResponseHandler.dispatch(session, response); + await drainAsyncTasks(); + + expect(updateMessageRequestDetails).not.toHaveBeenCalled(); + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + durationMs: expect.any(Number), + statusCode: 500, + errorMessage: "FAKE_200_JSON_ERROR_MESSAGE_NON_EMPTY", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + durationMs: expect.any(Number), + statusCode: 500, + errorMessage: "FAKE_200_JSON_ERROR_MESSAGE_NON_EMPTY", + providerId: 1, + providerChain: [ + expect.objectContaining({ + id: 1, + reason: "retry_failed", + statusCode: 500, + errorMessage: "FAKE_200_JSON_ERROR_MESSAGE_NON_EMPTY", + }), + ], + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + }); + + it.each([ + ["response timeout", "timeout"], + ["client abort", "client"], + ] as const)( + "uses the conditional fallback when the non-stream %s finalizer durable write rejects", + async (_name, abortSource) => { + vi.mocked(updateMessageRequestDetailsDurably).mockRejectedValueOnce( + new Error("durable finalizer acknowledgement failed") + ); + const clientController = new AbortController(); + const responseController = new AbortController(); + const session = createSession(clientController.signal); + Object.assign(session, { responseController }); + const response = createAbortableNonStreamResponse( + abortSource === "timeout" ? responseController.signal : clientController.signal + ); + + await ProxyResponseHandler.dispatch(session, response); + const abortError = new Error(`non-stream ${abortSource}`); + abortError.name = "AbortError"; + if (abortSource === "timeout") { + responseController.abort(abortError); + } else { + clientController.abort(abortError); + } + await drainAsyncTasks(); + + expect(updateMessageRequestDetails).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: abortSource === "timeout" ? 502 : 499, + ...(abortSource === "timeout" + ? { errorMessage: expect.stringContaining("non-stream timeout") } + : {}), + providerId: 1, + providerChain: + abortSource === "timeout" + ? [ + expect.objectContaining({ + id: 1, + statusCode: 502, + errorMessage: expect.stringContaining("non-stream timeout"), + }), + ] + : [], + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + } + ); + + it("rejects non-stream processing when both terminal persistence attempts fail", async () => { + vi.mocked(updateMessageRequestDetailsDurably).mockRejectedValueOnce( + new Error("durable acknowledgement failed") + ); + vi.mocked(updateMessageRequestDetailsIfUnfinalized).mockRejectedValueOnce( + new Error("conditional fallback failed") + ); + const session = createSession(new AbortController().signal); + const response = new Response('{"error":{"message":"provider failed"}}', { + status: 500, + headers: { "content-type": "application/json" }, + }); + + await ProxyResponseHandler.dispatch(session, response); + const processingTask = getRegisteredTask("non-stream-processing"); + + expect(processingTask).toBeDefined(); + await expect(processingTask).rejects.toThrow("conditional fallback failed"); + expect(updateMessageRequestDetails).not.toHaveBeenCalled(); + }); + + it("waits for durable Gemini non-stream failure details before mutating the Provider circuit", async () => { + const durableAck = createDeferred(); + vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce( + async (_id, details, options) => { + await durableAck.promise; + publishCommitObserver(details, options); + return true; + } + ); + const session = createSession(new AbortController().signal, { + providerType: "gemini", + originalFormat: "gemini", + endpoint: "/v1beta/models/gemini-2.0-flash:generateContent", + model: "gemini-2.0-flash", + }); + const response = new Response('{"error":{"message":"provider failed"}}', { + status: 500, + headers: { "content-type": "application/json" }, + }); + + const returned = await ProxyResponseHandler.dispatch(session, response); + expect(returned).toBe(response); + while (vi.mocked(updateMessageRequestDetailsDurably).mock.calls.length === 0) { + await new Promise((resolve) => setImmediate(resolve)); + } + + expect(session.getProviderChain()).toEqual([ + expect.objectContaining({ + id: 1, + reason: "retry_failed", + statusCode: 500, + }), + ]); + const recordFailureCallsBeforeAck = vi.mocked(recordFailure).mock.calls.length; + + durableAck.resolve(); + await drainAsyncTasks(); + + expect(recordFailureCallsBeforeAck).toBe(0); + expect(recordFailure).toHaveBeenCalledTimes(1); + expect(recordFailure).toHaveBeenCalledWith( + 1, + expect.objectContaining({ message: "FAKE_200_JSON_ERROR_MESSAGE_NON_EMPTY" }) + ); + }); + + it.each([ + ["ordinary", {}], + [ + "Gemini passthrough", + { + providerType: "gemini", + originalFormat: "gemini", + endpoint: "/v1beta/models/gemini-2.0-flash:generateContent", + model: "gemini-2.0-flash", + }, + ], + ] as const)( + "keeps non-stream 404 out of the Provider circuit for %s responses", + async (_name, overrides) => { + const session = createSession(new AbortController().signal, overrides); + const response = new Response('{"error":{"message":"model not found"}}', { + status: 404, + headers: { "content-type": "application/json" }, + }); + + await ProxyResponseHandler.dispatch(session, response); + await drainAsyncTasks(); + + expect(recordFailure).not.toHaveBeenCalled(); + expect(session.getProviderChain()).toEqual([ + expect.objectContaining({ + id: 1, + reason: "resource_not_found", + statusCode: 404, + }), + ]); + } + ); + + it("persists Gemini non-stream duration atomically with terminal stats", async () => { + const session = createSession(new AbortController().signal, { + providerType: "gemini", + originalFormat: "gemini", + endpoint: "/v1beta/models/gemini-2.0-flash:generateContent", + model: "gemini-2.0-flash", + }); + const response = new Response('{"candidates":[]}', { + status: 200, + headers: { "content-type": "application/json" }, + }); + + await ProxyResponseHandler.dispatch(session, response); + await drainAsyncTasks(); + + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ durationMs: expect.any(Number), statusCode: 200 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + }); + + it("persists one durable 502 before Provider circuit mutation on non-stream response timeout", async () => { + const durableAck = createDeferred(); + vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce( + async (_id, details, options) => { + await durableAck.promise; + publishCommitObserver(details, options); + return true; + } + ); + const responseController = new AbortController(); + const session = createSession(new AbortController().signal); + Object.assign(session, { responseController }); + const response = createAbortableNonStreamResponse(responseController.signal); + + await ProxyResponseHandler.dispatch(session, response); + const timeoutError = new Error("non-stream response timeout"); + timeoutError.name = "AbortError"; + responseController.abort(timeoutError); + while (vi.mocked(updateMessageRequestDetailsDurably).mock.calls.length === 0) { + await new Promise((resolve) => setImmediate(resolve)); + } + + const recordFailureCallsBeforeAck = vi.mocked(recordFailure).mock.calls.length; + const durablePayloadBeforeAck = vi.mocked(updateMessageRequestDetailsDurably).mock + .calls[0]?.[1]; + const ordinaryDetailsCallsBeforeAck = vi.mocked(updateMessageRequestDetails).mock.calls.length; + + durableAck.resolve(); + await drainAsyncTasks(); + + expect(recordFailureCallsBeforeAck).toBe(0); + expect(ordinaryDetailsCallsBeforeAck).toBe(0); + expect(durablePayloadBeforeAck).toEqual( + expect.objectContaining({ + statusCode: 502, + errorMessage: expect.stringContaining("non-stream response timeout"), + }) + ); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); + expect(recordFailure).toHaveBeenCalledTimes(1); + }); + + it("waits for durable non-stream details before updating the Codex cache binding", async () => { + const durableAck = createDeferred(); + vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce( + async (_id, details, options) => { + await durableAck.promise; + publishCommitObserver(details, options); + return true; + } + ); + vi.mocked(SessionManager.extractCodexPromptCacheKey).mockReturnValueOnce("cache-key-1"); + vi.mocked(SessionManager.updateSessionWithCodexCacheKey).mockResolvedValueOnce(undefined); + const session = createSession(new AbortController().signal); + session.sessionId = "codex-cache-binding-session"; + const response = new Response('{"id":"resp_1"}', { + status: 200, + headers: { "content-type": "application/json" }, + }); + + await ProxyResponseHandler.dispatch(session, response); + while (vi.mocked(updateMessageRequestDetailsDurably).mock.calls.length === 0) { + await new Promise((resolve) => setImmediate(resolve)); + } + + const cacheBindingCallsBeforeAck = vi.mocked(SessionManager.updateSessionWithCodexCacheKey).mock + .calls.length; + durableAck.resolve(); + await drainAsyncTasks(); + + expect(cacheBindingCallsBeforeAck).toBe(0); + expect(SessionManager.updateSessionWithCodexCacheKey).toHaveBeenCalledTimes(1); + expect(SessionManager.updateSessionWithCodexCacheKey).toHaveBeenCalledWith( + "codex-cache-binding-session", + "cache-key-1", + 1, + 2 + ); + }); + + it("publishes a successful stream Codex cache binding only after durable acknowledgement", async () => { + const durableAck = createDeferred(); + const cacheBinding = createDeferred(); + const cacheBindingStarted = createDeferred(); + vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce( + async (_id, details, options) => { + await durableAck.promise; + publishCommitObserver(details, options); + return true; + } + ); + vi.mocked(SessionManager.extractCodexPromptCacheKey).mockReturnValueOnce("stream-cache-key-1"); + vi.mocked(SessionManager.updateSessionWithCodexCacheKey).mockImplementationOnce(async () => { + cacheBindingStarted.resolve(); + await cacheBinding.promise; + }); + const session = createSession(new AbortController().signal); + session.sessionId = "stream-codex-cache-binding-session"; + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + try { + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); + while (vi.mocked(updateMessageRequestDetailsDurably).mock.calls.length === 0) { + await new Promise((resolve) => setImmediate(resolve)); + } + + expect(SessionManager.updateSessionWithCodexCacheKey).not.toHaveBeenCalled(); + + durableAck.resolve(); + await cacheBindingStarted.promise; + + expect(SessionManager.updateSessionWithCodexCacheKey).toHaveBeenCalledWith( + "stream-codex-cache-binding-session", + "stream-cache-key-1", + 1, + 2 + ); + const streamProcessingTask = getRegisteredTask("stream-processing"); + expect(streamProcessingTask).toBeDefined(); + await expectTaskToResolveWithoutWaiting(streamProcessingTask as Promise); + } finally { + durableAck.resolve(); + cacheBinding.resolve(); + await drainAsyncTasks(); + } + }); + + it("does not publish a stream Codex cache binding for a final non-2xx outcome", async () => { + vi.mocked(SessionManager.extractCodexPromptCacheKey).mockReturnValueOnce("stream-cache-key-2"); + const session = createSession(new AbortController().signal); + session.sessionId = "stream-codex-cache-binding-failure"; + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 500, + }); + const response = new Response(await createResponsesSse().text(), { + status: 500, + headers: { "content-type": "text/event-stream" }, + }); + + const downstream = await ProxyResponseHandler.dispatch(session, response); + await downstream.text(); + await drainAsyncTasks(); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ statusCode: 500 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect(SessionManager.updateSessionWithCodexCacheKey).not.toHaveBeenCalled(); + }); + + it("durably finalizes a Gemini non-stream passthrough body-read failure", async () => { + const session = createSession(new AbortController().signal, { + providerType: "gemini", + originalFormat: "gemini", + endpoint: "/v1beta/models/gemini-2.0-flash:generateContent", + model: "gemini-2.0-flash", + }); + const response = new Response( + new ReadableStream({ + start(controller) { + controller.error(new Error("Gemini non-stream body read failed")); + }, + }), + { + status: 200, + headers: { "content-type": "application/json" }, + } + ); + + await ProxyResponseHandler.dispatch(session, response); + await drainAsyncTasks(); + + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + durationMs: expect.any(Number), + statusCode: 502, + errorMessage: expect.stringContaining("Gemini non-stream body read failed"), + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect(recordFailure).toHaveBeenCalledTimes(1); + }); }); diff --git a/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts b/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts index 5b6f88a28..c3bdf3a5d 100644 --- a/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts +++ b/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts @@ -17,9 +17,23 @@ const asyncTasks: Promise[] = []; vi.mock("@/lib/async-task-manager", () => ({ AsyncTaskManager: { - register: (_taskId: string, promise: Promise) => { + register: ( + _taskId: string, + factory: (signal: AbortSignal) => Promise, + options?: string | { abortController?: AbortController } + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + let promise: Promise; + try { + promise = Promise.resolve(factory(controller.signal)); + } catch (error) { + promise = Promise.reject(error); + } asyncTasks.push(promise); - return new AbortController(); + return controller; }, touch: () => true, cleanup: () => {}, @@ -53,6 +67,8 @@ vi.mock("@/repository/message", () => ({ updateMessageRequestCost: vi.fn(), updateMessageRequestCostWithBreakdown: vi.fn(), updateMessageRequestDetails: vi.fn(), + updateMessageRequestDetailsDurably: vi.fn(), + updateMessageRequestDetailsIfUnfinalized: vi.fn(), updateMessageRequestDuration: vi.fn(), })); @@ -62,6 +78,8 @@ vi.mock("@/lib/session-manager", () => ({ storeSessionResponse: vi.fn(), clearSessionProvider: vi.fn(), extractCodexPromptCacheKey: vi.fn(), + updateSessionBindingSmart: vi.fn(), + updateSessionProvider: vi.fn(), updateSessionWithCodexCacheKey: vi.fn(), }, })); @@ -71,6 +89,7 @@ vi.mock("@/lib/rate-limit", () => ({ trackCost: vi.fn(), trackUserDailyCost: vi.fn(), decrementLeaseBudget: vi.fn(), + settleLeaseBudgets: vi.fn(), }, })); @@ -89,16 +108,21 @@ vi.mock("@/lib/proxy-status-tracker", () => ({ })); // Mock circuit breakers with tracked spies (vi.hoisted to avoid TDZ with vi.mock hoisting) -const { mockRecordFailure, mockRecordEndpointFailure, mockRecordEndpointSuccess } = vi.hoisted( - () => ({ - mockRecordFailure: vi.fn(), - mockRecordEndpointFailure: vi.fn(), - mockRecordEndpointSuccess: vi.fn(), - }) -); +const { + mockRecordFailure, + mockRecordSuccess, + mockRecordEndpointFailure, + mockRecordEndpointSuccess, +} = vi.hoisted(() => ({ + mockRecordFailure: vi.fn(), + mockRecordSuccess: vi.fn(), + mockRecordEndpointFailure: vi.fn(), + mockRecordEndpointSuccess: vi.fn(), +})); vi.mock("@/lib/circuit-breaker", () => ({ recordFailure: mockRecordFailure, + recordSuccess: mockRecordSuccess, })); vi.mock("@/lib/endpoint-circuit-breaker", () => ({ @@ -112,7 +136,11 @@ import { ProxySession } from "@/app/v1/_lib/proxy/session"; import { setDeferredStreamingFinalization } from "@/app/v1/_lib/proxy/stream-finalization"; import { getSystemSettings } from "@/repository/system-config"; import { findLatestPriceByModel } from "@/repository/model-price"; -import { updateMessageRequestDetails, updateMessageRequestDuration } from "@/repository/message"; +import { + updateMessageRequestDetails, + updateMessageRequestDetailsDurably, + updateMessageRequestDuration, +} from "@/repository/message"; import { SessionManager } from "@/lib/session-manager"; import { RateLimitService } from "@/lib/rate-limit"; import { SessionTracker } from "@/lib/session-tracker"; @@ -308,8 +336,10 @@ function createSuccessStreamResponse(): Response { } async function drainAsyncTasks(): Promise { - const tasks = asyncTasks.splice(0, asyncTasks.length); - await Promise.all(tasks); + while (asyncTasks.length > 0) { + const tasks = asyncTasks.splice(0, asyncTasks.length); + await Promise.all(tasks); + } } function setupCommonMocks() { @@ -327,17 +357,30 @@ function setupCommonMocks() { updatedAt: new Date(), }); vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(true); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(SessionManager.clearSessionProvider).mockResolvedValue(undefined); + vi.mocked(SessionManager.updateSessionUsage).mockResolvedValue(undefined); + vi.mocked(SessionManager.updateSessionBindingSmart).mockResolvedValue({ + updated: true, + reason: "test", + }); + vi.mocked(SessionManager.updateSessionProvider).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackCost).mockResolvedValue(undefined); vi.mocked(RateLimitService.trackUserDailyCost).mockResolvedValue(undefined); vi.mocked(RateLimitService.decrementLeaseBudget).mockResolvedValue({ success: true, newRemaining: 10, }); + vi.mocked(RateLimitService.settleLeaseBudgets).mockResolvedValue({ + requestId: "test", + status: "settled", + settlements: [], + }); vi.mocked(SessionTracker.refreshSession).mockResolvedValue(undefined); mockRecordFailure.mockResolvedValue(undefined); + mockRecordSuccess.mockResolvedValue(undefined); mockRecordEndpointFailure.mockResolvedValue(undefined); mockRecordEndpointSuccess.mockResolvedValue(undefined); } @@ -357,7 +400,8 @@ describe("Endpoint circuit breaker isolation", () => { setDeferredMeta(session, 42); const response = createFake200StreamResponse(); - await ProxyResponseHandler.dispatch(session, response); + const clientResponse = await ProxyResponseHandler.dispatch(session, response); + await clientResponse.text(); await drainAsyncTasks(); expect(mockRecordFailure).toHaveBeenCalledWith( @@ -365,7 +409,7 @@ describe("Endpoint circuit breaker isolation", () => { expect.objectContaining({ message: expect.stringContaining("FAKE_200") }) ); expect(mockRecordEndpointFailure).not.toHaveBeenCalled(); - expect(SessionManager.clearSessionProvider).toHaveBeenCalledWith("fake-session"); + expect(SessionManager.clearSessionProvider).toHaveBeenCalledWith("fake-session", 1); const chain = session.getProviderChain(); expect( @@ -385,7 +429,8 @@ describe("Endpoint circuit breaker isolation", () => { setDeferredMeta(session, 42); const response = createFake200StreamResponse(); - await ProxyResponseHandler.dispatch(session, response); + const clientResponse = await ProxyResponseHandler.dispatch(session, response); + await clientResponse.text(); await drainAsyncTasks(); expect(mockRecordFailure).toHaveBeenCalledWith( @@ -393,7 +438,7 @@ describe("Endpoint circuit breaker isolation", () => { expect.objectContaining({ message: expect.stringContaining("FAKE_200") }) ); expect(mockRecordEndpointFailure).not.toHaveBeenCalled(); - expect(SessionManager.clearSessionProvider).toHaveBeenCalledWith("fake-session"); + expect(SessionManager.clearSessionProvider).toHaveBeenCalledWith("fake-session", 1); expect(SessionManager.updateSessionUsage).not.toHaveBeenCalled(); expect(SessionTracker.refreshSession).not.toHaveBeenCalled(); }); @@ -403,12 +448,13 @@ describe("Endpoint circuit breaker isolation", () => { setDeferredMeta(session, 42); const response = createFake200StreamResponse("model not found"); - await ProxyResponseHandler.dispatch(session, response); + const clientResponse = await ProxyResponseHandler.dispatch(session, response); + await clientResponse.text(); await drainAsyncTasks(); expect(mockRecordFailure).not.toHaveBeenCalled(); expect(mockRecordEndpointFailure).not.toHaveBeenCalled(); - expect(SessionManager.clearSessionProvider).toHaveBeenCalledWith("fake-session"); + expect(SessionManager.clearSessionProvider).toHaveBeenCalledWith("fake-session", 1); const chain = session.getProviderChain(); expect( @@ -439,7 +485,8 @@ describe("Endpoint circuit breaker isolation", () => { }); const response = createNon200StreamResponse(429); - await ProxyResponseHandler.dispatch(session, response); + const clientResponse = await ProxyResponseHandler.dispatch(session, response); + await clientResponse.text(); await drainAsyncTasks(); expect(mockRecordFailure).toHaveBeenCalledWith(1, expect.any(Error)); @@ -451,11 +498,26 @@ describe("Endpoint circuit breaker isolation", () => { setDeferredMeta(session, 42); const response = createSuccessStreamResponse(); - await ProxyResponseHandler.dispatch(session, response); + const clientResponse = await ProxyResponseHandler.dispatch(session, response); + await clientResponse.text(); await drainAsyncTasks(); expect(mockRecordEndpointSuccess).toHaveBeenCalledWith(42); + expect(mockRecordSuccess).toHaveBeenCalledWith(1); expect(mockRecordEndpointFailure).not.toHaveBeenCalled(); + expect(SessionManager.updateSessionBindingSmart).toHaveBeenCalledWith( + "fake-session", + 1, + 10, + true, + false, + 456 + ); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 1, + expect.objectContaining({ statusCode: 200, providerId: 1 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); }); it("streaming success without endpointId should NOT call any endpoint circuit breaker function", async () => { @@ -463,7 +525,8 @@ describe("Endpoint circuit breaker isolation", () => { setDeferredMeta(session, null); const response = createSuccessStreamResponse(); - await ProxyResponseHandler.dispatch(session, response); + const clientResponse = await ProxyResponseHandler.dispatch(session, response); + await clientResponse.text(); await drainAsyncTasks(); expect(mockRecordEndpointSuccess).not.toHaveBeenCalled(); diff --git a/tests/unit/proxy/response-handler-exported-finalizers.test.ts b/tests/unit/proxy/response-handler-exported-finalizers.test.ts new file mode 100644 index 000000000..2fc324b1f --- /dev/null +++ b/tests/unit/proxy/response-handler-exported-finalizers.test.ts @@ -0,0 +1,202 @@ +import { Context } from "hono"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { + finalizeHedgeLoserBilling, + finalizeRequestStats, +} from "@/app/v1/_lib/proxy/response-handler"; +import { ProxySession } from "@/app/v1/_lib/proxy/session"; +import type { Provider } from "@/types/provider"; + +const mocks = vi.hoisted(() => ({ + addLoserCost: vi.fn<(id: number, cost: object, entry: object) => Promise>(), + durable: vi.fn<(id: number, details: object) => Promise>(), + updateCost: vi.fn<(id: number, cost: object, breakdown: object) => Promise>(), +})); + +vi.mock("@/lib/logger", () => ({ + logger: { + debug: vi.fn(), + error: vi.fn(), + info: vi.fn(), + trace: vi.fn(), + warn: vi.fn(), + }, +})); +vi.mock("@/repository/message", () => ({ + addMessageRequestHedgeLoserCost: mocks.addLoserCost, + updateMessageRequestCostWithBreakdown: mocks.updateCost, + updateMessageRequestDetails: vi.fn(), + updateMessageRequestDetailsDurably: mocks.durable, + updateMessageRequestDetailsIfUnfinalized: vi.fn(), + updateMessageRequestDuration: vi.fn(), + updateMessageRequestWinnerCost: vi.fn(), +})); + +const CREATED_AT = new Date(0); + +function createProvider(): Provider { + return { + activeTimeEnd: null, + activeTimeStart: null, + allowedClients: [], + allowedModels: null, + anthropicAdaptiveThinking: null, + anthropicMaxTokensPreference: null, + anthropicThinkingBudgetPreference: null, + blockedClients: [], + cacheTtlPreference: null, + cc: 0, + circuitBreakerFailureThreshold: 5, + circuitBreakerHalfOpenSuccessThreshold: 2, + circuitBreakerOpenDuration: 1_800_000, + codexImageGenerationPreference: null, + codexParallelToolCallsPreference: null, + codexReasoningEffortPreference: null, + codexReasoningSummaryPreference: null, + codexServiceTierPreference: null, + codexTextVerbosityPreference: null, + context1mPreference: null, + costMultiplier: 1, + createdAt: CREATED_AT, + customHeaders: null, + dailyResetMode: "fixed", + dailyResetTime: "00:00", + disableSessionReuse: false, + faviconUrl: null, + firstByteTimeoutStreamingMs: 0, + geminiGoogleSearchPreference: null, + groupPriorities: null, + groupTag: null, + id: 13, + isEnabled: true, + key: "provider-key", + limit5hResetMode: "fixed", + limit5hUsd: null, + limitConcurrentSessions: 0, + limitDailyUsd: null, + limitMonthlyUsd: null, + limitTotalUsd: null, + limitWeeklyUsd: null, + maxRetryAttempts: null, + mcpPassthroughType: "none", + mcpPassthroughUrl: null, + modelRedirects: null, + name: "finalizer-provider", + preserveClientIp: false, + priority: 1, + providerType: "claude", + providerVendorId: null, + proxyFallbackToDirect: false, + proxyUrl: null, + requestTimeoutNonStreamingMs: 0, + rpd: 0, + rpm: 0, + streamingIdleTimeoutMs: 0, + swapCacheTtlBilling: false, + totalCostResetAt: null, + tpm: 0, + updatedAt: CREATED_AT, + url: "https://provider.test", + websiteUrl: null, + weight: 1, + } satisfies Provider; +} + +async function createSession(provider: Provider | null): Promise { + const request = new Request("https://hub.test/v1/messages", { + body: JSON.stringify({ messages: [], model: "claude-test", stream: false }), + headers: { "content-type": "application/json" }, + method: "POST", + }); + const session = await ProxySession.fromContext(new Context(request)); + session.setProvider(provider); + session.setOriginalModel("claude-test"); + if (provider) { + Object.defineProperty(session, "messageContext", { + value: { createdAt: CREATED_AT, id: 71 }, + writable: true, + }); + Object.defineProperty(session, "getResolvedPricingByBillingSource", { + value: vi.fn(async () => ({ + priceData: { input_cost_per_token: 1, output_cost_per_token: 10 }, + resolvedModelName: "claude-test", + resolvedPricingProviderKey: "anthropic", + source: "official_fallback" as const, + })), + }); + } + return session; +} + +describe("exported response finalizers", () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.addLoserCost.mockResolvedValue(undefined); + mocks.durable.mockResolvedValue(undefined); + mocks.updateCost.mockResolvedValue(undefined); + }); + + it("skips request finalization without provider and message context", async () => { + const session = await createSession(null); + + const usage = await finalizeRequestStats(session, "{}", 200, 15); + + expect(usage).toBeNull(); + expect(mocks.durable).not.toHaveBeenCalled(); + }); + + it("returns parsed usage and durably persists request statistics", async () => { + const session = await createSession(createProvider()); + const responseText = JSON.stringify({ usage: { input_tokens: 2, output_tokens: 3 } }); + + const usage = await finalizeRequestStats(session, responseText, 200, 15); + + expect(usage).toMatchObject({ input_tokens: 2, output_tokens: 3 }); + expect(mocks.durable).toHaveBeenCalledWith( + 71, + expect.objectContaining({ inputTokens: 2, outputTokens: 3, statusCode: 200 }) + ); + }); + + it("skips incomplete hedge drains that contain no usage", async () => { + const provider = createProvider(); + const session = await createSession(provider); + + const billed = await finalizeHedgeLoserBilling({ + allContent: "partial", + attemptNumber: 2, + drainComplete: false, + loserSession: session, + messageRequestCreatedAtMs: 0, + messageRequestId: 71, + provider, + upstreamStatusCode: 200, + }); + + expect(billed).toBeNull(); + expect(mocks.addLoserCost).not.toHaveBeenCalled(); + }); + + it("adds a complete hedge loser's calculated cost to the original request", async () => { + const provider = createProvider(); + const session = await createSession(provider); + + const billed = await finalizeHedgeLoserBilling({ + allContent: JSON.stringify({ usage: { input_tokens: 2, output_tokens: 3 } }), + attemptNumber: 2, + drainComplete: true, + loserSession: session, + messageRequestCreatedAtMs: 0, + messageRequestId: 71, + provider, + upstreamStatusCode: 200, + }); + + expect(billed).toBe("32"); + expect(mocks.addLoserCost).toHaveBeenCalledWith( + 71, + expect.objectContaining({ toString: expect.any(Function) }), + expect.objectContaining({ attemptNumber: 2, costUsd: "32", providerId: 13 }) + ); + }); +}); diff --git a/tests/unit/proxy/response-handler-gemini-stream-passthrough-timeouts.test.ts b/tests/unit/proxy/response-handler-gemini-stream-passthrough-timeouts.test.ts index c786c0e27..ba7bd9b13 100644 --- a/tests/unit/proxy/response-handler-gemini-stream-passthrough-timeouts.test.ts +++ b/tests/unit/proxy/response-handler-gemini-stream-passthrough-timeouts.test.ts @@ -5,13 +5,27 @@ import { ProxyForwarder } from "@/app/v1/_lib/proxy/forwarder"; import { resolveEndpointPolicy } from "@/app/v1/_lib/proxy/endpoint-policy"; import { ProxyResponseHandler } from "@/app/v1/_lib/proxy/response-handler"; import { ProxySession } from "@/app/v1/_lib/proxy/session"; +import { setDeferredStreamingFinalization } from "@/app/v1/_lib/proxy/stream-finalization"; import { AsyncTaskManager } from "@/lib/async-task-manager"; import { SessionManager } from "@/lib/session-manager"; -import { updateMessageRequestDetails } from "@/repository/message"; +import { + updateMessageRequestDetails, + updateMessageRequestDetailsDurably, +} from "@/repository/message"; import type { Provider } from "@/types/provider"; const asyncTasks: Promise[] = []; +async function expectAllFulfilled(tasks: readonly Promise[]): Promise { + const settlements = await Promise.allSettled(tasks); + const rejections = settlements + .filter((settlement): settlement is PromiseRejectedResult => settlement.status === "rejected") + .map((settlement) => settlement.reason); + if (rejections.length > 0) { + throw new AggregateError(rejections, "Unexpected async task rejection"); + } +} + const mocks = vi.hoisted(() => { return { isHttp2Enabled: vi.fn(async () => false), @@ -21,6 +35,17 @@ const mocks = vi.hoisted(() => { beforeEach(() => { mocks.isHttp2Enabled.mockReset(); mocks.isHttp2Enabled.mockResolvedValue(false); + vi.mocked(updateMessageRequestDetailsDurably).mockImplementation( + async (_id, details, options) => { + try { + const result = options?.onCommitted?.(details); + if (result) void Promise.resolve(result).catch(() => undefined); + } catch { + // Test mock mirrors the repository's commit-observer boundary. + } + return true; + } + ); }); vi.mock("@/lib/config", async (importOriginal) => { @@ -39,10 +64,26 @@ vi.mock("@/app/v1/_lib/proxy/response-fixer", () => ({ vi.mock("@/lib/async-task-manager", () => ({ AsyncTaskManager: { - register: vi.fn((_taskId: string, promise: Promise) => { - asyncTasks.push(promise); - return new AbortController(); - }), + register: vi.fn( + ( + _taskId: string, + factory: (signal: AbortSignal) => Promise, + options?: string | { abortController?: AbortController } + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + let promise: Promise; + try { + promise = Promise.resolve(factory(controller.signal)); + } catch (error) { + promise = Promise.reject(error); + } + asyncTasks.push(promise); + return controller; + } + ), touch: () => true, cleanup: () => {}, cancel: () => {}, @@ -59,9 +100,20 @@ vi.mock("@/lib/logger", () => ({ }, })); +vi.mock("@/lib/circuit-breaker", () => ({ + recordFailure: vi.fn(async () => undefined), + recordSuccess: vi.fn(async () => undefined), +})); + +vi.mock("@/lib/endpoint-circuit-breaker", () => ({ + recordEndpointFailure: vi.fn(async () => undefined), + recordEndpointSuccess: vi.fn(async () => undefined), +})); + vi.mock("@/repository/message", () => ({ updateMessageRequestCost: vi.fn(), updateMessageRequestDetails: vi.fn(), + updateMessageRequestDetailsDurably: vi.fn(), updateMessageRequestDuration: vi.fn(), })); @@ -80,6 +132,8 @@ vi.mock("@/lib/session-manager", () => ({ storeSessionResponse: vi.fn(), updateSessionUsage: vi.fn(async () => undefined), clearSessionProvider: vi.fn(), + updateSessionBindingSmart: vi.fn(async () => ({ updated: false, reason: "test" })), + updateSessionProvider: vi.fn(async () => undefined), storeSessionRequestPhaseSnapshot: vi.fn(async () => undefined), storeSessionResponsePhaseSnapshot: vi.fn(async () => undefined), storeSessionUpstreamRequestMeta: vi.fn(async () => undefined), @@ -350,7 +404,7 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { ).toBeNull(); }); - test("Gemini 流式透传返回原始响应前应丢弃未消费的 before-snapshot 响应分支", async () => { + test("Gemini 流式透传应在泵读取前丢弃 before-snapshot 并保留返回正文", async () => { asyncTasks.length = 0; const cancel = vi.fn(async () => undefined); const session = createSession({ @@ -369,10 +423,11 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { body: { cancel }, }; + const expectedBody = 'data: {"provider":"gemini"}\n\n'; const upstreamResponse = new Response( new ReadableStream({ start(controller) { - controller.enqueue(new TextEncoder().encode('data: {"provider":"gemini"}\n\n')); + controller.enqueue(new TextEncoder().encode(expectedBody)); controller.close(); }, }), @@ -388,14 +443,14 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { } ).handleStream(session, upstreamResponse); - expect(returned).toBe(upstreamResponse); + await expect(returned.text()).resolves.toBe(expectedBody); expect(cancel).toHaveBeenCalledOnce(); expect( (session as ProxySession & { detailSnapshotResponseBeforeSource?: unknown }) .detailSnapshotResponseBeforeSource ).toBeNull(); - await Promise.allSettled(asyncTasks); + await expectAllFulfilled(asyncTasks); }); test("Gemini 流式透传禁用 idle timeout 时不应回落到默认 stale cleanup", async () => { @@ -425,7 +480,7 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { ).handleStream(session, upstreamResponse); await returned.text(); - await Promise.allSettled(asyncTasks); + await expectAllFulfilled(asyncTasks); const statsRegisterCall = vi.mocked(AsyncTaskManager.register).mock.calls.find((call) => { const options = call[2] as { taskType?: string } | undefined; @@ -504,7 +559,7 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { } finally { clientAbortController.abort(new Error("test_cleanup")); await close(); - await Promise.allSettled(asyncTasks); + await expectAllFulfilled(asyncTasks); } }); @@ -570,7 +625,7 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { } finally { clientAbortController.abort(new Error("test_cleanup")); await close(); - await Promise.allSettled(asyncTasks); + await expectAllFulfilled(asyncTasks); } }); @@ -635,11 +690,11 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { } finally { clientAbortController.abort(new Error("test_cleanup")); await close(); - await Promise.allSettled(asyncTasks); + await expectAllFulfilled(asyncTasks); } }); - test("客户端中断流式透传后应清理 session provider 绑定,避免下次继续复用旧供应商", async () => { + test("客户端中断流式透传后应继续 drain 并提交尾部 usage", async () => { asyncTasks.length = 0; const { baseUrl, close } = await startSseServer((_req, res) => { res.writeHead(200, { @@ -651,11 +706,13 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { res.write('data: {"x":1}\n\n'); setTimeout(() => { try { - res.write('data: {"x":2}\n\n'); + res.end( + 'data: {"usageMetadata":{"promptTokenCount":3,"candidatesTokenCount":2},"finishReason":"STOP"}\n\n' + ); } catch { // ignore } - }, 1000); + }, 20); }); const clientAbortController = new AbortController(); @@ -697,22 +754,29 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { expect(first.done).toBe(false); clientAbortController.abort(new Error("client_cancelled")); - await Promise.allSettled(asyncTasks); - - expect(vi.mocked(SessionManager.clearSessionProvider)).toHaveBeenCalledWith( - "gemini-abort-session" + await expectAllFulfilled(asyncTasks); + + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 4, + expect.objectContaining({ + statusCode: 200, + inputTokens: 3, + outputTokens: 2, + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); + expect(vi.mocked(SessionManager.clearSessionProvider)).not.toHaveBeenCalled(); } finally { clientAbortController.abort(new Error("test_cleanup")); await close(); - await Promise.allSettled(asyncTasks); + await expectAllFulfilled(asyncTasks); } }); test("Gemini 流式透传超大单 chunk 应保留尾部 usage 且不把截断快照作为完整正文存储", async () => { asyncTasks.length = 0; vi.mocked(SessionManager.storeSessionResponse).mockClear(); - vi.mocked(updateMessageRequestDetails).mockClear(); + vi.mocked(updateMessageRequestDetailsDurably).mockClear(); const clientAbortController = new AbortController(); const provider = createProvider({ @@ -756,16 +820,97 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { ).handleStream(session, upstreamResponse); await returned.text(); - await Promise.allSettled(asyncTasks); + await expectAllFulfilled(asyncTasks); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( 77, expect.objectContaining({ statusCode: 200, inputTokens: 463, outputTokens: 11, - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); expect(SessionManager.storeSessionResponse).not.toHaveBeenCalled(); }); + + test("Gemini 终态副作用挂起时应保持 shutdown ownership 直到真实 I/O 完成", async () => { + vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); + asyncTasks.length = 0; + vi.mocked(updateMessageRequestDetailsDurably).mockClear(); + vi.mocked(SessionManager.clearSessionProvider).mockClear(); + const releaseAgent = vi.fn(); + let releaseSideEffect: () => void = () => {}; + vi.mocked(SessionManager.clearSessionProvider).mockImplementationOnce(() => { + return new Promise((resolve) => { + releaseSideEffect = resolve; + }); + }); + + try { + const session = createSession({ + clientAbortSignal: new AbortController().signal, + messageId: 88, + userId: 1, + }); + const provider = createProvider({ + firstByteTimeoutStreamingMs: 1000, + streamingIdleTimeoutMs: 0, + }); + session.setProvider(provider); + session.setSessionId("gemini-post-terminal-deadline"); + Object.assign(session, { releaseAgent }); + setDeferredStreamingFinalization(session, { + providerId: provider.id, + providerName: provider.name, + providerPriority: provider.priority, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 500, + }); + + const bodyText = + 'data: {"usageMetadata":{"promptTokenCount":463,"candidatesTokenCount":11}}\n\n'; + const response = new Response(bodyText, { + status: 500, + headers: { "content-type": "text/event-stream" }, + }); + const returned = await ( + ProxyResponseHandler as unknown as { + handleStream: (session: ProxySession, response: Response) => Promise; + } + ).handleStream(session, response); + + await returned.text(); + for ( + let attempt = 0; + attempt < 20 && vi.mocked(SessionManager.clearSessionProvider).mock.calls.length === 0; + attempt++ + ) { + await vi.advanceTimersByTimeAsync(1); + } + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); + expect(SessionManager.clearSessionProvider).toHaveBeenCalledTimes(1); + expect(releaseAgent).toHaveBeenCalledTimes(1); + + let tasksSettled = false; + const pendingTasks = expectAllFulfilled(asyncTasks.splice(0, asyncTasks.length)).then(() => { + tasksSettled = true; + }); + await vi.advanceTimersByTimeAsync(120_000); + expect(tasksSettled).toBe(false); + + releaseSideEffect(); + await pendingTasks; + + expect(releaseAgent).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); }); diff --git a/tests/unit/proxy/response-handler-gemini-terminal.test.ts b/tests/unit/proxy/response-handler-gemini-terminal.test.ts new file mode 100644 index 000000000..c6d0706e8 --- /dev/null +++ b/tests/unit/proxy/response-handler-gemini-terminal.test.ts @@ -0,0 +1,249 @@ +import { Context } from "hono"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { ProxyResponseHandler } from "@/app/v1/_lib/proxy/response-handler"; +import { ProxySession, type MessageContext } from "@/app/v1/_lib/proxy/session"; +import type { Provider } from "@/types/provider"; +import type { User } from "@/types/user"; + +type TaskOptions = { readonly abortController?: AbortController }; +type TerminalWriterOptions = { readonly onCommitted?: () => void | Promise }; +const mocks = vi.hoisted(() => ({ + conditional: + vi.fn<(id: number, details: object, options?: TerminalWriterOptions) => Promise>(), + details: vi.fn<(id: number, details: object) => Promise>(), + durable: + vi.fn<(id: number, details: object, options?: TerminalWriterOptions) => Promise>(), + tasks: Array.from>([]), +})); + +vi.mock("@/app/v1/_lib/proxy/response-fixer", () => ({ + ResponseFixer: { process: async (_session: ProxySession, response: Response) => response }, +})); +vi.mock("@/lib/async-task-manager", () => ({ + AsyncTaskManager: { + register: ( + _id: string, + factory: (signal: AbortSignal) => Promise, + options: string | TaskOptions = "unknown" + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + const task = Promise.resolve().then(() => factory(controller.signal)); + mocks.tasks.push(task); + return controller; + }, + touch: vi.fn(() => true), + }, +})); +vi.mock("@/lib/config/system-settings-cache", () => ({ + getCachedSystemSettings: vi.fn(async () => ({ billNonSuccessfulRequests: false })), +})); +vi.mock("@/lib/langfuse/emit-proxy-trace", () => ({ emitProxyLangfuseTrace: vi.fn() })); +vi.mock("@/lib/logger", () => ({ + logger: { + debug: vi.fn(), + error: vi.fn(), + info: vi.fn(), + trace: vi.fn(), + warn: vi.fn(), + }, +})); +vi.mock("@/lib/proxy-status-tracker", () => ({ + ProxyStatusTracker: { getInstance: () => ({ endRequest: vi.fn() }) }, +})); +vi.mock("@/repository/message", () => ({ + addMessageRequestHedgeLoserCost: vi.fn(), + updateMessageRequestCostWithBreakdown: vi.fn(), + updateMessageRequestDetails: mocks.details, + updateMessageRequestDetailsDurably: mocks.durable, + updateMessageRequestDetailsIfUnfinalized: mocks.conditional, + updateMessageRequestDuration: vi.fn(), + updateMessageRequestWinnerCost: vi.fn(), +})); + +type TerminalMessage = Pick & { + readonly user: Pick; +}; +const MESSAGE = { + createdAt: new Date(0), + id: 61, + user: { id: 31 }, +} satisfies TerminalMessage; + +type GeminiProvider = Pick< + Provider, + | "costMultiplier" + | "id" + | "name" + | "providerType" + | "streamingIdleTimeoutMs" + | "swapCacheTtlBilling" +>; + +function createProvider(streamingIdleTimeoutMs = 0): GeminiProvider { + return { + costMultiplier: 1, + id: 9, + name: "gemini-terminal-provider", + providerType: "gemini", + streamingIdleTimeoutMs, + swapCacheTtlBilling: false, + } satisfies GeminiProvider; +} + +async function createSession(options: { + readonly idleMs?: number; + readonly responseController?: AbortController; +}): Promise { + const request = new Request("https://hub.test/v1/messages", { + body: JSON.stringify({ messages: [], stream: true }), + headers: { "content-type": "application/json" }, + method: "POST", + }); + const session = await ProxySession.fromContext(new Context(request)); + Object.defineProperty(session, "provider", { + value: createProvider(options.idleMs), + writable: true, + }); + Object.defineProperty(session, "messageContext", { value: MESSAGE, writable: true }); + session.providerType = "gemini"; + session.setOriginalFormat("gemini"); + if (options.responseController) { + Object.defineProperty(session, "responseController", { value: options.responseController }); + } + return session; +} + +const geminiResponse = (body: BodyInit) => + new Response(body, { headers: { "content-type": "text/event-stream" } }); + +async function settleTasks(): Promise { + while (mocks.tasks.length > 0) { + const results = await Promise.allSettled(mocks.tasks.splice(0, mocks.tasks.length)); + const errors = results.flatMap((result) => + result.status === "rejected" ? [result.reason] : [] + ); + if (errors.length > 0) { + throw new AggregateError(errors, "Gemini terminal background tasks failed"); + } + } +} + +describe("ProxyResponseHandler.dispatch Gemini terminal behavior", () => { + afterEach(() => vi.useRealTimers()); + + beforeEach(() => { + vi.useRealTimers(); + mocks.tasks.length = 0; + vi.clearAllMocks(); + mocks.conditional.mockResolvedValue(true); + mocks.details.mockResolvedValue(undefined); + mocks.durable.mockResolvedValue(true); + }); + + it("drains the Gemini source after the returned body is cancelled", async () => { + const cancelSource = vi.fn(); + let sourceController: ReadableStreamDefaultController | null = null; + const encoder = new TextEncoder(); + const source = new ReadableStream({ + cancel: cancelSource, + start(controller) { + sourceController = controller; + controller.enqueue(encoder.encode('{"chunk":1}\n')); + }, + }); + const session = await createSession({}); + const returned = await ProxyResponseHandler.dispatch(session, geminiResponse(source)); + + await returned.body?.cancel(new Error("client cancelled body")); + sourceController?.enqueue(encoder.encode('{"usageMetadata":{"promptTokenCount":1}}\n')); + sourceController?.close(); + await settleTasks(); + + expect(cancelSource).not.toHaveBeenCalled(); + expect(mocks.durable).toHaveBeenCalledWith( + MESSAGE.id, + expect.objectContaining({ + durationMs: expect.any(Number), + inputTokens: 1, + statusCode: 499, + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + }); + + it("resets the Gemini idle window after every received chunk", async () => { + vi.useFakeTimers(); + const responseController = new AbortController(); + const source = new ReadableStream({ + start(controller) { + const encoder = new TextEncoder(); + controller.enqueue(encoder.encode('{"chunk":1}\n')); + setTimeout(() => controller.enqueue(encoder.encode('{"chunk":2}\n')), 80); + setTimeout(() => { + controller.enqueue(encoder.encode('{"finishReason":"STOP"}\n')); + controller.close(); + }, 160); + }, + }); + const session = await createSession({ idleMs: 100, responseController }); + const returned = await ProxyResponseHandler.dispatch(session, geminiResponse(source)); + const body = returned.text(); + + await vi.advanceTimersByTimeAsync(160); + await expect(body).resolves.toContain('"finishReason":"STOP"'); + await settleTasks(); + + expect(responseController.signal.aborted).toBe(false); + }); + + it("uses conditional persistence when Gemini durable finalization fails", async () => { + mocks.durable.mockRejectedValue(new Error("durable unavailable")); + const session = await createSession({}); + const returned = await ProxyResponseHandler.dispatch( + session, + geminiResponse('{"finishReason":"STOP"}\n') + ); + + await returned.text(); + await settleTasks(); + + expect(mocks.conditional).toHaveBeenCalledWith( + MESSAGE.id, + expect.objectContaining({ durationMs: expect.any(Number), statusCode: 500 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + }); + + it("settles Gemini terminal work when fallback persistence hangs", async () => { + vi.useFakeTimers(); + mocks.durable.mockRejectedValue(new Error("durable unavailable")); + mocks.details.mockImplementation(() => new Promise(() => {})); + mocks.conditional.mockImplementation(() => new Promise(() => {})); + const session = await createSession({}); + const returned = await ProxyResponseHandler.dispatch( + session, + geminiResponse('{"finishReason":"STOP"}\n') + ); + await returned.text(); + const terminalTask = mocks.tasks[0]; + let settled = false; + void terminalTask?.then( + () => { + settled = true; + }, + () => { + settled = true; + } + ); + + await vi.advanceTimersByTimeAsync(5_001); + + expect(mocks.details.mock.calls.length + mocks.conditional.mock.calls.length).toBeGreaterThan( + 0 + ); + expect(settled).toBe(true); + }); +}); diff --git a/tests/unit/proxy/response-handler-hedge-loser-priority.test.ts b/tests/unit/proxy/response-handler-hedge-loser-priority.test.ts index d9f2d354d..fae16b4ca 100644 --- a/tests/unit/proxy/response-handler-hedge-loser-priority.test.ts +++ b/tests/unit/proxy/response-handler-hedge-loser-priority.test.ts @@ -1,3 +1,4 @@ +import { Context } from "hono"; import { beforeEach, describe, expect, it, vi } from "vitest"; const mocks = vi.hoisted(() => ({ @@ -7,6 +8,11 @@ const mocks = vi.hoisted(() => ({ trackCost: vi.fn(async () => {}), trackUserDailyCost: vi.fn(async () => {}), decrementLeaseBudget: vi.fn(async () => {}), + settleLeaseBudgets: vi.fn(async () => ({ + requestId: "test", + status: "settled" as const, + settlements: [], + })), })); vi.mock("@/repository/message", () => ({ @@ -30,7 +36,18 @@ vi.mock("@/lib/logger", () => ({ vi.mock("@/lib/async-task-manager", () => ({ AsyncTaskManager: { - register: () => new AbortController(), + register: ( + _taskId: string, + factory: (signal: AbortSignal) => Promise, + options?: string | { abortController?: AbortController } + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + void Promise.resolve(factory(controller.signal)).catch(() => {}); + return controller; + }, touch: vi.fn(() => true), cleanup: vi.fn(), cancel: vi.fn(), @@ -47,6 +64,7 @@ vi.mock("@/lib/rate-limit", () => ({ trackCost: mocks.trackCost, trackUserDailyCost: mocks.trackUserDailyCost, decrementLeaseBudget: mocks.decrementLeaseBudget, + settleLeaseBudgets: mocks.settleLeaseBudgets, }, })); @@ -59,6 +77,7 @@ vi.mock(import("@/lib/utils/performance-formatter"), async (importOriginal) => { }); import { finalizeHedgeLoserBilling } from "@/app/v1/_lib/proxy/response-handler"; +import { ProxySession } from "@/app/v1/_lib/proxy/session"; import type { Provider } from "@/types/provider"; function createCodexProvider(overrides: Partial = {}): Provider { @@ -123,7 +142,7 @@ function createCodexProvider(overrides: Partial = {}): Provider { }; } -function createLoserSession( +async function createLoserSession( provider: Provider, overrides: { sessionId?: string | null; @@ -131,53 +150,46 @@ function createLoserSession( groupCostMultiplier?: number; } = {} ) { - return { - provider, - request: { - model: "winner-default-model", - message: { - model: "winner-default-model", - service_tier: "default", - }, - }, - sessionId: overrides.sessionId ?? "session-1", - messageContext: { id: 123, createdAt: new Date("2026-06-08T00:00:00.000Z") }, - authState: { - key: { - id: 10, - limit5hResetMode: "rolling", - dailyResetTime: "00:00", - dailyResetMode: "fixed", - }, - user: { - id: 20, - limit5hResetMode: "rolling", - dailyResetTime: "00:00", - dailyResetMode: "fixed", - }, - }, - getEndpoint: () => "/v1/responses", - getOriginalModel: () => "winner-default-model", - getCurrentModel: () => "winner-default-model", - getContext1mApplied: () => overrides.context1mApplied ?? false, - setContext1mApplied: vi.fn(), - getGroupCostMultiplier: () => overrides.groupCostMultiplier ?? 1, - getSpecialSettings: () => [], - addSpecialSetting: vi.fn(), - shouldTrackSessionObservability: () => false, - getCodexPriorityBillingSource: vi.fn(async () => "requested"), - getResolvedPricingByBillingSource: vi.fn(async () => ({ - resolvedModelName: "gpt-5.5", - resolvedPricingProviderKey: "openai", - source: "official_fallback", - priceData: { - input_cost_per_token: 1, - output_cost_per_token: 10, - input_cost_per_token_priority: 2, - output_cost_per_token_priority: 20, - }, - })), + const context = new Context( + new Request("http://localhost/v1/responses", { + method: "POST", + body: JSON.stringify({ model: "winner-default-model", service_tier: "default" }), + headers: { "content-type": "application/json" }, + }) + ); + const session = await ProxySession.fromContext(context); + session.provider = provider; + session.sessionId = overrides.sessionId ?? "session-1"; + session.messageContext = { + id: 123, + createdAt: new Date("2026-06-08T00:00:00.000Z"), + user: { id: 20, limit5hResetMode: "rolling", dailyResetTime: "00:00", dailyResetMode: "fixed" }, + key: { id: 10, limit5hResetMode: "rolling", dailyResetTime: "00:00", dailyResetMode: "fixed" }, + apiKey: "test-api-key", + }; + session.authState = { + success: true, + apiKey: "test-api-key", + key: session.messageContext.key, + user: session.messageContext.user, }; + session.getCodexPriorityBillingSource = vi.fn(async () => "requested"); + session.getResolvedPricingByBillingSource = vi.fn(async () => ({ + resolvedModelName: "gpt-5.5", + resolvedPricingProviderKey: "openai", + source: "official_fallback", + priceData: { + input_cost_per_token: 1, + output_cost_per_token: 10, + input_cost_per_token_priority: 2, + output_cost_per_token_priority: 20, + }, + })); + if (overrides.context1mApplied) session.setContext1mApplied(true); + if (overrides.groupCostMultiplier !== undefined) { + session.setGroupCostMultiplier(overrides.groupCostMultiplier); + } + return session; } describe("finalizeHedgeLoserBilling Codex priority snapshot", () => { @@ -188,11 +200,12 @@ describe("finalizeHedgeLoserBilling Codex priority snapshot", () => { mocks.trackCost.mockClear(); mocks.trackUserDailyCost.mockClear(); mocks.decrementLeaseBudget.mockClear(); + mocks.settleLeaseBudgets.mockClear(); }); it("uses the initial loser's captured requested service tier after winner session sync", async () => { const provider = createCodexProvider(); - const loserSession = createLoserSession(provider); + const loserSession = await createLoserSession(provider); const responseBody = JSON.stringify({ usage: { input_tokens: 100, @@ -202,7 +215,8 @@ describe("finalizeHedgeLoserBilling Codex priority snapshot", () => { const billed = await finalizeHedgeLoserBilling({ messageRequestId: 123, - loserSession: loserSession as any, + messageRequestCreatedAtMs: new Date("2026-06-08T00:00:00.000Z").getTime(), + loserSession, provider, attemptNumber: 1, upstreamStatusCode: 200, @@ -233,7 +247,7 @@ describe("finalizeHedgeLoserBilling Codex priority snapshot", () => { name: "winner-polluted-provider", costMultiplier: 10, }); - const loserSession = createLoserSession(winnerProvider, { + const loserSession = await createLoserSession(winnerProvider, { context1mApplied: true, groupCostMultiplier: 99, }); @@ -246,7 +260,8 @@ describe("finalizeHedgeLoserBilling Codex priority snapshot", () => { const billed = await finalizeHedgeLoserBilling({ messageRequestId: 123, - loserSession: loserSession as any, + messageRequestCreatedAtMs: new Date("2026-06-08T00:00:00.000Z").getTime(), + loserSession, provider: loserProvider, attemptNumber: 1, upstreamStatusCode: 200, @@ -270,17 +285,50 @@ describe("finalizeHedgeLoserBilling Codex priority snapshot", () => { 2400, expect.objectContaining({ userId: 20, - requestId: 123, + userResetTime: "00:00", + userResetMode: "fixed", + requestId: "123:hedge-loser:11:1", }) ); - expect(mocks.trackUserDailyCost).toHaveBeenCalledWith( - 20, - 2400, - "00:00", - "fixed", + expect(mocks.trackUserDailyCost).not.toHaveBeenCalled(); + expect(mocks.settleLeaseBudgets).toHaveBeenCalledWith( expect.objectContaining({ - requestId: 123, + requestId: "123:hedge-loser:11:1", + cost: 2400, }) ); }); + + it("tracks an alternative loser even when its shadow session has no request context", async () => { + const loserProvider = createCodexProvider({ id: 12, name: "shadow-loser" }); + const loserSession = await createLoserSession(loserProvider, { sessionId: null }); + loserSession.sessionId = null; + loserSession.messageContext = null; + + const billed = await finalizeHedgeLoserBilling({ + messageRequestId: 124, + messageRequestCreatedAtMs: new Date("2026-06-08T00:00:01.000Z").getTime(), + loserSession, + provider: loserProvider, + attemptNumber: 2, + upstreamStatusCode: 200, + allContent: JSON.stringify({ usage: { input_tokens: 100, output_tokens: 10 } }), + drainComplete: true, + }); + + expect(billed).toBe("200"); + expect(mocks.trackCost).toHaveBeenCalledWith( + 10, + loserProvider.id, + "", + 200, + expect.objectContaining({ + requestId: "124:hedge-loser:12:2", + createdAtMs: new Date("2026-06-08T00:00:01.000Z").getTime(), + }) + ); + expect(mocks.settleLeaseBudgets).toHaveBeenCalledWith( + expect.objectContaining({ requestId: "124:hedge-loser:12:2", cost: 200 }) + ); + }); }); diff --git a/tests/unit/proxy/response-handler-lease-decrement.test.ts b/tests/unit/proxy/response-handler-lease-decrement.test.ts index 28bc86d65..dd76fdc05 100644 --- a/tests/unit/proxy/response-handler-lease-decrement.test.ts +++ b/tests/unit/proxy/response-handler-lease-decrement.test.ts @@ -1,11 +1,8 @@ /** - * TDD: RED Phase - Tests for lease budget decrement in response-handler.ts + * TDD: Tests for atomic lease budget settlement in response-handler.ts * - * Tests that decrementLeaseBudget is called correctly after trackCostToRedis completes. - * - All windows: 5h, daily, weekly, monthly - * - All entity types: key, user, provider - * - Zero-cost requests should NOT trigger decrement - * - Function runs once per request (no duplicates) + * Tests that settleLeaseBudgets is called once after trackCostToRedis completes. + * The service expands the explicit key/user/provider entities into all twelve windows. */ import { beforeEach, describe, expect, it, vi } from "vitest"; @@ -14,12 +11,34 @@ import type { ModelPriceData } from "@/types/model-price"; // Track async tasks for draining const asyncTasks: Promise[] = []; +const asyncTaskControllers = new Map, AbortController>(); +let asyncTaskAdmissionOpen = true; vi.mock("@/lib/async-task-manager", () => ({ AsyncTaskManager: { - register: (_taskId: string, promise: Promise) => { + register: ( + _taskId: string, + factory: (signal: AbortSignal) => Promise, + options?: string | { abortController?: AbortController } + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + if (!asyncTaskAdmissionOpen) { + controller.abort(); + return controller; + } + + let promise: Promise; + try { + promise = Promise.resolve(factory(controller.signal)); + } catch (error) { + promise = Promise.reject(error); + } asyncTasks.push(promise); - return new AbortController(); + asyncTaskControllers.set(promise, controller); + return controller; }, touch: vi.fn(() => true), cleanup: () => {}, @@ -53,6 +72,8 @@ vi.mock("@/repository/message", () => ({ updateMessageRequestCost: vi.fn(), updateMessageRequestCostWithBreakdown: vi.fn(), updateMessageRequestDetails: vi.fn(), + updateMessageRequestDetailsDurably: vi.fn(), + updateMessageRequestDetailsIfUnfinalized: vi.fn(), updateMessageRequestDuration: vi.fn(), })); @@ -71,6 +92,7 @@ vi.mock("@/lib/rate-limit", () => ({ trackCost: vi.fn(), trackUserDailyCost: vi.fn(), decrementLeaseBudget: vi.fn(), + settleLeaseBudgets: vi.fn(), }, })); @@ -97,6 +119,7 @@ import { SessionTracker } from "@/lib/session-tracker"; import { updateMessageRequestCost, updateMessageRequestDetails, + updateMessageRequestDetailsDurably, updateMessageRequestDuration, } from "@/repository/message"; import { findLatestPriceByModel } from "@/repository/model-price"; @@ -317,13 +340,204 @@ function createStreamResponse(usage: { input_tokens: number; output_tokens: numb } async function drainAsyncTasks(): Promise { - const tasks = asyncTasks.splice(0, asyncTasks.length); - await Promise.all(tasks); + const errors: unknown[] = []; + const maxDrainRounds = 100; + let round = 0; + + while (asyncTasks.length > 0) { + if (round >= maxDrainRounds) { + asyncTaskAdmissionOpen = false; + const overflowTasks = asyncTasks.splice(0, asyncTasks.length); + for (const task of overflowTasks) { + asyncTaskControllers.get(task)?.abort(); + } + const overflowResults = await Promise.allSettled(overflowTasks); + for (let index = 0; index < overflowResults.length; index += 1) { + asyncTaskControllers.delete(overflowTasks[index]); + const result = overflowResults[index]; + if (result.status === "rejected") { + errors.push(result.reason); + } + } + errors.push(new Error(`Async task drain exceeded ${maxDrainRounds} rounds`)); + break; + } + round += 1; + + const tasks = asyncTasks.splice(0, asyncTasks.length); + const results = await Promise.allSettled(tasks); + for (let index = 0; index < results.length; index += 1) { + asyncTaskControllers.delete(tasks[index]); + const result = results[index]; + if (result.status === "rejected") { + errors.push(result.reason); + } + } + } + + if (errors.length > 0) { + throw new AggregateError(errors, "Async task drain failed"); + } } beforeEach(() => { vi.clearAllMocks(); + asyncTaskAdmissionOpen = false; + for (const controller of asyncTaskControllers.values()) { + controller.abort(); + } asyncTasks.splice(0, asyncTasks.length); + asyncTaskControllers.clear(); + asyncTaskAdmissionOpen = true; +}); + +describe("drainAsyncTasks", () => { + it("waits for a tail task registered while draining the primary task", async () => { + let markTailStarted: () => void = () => {}; + let releaseTail: () => void = () => {}; + const tailStarted = new Promise((resolve) => { + markTailStarted = resolve; + }); + const tailCompleted = vi.fn(); + + AsyncTaskManager.register("primary", async () => { + await Promise.resolve(); + AsyncTaskManager.register("tail", async () => { + markTailStarted(); + await new Promise((resolve) => { + releaseTail = resolve; + }); + tailCompleted(); + }); + }); + + const drainPromise = drainAsyncTasks(); + await tailStarted; + + try { + const outcome = await Promise.race([ + drainPromise.then(() => "drained" as const), + new Promise<"pending">((resolve) => { + setTimeout(() => resolve("pending"), 0); + }), + ]); + + expect(outcome).toBe("pending"); + } finally { + releaseTail(); + } + + await drainPromise; + expect(tailCompleted).toHaveBeenCalledTimes(1); + }); + + it("waits for sibling tail work before reporting tail rejections", async () => { + const tailError = new Error("tail task failed"); + let markPendingTailStarted: () => void = () => {}; + let releasePendingTail: () => void = () => {}; + const pendingTailStarted = new Promise((resolve) => { + markPendingTailStarted = resolve; + }); + + AsyncTaskManager.register("primary", async () => { + await Promise.resolve(); + AsyncTaskManager.register("rejecting-tail", async () => { + throw tailError; + }); + void asyncTasks.at(-1)?.catch(() => {}); + AsyncTaskManager.register("pending-tail", async () => { + markPendingTailStarted(); + await new Promise((resolve) => { + releasePendingTail = resolve; + }); + }); + }); + + const drainPromise = drainAsyncTasks(); + await pendingTailStarted; + + try { + const earlyOutcome = await Promise.race([ + drainPromise.then( + () => "resolved" as const, + () => "rejected" as const + ), + new Promise<"pending">((resolve) => { + setTimeout(() => resolve("pending"), 0); + }), + ]); + + expect(earlyOutcome).toBe("pending"); + } finally { + releasePendingTail(); + } + + const rejection = await drainPromise.then( + () => undefined, + (error: unknown) => error + ); + expect(rejection).toBeInstanceOf(AggregateError); + expect((rejection as AggregateError).errors).toEqual([tailError]); + }); + + it("closes admission and observes overflow work when the drain guard trips", async () => { + const overflowError = new Error("overflow task aborted"); + const blockedTailStarted = vi.fn(); + let overflowController: AbortController | undefined; + + const registerGeneration = (generation: number): void => { + const controller = AsyncTaskManager.register(`generation-${generation}`, async (signal) => { + await new Promise((resolve) => { + setTimeout(resolve, 0); + }); + + if (generation <= 100) { + registerGeneration(generation + 1); + return; + } + + await new Promise((_resolve, reject) => { + const rejectOnAbort = () => { + AsyncTaskManager.register("blocked-overflow-tail", async () => { + blockedTailStarted(); + }); + reject(overflowError); + }; + + if (signal.aborted) { + rejectOnAbort(); + return; + } + signal.addEventListener("abort", rejectOnAbort, { once: true }); + }); + }); + + if (generation === 101) { + overflowController = controller; + } + }; + + registerGeneration(1); + const rejection = await drainAsyncTasks().then( + () => undefined, + (error: unknown) => error + ); + + try { + expect(rejection).toBeInstanceOf(AggregateError); + expect((rejection as AggregateError).errors).toContain(overflowError); + expect((rejection as AggregateError).errors).toEqual( + expect.arrayContaining([ + expect.objectContaining({ message: "Async task drain exceeded 100 rounds" }), + ]) + ); + expect(blockedTailStarted).not.toHaveBeenCalled(); + expect(asyncTasks).toHaveLength(0); + } finally { + overflowController?.abort(); + await Promise.allSettled(asyncTasks.splice(0, asyncTasks.length)); + } + }); }); describe("Lease Budget Decrement after trackCostToRedis", () => { @@ -335,7 +549,7 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { vi.mocked(findLatestPriceByModel).mockResolvedValue( makePriceRecord(originalModel, testPriceData) ); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(true); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponsePhaseSnapshot).mockResolvedValue(undefined); @@ -345,10 +559,15 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { success: true, newRemaining: 10, }); + vi.mocked(RateLimitService.settleLeaseBudgets).mockResolvedValue({ + requestId: "test", + status: "settled", + settlements: [], + }); vi.mocked(SessionTracker.refreshSession).mockResolvedValue(undefined); }); - it("should call decrementLeaseBudget for all windows and entity types (non-stream)", async () => { + it("should settle all windows and entity types in one call (non-stream)", async () => { const session = createSession({ originalModel, redirectedModel: originalModel, @@ -363,33 +582,17 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { // Expected cost: (1000 * 0.000003) + (500 * 0.000015) = 0.003 + 0.0075 = 0.0105 const expectedCost = 0.0105; - // Should be called 12 times: - // 4 windows x 3 entity types = 12 calls - // Windows: 5h, daily, weekly, monthly - // Entity types: key(456), user(123), provider(99) - expect(RateLimitService.decrementLeaseBudget).toHaveBeenCalled(); - - const calls = vi.mocked(RateLimitService.decrementLeaseBudget).mock.calls; - expect(calls.length).toBe(12); - - // Verify all windows are covered for each entity type - const windows = ["5h", "daily", "weekly", "monthly"]; - const entities = [ - { id: 456, type: "key" }, - { id: 123, type: "user" }, - { id: 99, type: "provider" }, - ]; - - for (const entity of entities) { - for (const window of windows) { - const matchingCall = calls.find( - (call) => call[0] === entity.id && call[1] === entity.type && call[2] === window - ); - expect(matchingCall).toBeDefined(); - // Cost should be approximately 0.0105 - expect(matchingCall![3]).toBeCloseTo(expectedCost, 4); - } - } + expect(RateLimitService.settleLeaseBudgets).toHaveBeenCalledTimes(1); + expect(RateLimitService.settleLeaseBudgets).toHaveBeenCalledWith({ + requestId: 5001, + cost: expectedCost, + entities: { + key: { id: 456, resetModes: { "5h": undefined, daily: "fixed" } }, + user: { id: 123, resetModes: { "5h": undefined, daily: "fixed" } }, + provider: { id: 99, resetModes: { "5h": undefined, daily: "fixed" } }, + }, + }); + expect(RateLimitService.decrementLeaseBudget).not.toHaveBeenCalled(); }); it("should refresh task activity while reading chunked non-stream response bodies", async () => { @@ -422,17 +625,18 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { }), session.requestSequence ); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( messageId, expect.objectContaining({ statusCode: 200, inputTokens: usage.input_tokens, outputTokens: usage.output_tokens, - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); }); - it("should call decrementLeaseBudget for all windows and entity types (stream)", async () => { + it("should settle all windows and entity types in one call (stream)", async () => { const session = createSession({ originalModel, redirectedModel: originalModel, @@ -445,14 +649,11 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { await clientResponse.text(); await drainAsyncTasks(); - expect(RateLimitService.decrementLeaseBudget).toHaveBeenCalled(); - const calls = vi.mocked(RateLimitService.decrementLeaseBudget).mock.calls; - - // Should have exactly 12 calls (4 windows x 3 entity types) - expect(calls.length).toBe(12); + expect(RateLimitService.settleLeaseBudgets).toHaveBeenCalledTimes(1); + expect(RateLimitService.decrementLeaseBudget).not.toHaveBeenCalled(); }); - it("should NOT call decrementLeaseBudget when cost is zero", async () => { + it("should NOT settle lease budgets when cost is zero", async () => { // Mock price data that results in zero cost const zeroPriceData: ModelPriceData = { input_cost_per_token: 0, @@ -490,7 +691,8 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { await ProxyResponseHandler.dispatch(session, response); await drainAsyncTasks(); - // Zero cost should NOT trigger decrement + // Zero cost should NOT trigger settlement. + expect(RateLimitService.settleLeaseBudgets).not.toHaveBeenCalled(); expect(RateLimitService.decrementLeaseBudget).not.toHaveBeenCalled(); }); @@ -511,18 +713,20 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { expect(RateLimitService.trackCost).not.toHaveBeenCalled(); expect(RateLimitService.trackUserDailyCost).not.toHaveBeenCalled(); + expect(RateLimitService.settleLeaseBudgets).not.toHaveBeenCalled(); expect(RateLimitService.decrementLeaseBudget).not.toHaveBeenCalled(); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( 5999, expect.objectContaining({ statusCode: 200, inputTokens: usage.input_tokens, outputTokens: usage.output_tokens, - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); }); - it("should call decrementLeaseBudget exactly once per request (no duplicates)", async () => { + it("should call settleLeaseBudgets exactly once per request", async () => { const session = createSession({ originalModel, redirectedModel: originalModel, @@ -534,16 +738,10 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { await ProxyResponseHandler.dispatch(session, response); await drainAsyncTasks(); - // Each window/entity combo should be called exactly once - const calls = vi.mocked(RateLimitService.decrementLeaseBudget).mock.calls; - - // Create a unique key for each call to check for duplicates - const callKeys = calls.map((call) => `${call[0]}-${call[1]}-${call[2]}`); - const uniqueKeys = new Set(callKeys); - - // No duplicates: unique keys should equal total calls - expect(uniqueKeys.size).toBe(calls.length); - expect(calls.length).toBe(12); // 4 windows x 3 entities + expect(RateLimitService.settleLeaseBudgets).toHaveBeenCalledTimes(1); + expect(RateLimitService.settleLeaseBudgets).toHaveBeenCalledWith( + expect.objectContaining({ requestId: 5004 }) + ); }); it("should use correct entity IDs from session", async () => { @@ -607,27 +805,20 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { await ProxyResponseHandler.dispatch(session, response); await drainAsyncTasks(); - const calls = vi.mocked(RateLimitService.decrementLeaseBudget).mock.calls; - - // Verify key ID - const keyCalls = calls.filter((c) => c[1] === "key"); - expect(keyCalls.every((c) => c[0] === customKeyId)).toBe(true); - expect(keyCalls.length).toBe(4); - - // Verify user ID - const userCalls = calls.filter((c) => c[1] === "user"); - expect(userCalls.every((c) => c[0] === customUserId)).toBe(true); - expect(userCalls.length).toBe(4); - - // Verify provider ID - const providerCalls = calls.filter((c) => c[1] === "provider"); - expect(providerCalls.every((c) => c[0] === customProviderId)).toBe(true); - expect(providerCalls.length).toBe(4); + expect(RateLimitService.settleLeaseBudgets).toHaveBeenCalledWith( + expect.objectContaining({ + requestId: 5005, + entities: expect.objectContaining({ + key: expect.objectContaining({ id: customKeyId }), + user: expect.objectContaining({ id: customUserId }), + provider: expect.objectContaining({ id: customProviderId }), + }), + }) + ); }); - it("should use fire-and-forget pattern (not block on decrement failures)", async () => { - // Mock decrementLeaseBudget to fail - vi.mocked(RateLimitService.decrementLeaseBudget).mockRejectedValue( + it("should preserve fail-open completion when atomic settlement unexpectedly rejects", async () => { + vi.mocked(RateLimitService.settleLeaseBudgets).mockRejectedValue( new Error("Redis connection failed") ); @@ -640,11 +831,10 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { const response = createNonStreamResponse(usage); - // Should NOT throw even if decrementLeaseBudget fails + // Should NOT throw even if the settlement wrapper fails unexpectedly. await expect(ProxyResponseHandler.dispatch(session, response)).resolves.toBeDefined(); await drainAsyncTasks(); - // Verify decrement was attempted - expect(RateLimitService.decrementLeaseBudget).toHaveBeenCalled(); + expect(RateLimitService.settleLeaseBudgets).toHaveBeenCalledTimes(1); }); }); diff --git a/tests/unit/proxy/response-handler-non200.test.ts b/tests/unit/proxy/response-handler-non200.test.ts index ef62b25a7..19ad99af6 100644 --- a/tests/unit/proxy/response-handler-non200.test.ts +++ b/tests/unit/proxy/response-handler-non200.test.ts @@ -19,9 +19,23 @@ const asyncTasks: Promise[] = []; vi.mock("@/lib/async-task-manager", () => ({ AsyncTaskManager: { - register: (_taskId: string, promise: Promise) => { + register: ( + _taskId: string, + factory: (signal: AbortSignal) => Promise, + options?: string | { abortController?: AbortController } + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + let promise: Promise; + try { + promise = Promise.resolve(factory(controller.signal)); + } catch (error) { + promise = Promise.reject(error); + } asyncTasks.push(promise); - return new AbortController(); + return controller; }, touch: () => true, cleanup: () => {}, diff --git a/tests/unit/proxy/response-handler-nonstream-terminal.test.ts b/tests/unit/proxy/response-handler-nonstream-terminal.test.ts new file mode 100644 index 000000000..f72581739 --- /dev/null +++ b/tests/unit/proxy/response-handler-nonstream-terminal.test.ts @@ -0,0 +1,282 @@ +import { Context } from "hono"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { ProxyResponseHandler } from "@/app/v1/_lib/proxy/response-handler"; +import { ProxySession, type MessageContext } from "@/app/v1/_lib/proxy/session"; +import type { Key } from "@/types/key"; +import type { Provider } from "@/types/provider"; +import type { User } from "@/types/user"; + +type TaskOptions = { readonly abortController?: AbortController; readonly taskType?: string }; +type TerminalWriterOptions = { readonly onCommitted?: () => void | Promise }; + +const mocks = vi.hoisted(() => ({ + conditional: + vi.fn<(id: number, details: object, options?: TerminalWriterOptions) => Promise>(), + durable: + vi.fn<(id: number, details: object, options?: TerminalWriterOptions) => Promise>(), + recordFailure: vi.fn<(providerId: number, error: Error) => Promise>(), + tasks: Array.from>([]), + trackerEnd: vi.fn(), +})); + +vi.mock("@/app/v1/_lib/proxy/response-fixer", () => ({ + ResponseFixer: { process: async (_session: ProxySession, response: Response) => response }, +})); +vi.mock("@/lib/async-task-manager", () => ({ + AsyncTaskManager: { + cancel: vi.fn(), + cleanup: vi.fn(), + register: ( + _id: string, + factory: (signal: AbortSignal) => Promise, + options: string | TaskOptions = "unknown" + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + const task = Promise.resolve().then(() => factory(controller.signal)); + mocks.tasks.push(task); + return controller; + }, + touch: vi.fn(() => true), + }, +})); +vi.mock("@/lib/circuit-breaker", () => ({ + recordFailure: mocks.recordFailure, + recordSuccess: vi.fn(), +})); +vi.mock("@/lib/config/system-settings-cache", () => ({ + getCachedSystemSettings: vi.fn(async () => ({ billNonSuccessfulRequests: false })), +})); +vi.mock("@/lib/langfuse/emit-proxy-trace", () => ({ emitProxyLangfuseTrace: vi.fn() })); +vi.mock("@/lib/logger", () => ({ + logger: { + debug: vi.fn(), + error: vi.fn(), + info: vi.fn(), + trace: vi.fn(), + warn: vi.fn(), + }, +})); +vi.mock("@/lib/proxy-status-tracker", () => ({ + ProxyStatusTracker: { getInstance: () => ({ endRequest: mocks.trackerEnd }) }, +})); +vi.mock("@/repository/message", () => ({ + addMessageRequestHedgeLoserCost: vi.fn(), + updateMessageRequestCostWithBreakdown: vi.fn(), + updateMessageRequestDetails: vi.fn(), + updateMessageRequestDetailsDurably: mocks.durable, + updateMessageRequestDetailsIfUnfinalized: mocks.conditional, + updateMessageRequestDuration: vi.fn(), + updateMessageRequestWinnerCost: vi.fn(), +})); + +const CREATED_AT = new Date(0); +const USER = { + createdAt: CREATED_AT, + dailyQuota: null, + dailyResetMode: "fixed", + dailyResetTime: "00:00", + description: "terminal test user", + id: 11, + isEnabled: true, + limit5hResetMode: "fixed", + name: "terminal-user", + providerGroup: null, + role: "user", + rpm: null, + updatedAt: CREATED_AT, +} satisfies User; +const KEY = { + cacheTtlPreference: null, + canLoginWebUi: false, + createdAt: CREATED_AT, + dailyResetMode: "fixed", + dailyResetTime: "00:00", + id: 12, + isEnabled: true, + key: "sk-terminal", + limit5hResetMode: "fixed", + limit5hUsd: null, + limitConcurrentSessions: 0, + limitDailyUsd: null, + limitMonthlyUsd: null, + limitWeeklyUsd: null, + name: "terminal-key", + providerGroup: null, + updatedAt: CREATED_AT, + userId: USER.id, +} satisfies Key; +const MESSAGE = { + apiKey: KEY.key, + createdAt: CREATED_AT, + id: 41, + key: KEY, + user: USER, +} satisfies MessageContext; + +function createProvider(): Provider { + return { + activeTimeEnd: null, + activeTimeStart: null, + allowedClients: [], + allowedModels: null, + anthropicAdaptiveThinking: null, + anthropicMaxTokensPreference: null, + anthropicThinkingBudgetPreference: null, + blockedClients: [], + cacheTtlPreference: null, + cc: 0, + circuitBreakerFailureThreshold: 5, + circuitBreakerHalfOpenSuccessThreshold: 2, + circuitBreakerOpenDuration: 1_800_000, + codexImageGenerationPreference: null, + codexParallelToolCallsPreference: null, + codexReasoningEffortPreference: null, + codexReasoningSummaryPreference: null, + codexServiceTierPreference: null, + codexTextVerbosityPreference: null, + context1mPreference: null, + costMultiplier: 1, + createdAt: CREATED_AT, + customHeaders: null, + dailyResetMode: "fixed", + dailyResetTime: "00:00", + disableSessionReuse: false, + faviconUrl: null, + firstByteTimeoutStreamingMs: 0, + geminiGoogleSearchPreference: null, + groupPriorities: null, + groupTag: null, + id: 7, + isEnabled: true, + key: "provider-key", + limit5hResetMode: "fixed", + limit5hUsd: null, + limitConcurrentSessions: 0, + limitDailyUsd: null, + limitMonthlyUsd: null, + limitTotalUsd: null, + limitWeeklyUsd: null, + maxRetryAttempts: null, + mcpPassthroughType: "none", + mcpPassthroughUrl: null, + modelRedirects: null, + name: "nonstream-terminal-provider", + preserveClientIp: false, + priority: 1, + providerType: "claude", + providerVendorId: null, + proxyFallbackToDirect: false, + proxyUrl: null, + requestTimeoutNonStreamingMs: 0, + rpd: 0, + rpm: 0, + streamingIdleTimeoutMs: 0, + swapCacheTtlBilling: false, + totalCostResetAt: null, + tpm: 0, + updatedAt: CREATED_AT, + url: "https://provider.test", + websiteUrl: null, + weight: 1, + } satisfies Provider; +} + +async function createSession(releaseAgent: () => void): Promise { + const request = new Request("https://hub.test/v1/messages", { + body: JSON.stringify({ messages: [], stream: false }), + headers: { "content-type": "application/json" }, + method: "POST", + }); + const session = await ProxySession.fromContext(new Context(request)); + session.setProvider(createProvider()); + session.setMessageContext(MESSAGE); + Object.defineProperty(session, "releaseAgent", { value: releaseAgent, writable: true }); + return session; +} + +async function settleTasks(): Promise[]> { + const settlements = Array.from>([]); + while (mocks.tasks.length > 0) { + settlements.push(...(await Promise.allSettled(mocks.tasks.splice(0, mocks.tasks.length)))); + } + return settlements; +} + +describe("ProxyResponseHandler.dispatch nonstream terminal behavior", () => { + beforeEach(() => { + mocks.tasks.length = 0; + vi.clearAllMocks(); + mocks.conditional.mockImplementation(async (_id, _details, options) => { + try { + const result = options?.onCommitted?.(); + if (result) void Promise.resolve(result).catch(() => undefined); + } catch { + // Test mock mirrors the repository's commit-observer boundary. + } + return true; + }); + mocks.durable.mockImplementation(async (_id, _details, options) => { + try { + const result = options?.onCommitted?.(); + if (result) void Promise.resolve(result).catch(() => undefined); + } catch { + // Test mock mirrors the repository's commit-observer boundary. + } + return true; + }); + mocks.recordFailure.mockResolvedValue(undefined); + }); + + it("uses conditional persistence before recording a nonstream provider failure", async () => { + mocks.durable.mockRejectedValueOnce(new Error("primary unavailable")); + const releaseAgent = vi.fn(); + const session = await createSession(releaseAgent); + + const returned = await ProxyResponseHandler.dispatch( + session, + new Response('{"error":{"message":"unavailable"}}', { + status: 503, + headers: { "content-type": "application/json" }, + }) + ); + await returned.text(); + const settlements = await settleTasks(); + + expect(settlements.every(({ status }) => status === "fulfilled")).toBe(true); + expect(mocks.conditional).toHaveBeenCalledWith( + 41, + expect.objectContaining({ durationMs: expect.any(Number), statusCode: 503 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect(mocks.recordFailure).toHaveBeenCalledWith(7, expect.any(Error)); + expect(mocks.conditional.mock.invocationCallOrder[0]).toBeLessThan( + mocks.recordFailure.mock.invocationCallOrder[0] + ); + expect(releaseAgent).toHaveBeenCalledOnce(); + }); + + it("rejects the managed task without mutating the circuit when both writes fail", async () => { + mocks.durable.mockRejectedValueOnce(new Error("primary unavailable")); + mocks.conditional.mockRejectedValueOnce(new Error("fallback unavailable")); + const releaseAgent = vi.fn(); + const session = await createSession(releaseAgent); + + const returned = await ProxyResponseHandler.dispatch( + session, + new Response('{"error":{"message":"unavailable"}}', { + status: 503, + headers: { "content-type": "application/json" }, + }) + ); + await returned.text(); + const settlements = await settleTasks(); + + expect(settlements.some(({ status }) => status === "rejected")).toBe(true); + expect(mocks.recordFailure).not.toHaveBeenCalled(); + expect(mocks.trackerEnd).toHaveBeenCalledWith(USER.id, MESSAGE.id); + expect(releaseAgent).toHaveBeenCalledOnce(); + }); +}); diff --git a/tests/unit/proxy/response-handler-public-dispatch.test.ts b/tests/unit/proxy/response-handler-public-dispatch.test.ts new file mode 100644 index 000000000..3fbe94539 --- /dev/null +++ b/tests/unit/proxy/response-handler-public-dispatch.test.ts @@ -0,0 +1,213 @@ +import { Context } from "hono"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { ProxyResponseHandler } from "@/app/v1/_lib/proxy/response-handler"; +import { ProxySession } from "@/app/v1/_lib/proxy/session"; +import type { Provider } from "@/types/provider"; + +type TaskOptions = { + readonly abortController?: AbortController; + readonly staleTimeoutMs?: number; + readonly taskType?: string; +}; + +const state = vi.hoisted(() => ({ + tasks: Array.from>([]), + taskTypes: Array.from([]), +})); + +vi.mock("@/app/v1/_lib/proxy/response-fixer", () => ({ + ResponseFixer: { process: async (_session: ProxySession, response: Response) => response }, +})); +vi.mock("@/lib/async-task-manager", () => ({ + AsyncTaskManager: { + cancel: vi.fn(), + cleanup: vi.fn(), + register: ( + _taskId: string, + factory: (signal: AbortSignal) => Promise, + options: string | TaskOptions = "unknown" + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + const task = Promise.resolve().then(() => factory(controller.signal)); + state.tasks.push(task); + state.taskTypes.push(typeof options === "object" ? (options.taskType ?? "unknown") : options); + return controller; + }, + touch: vi.fn(() => true), + }, +})); +vi.mock("@/lib/langfuse/emit-proxy-trace", () => ({ emitProxyLangfuseTrace: vi.fn() })); +vi.mock("@/lib/logger", () => ({ + logger: { + debug: vi.fn(), + error: vi.fn(), + info: vi.fn(), + trace: vi.fn(), + warn: vi.fn(), + }, +})); +vi.mock("@/repository/message", () => ({ + addMessageRequestHedgeLoserCost: vi.fn(), + updateMessageRequestCostWithBreakdown: vi.fn(), + updateMessageRequestDetails: vi.fn(), + updateMessageRequestDetailsDurably: vi.fn(), + updateMessageRequestDetailsIfUnfinalized: vi.fn(), + updateMessageRequestDuration: vi.fn(), + updateMessageRequestWinnerCost: vi.fn(), +})); + +function createProvider(): Provider { + const now = new Date(0); + return { + activeTimeEnd: null, + activeTimeStart: null, + allowedClients: [], + allowedModels: null, + anthropicAdaptiveThinking: null, + anthropicMaxTokensPreference: null, + anthropicThinkingBudgetPreference: null, + blockedClients: [], + cacheTtlPreference: null, + cc: 0, + circuitBreakerFailureThreshold: 5, + circuitBreakerHalfOpenSuccessThreshold: 2, + circuitBreakerOpenDuration: 1_800_000, + codexImageGenerationPreference: null, + codexParallelToolCallsPreference: null, + codexReasoningEffortPreference: null, + codexReasoningSummaryPreference: null, + codexServiceTierPreference: null, + codexTextVerbosityPreference: null, + context1mPreference: null, + costMultiplier: 1, + createdAt: now, + customHeaders: null, + dailyResetMode: "fixed", + dailyResetTime: "00:00", + disableSessionReuse: false, + faviconUrl: null, + firstByteTimeoutStreamingMs: 0, + geminiGoogleSearchPreference: null, + groupPriorities: null, + groupTag: null, + id: 1, + isEnabled: true, + key: "test-key", + limit5hResetMode: "fixed", + limit5hUsd: null, + limitConcurrentSessions: 0, + limitDailyUsd: null, + limitMonthlyUsd: null, + limitTotalUsd: null, + limitWeeklyUsd: null, + maxRetryAttempts: null, + mcpPassthroughType: "none", + mcpPassthroughUrl: null, + modelRedirects: null, + name: "public-dispatch-provider", + preserveClientIp: false, + priority: 1, + providerType: "claude", + providerVendorId: null, + proxyFallbackToDirect: false, + proxyUrl: null, + requestTimeoutNonStreamingMs: 0, + rpd: 0, + rpm: 0, + streamingIdleTimeoutMs: 0, + swapCacheTtlBilling: false, + totalCostResetAt: null, + tpm: 0, + updatedAt: now, + url: "https://provider.test", + websiteUrl: null, + weight: 1, + } satisfies Provider; +} + +async function createSession(stream: boolean, provider: Provider | null): Promise { + const request = new Request("https://hub.test/v1/messages", { + body: JSON.stringify({ messages: [], stream }), + headers: { "content-type": "application/json" }, + method: "POST", + }); + const session = await ProxySession.fromContext(new Context(request)); + session.setProvider(provider); + return session; +} + +async function settleTasks(): Promise { + while (state.tasks.length > 0) { + const tasks = state.tasks.splice(0, state.tasks.length); + await Promise.all(tasks); + } +} + +describe("ProxyResponseHandler.dispatch public routing", () => { + beforeEach(() => { + state.tasks.length = 0; + state.taskTypes.length = 0; + vi.clearAllMocks(); + }); + + it("returns a nonstream response and releases the transport when no provider exists", async () => { + const releaseAgent = vi.fn(); + const session = await createSession(false, null); + Object.defineProperty(session, "releaseAgent", { + configurable: true, + value: releaseAgent, + writable: true, + }); + const upstream = new Response('{"ok":true}', { + headers: { "content-type": "application/json" }, + }); + + const returned = await ProxyResponseHandler.dispatch(session, upstream); + + await expect(returned.text()).resolves.toBe('{"ok":true}'); + expect(releaseAgent).toHaveBeenCalledOnce(); + expect(state.tasks).toEqual([]); + }); + + it("routes a provider nonstream response through the managed terminal task", async () => { + const releaseAgent = vi.fn(); + const session = await createSession(false, createProvider()); + Object.defineProperty(session, "releaseAgent", { + configurable: true, + value: releaseAgent, + writable: true, + }); + const upstream = new Response('{"result":"accepted"}', { + headers: { "content-type": "application/json" }, + }); + + const returned = await ProxyResponseHandler.dispatch(session, upstream); + await expect(returned.text()).resolves.toBe('{"result":"accepted"}'); + await settleTasks(); + + expect(state.taskTypes).toContain("non-stream-processing"); + expect(releaseAgent).toHaveBeenCalledOnce(); + }); + + it("routes SSE through the stream boundary and releases an incomplete session", async () => { + const releaseAgent = vi.fn(); + const session = await createSession(true, createProvider()); + Object.defineProperty(session, "releaseAgent", { + configurable: true, + value: releaseAgent, + writable: true, + }); + const upstream = new Response('event: message_stop\ndata: {"type":"message_stop"}\n\n', { + headers: { "content-type": "text/event-stream" }, + }); + + const returned = await ProxyResponseHandler.dispatch(session, upstream); + + await expect(returned.text()).resolves.toContain("message_stop"); + expect(releaseAgent).toHaveBeenCalledOnce(); + expect(state.tasks).toEqual([]); + }); +}); diff --git a/tests/unit/proxy/response-handler-stream-terminal.test.ts b/tests/unit/proxy/response-handler-stream-terminal.test.ts new file mode 100644 index 000000000..efe5ddd8e --- /dev/null +++ b/tests/unit/proxy/response-handler-stream-terminal.test.ts @@ -0,0 +1,280 @@ +import { Context } from "hono"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { ProxyResponseHandler } from "@/app/v1/_lib/proxy/response-handler"; +import { ProxySession, type MessageContext } from "@/app/v1/_lib/proxy/session"; +import type { Key } from "@/types/key"; +import type { Provider } from "@/types/provider"; +import type { User } from "@/types/user"; + +type TaskOptions = { readonly abortController?: AbortController }; + +const mocks = vi.hoisted(() => ({ + durable: + vi.fn< + (id: number, details: object, options?: { onCommitted?: () => void }) => Promise + >(), + tasks: Array.from>([]), + trackerEnd: vi.fn(), +})); + +vi.mock("@/app/v1/_lib/proxy/response-fixer", () => ({ + ResponseFixer: { process: async (_session: ProxySession, response: Response) => response }, +})); +vi.mock("@/lib/async-task-manager", () => ({ + AsyncTaskManager: { + register: ( + _id: string, + factory: (signal: AbortSignal) => Promise, + options: string | TaskOptions = "unknown" + ) => { + const controller = + typeof options === "object" && options.abortController + ? options.abortController + : new AbortController(); + const task = Promise.resolve().then(() => factory(controller.signal)); + mocks.tasks.push(task); + return controller; + }, + touch: vi.fn(() => true), + }, +})); +vi.mock("@/lib/config/system-settings-cache", () => ({ + getCachedSystemSettings: vi.fn(async () => ({ billNonSuccessfulRequests: false })), +})); +vi.mock("@/lib/langfuse/emit-proxy-trace", () => ({ emitProxyLangfuseTrace: vi.fn() })); +vi.mock("@/lib/logger", () => ({ + logger: { + debug: vi.fn(), + error: vi.fn(), + info: vi.fn(), + trace: vi.fn(), + warn: vi.fn(), + }, +})); +vi.mock("@/lib/proxy-status-tracker", () => ({ + ProxyStatusTracker: { getInstance: () => ({ endRequest: mocks.trackerEnd }) }, +})); +vi.mock("@/repository/message", () => ({ + addMessageRequestHedgeLoserCost: vi.fn(), + updateMessageRequestCostWithBreakdown: vi.fn(), + updateMessageRequestDetails: vi.fn(), + updateMessageRequestDetailsDurably: mocks.durable, + updateMessageRequestDetailsIfUnfinalized: vi.fn(), + updateMessageRequestDuration: vi.fn(), + updateMessageRequestWinnerCost: vi.fn(), +})); + +const CREATED_AT = new Date(0); +const USER = { + createdAt: CREATED_AT, + dailyQuota: null, + dailyResetMode: "fixed", + dailyResetTime: "00:00", + description: "stream test user", + id: 21, + isEnabled: true, + limit5hResetMode: "fixed", + name: "stream-user", + providerGroup: null, + role: "user", + rpm: null, + updatedAt: CREATED_AT, +} satisfies User; +const KEY = { + cacheTtlPreference: null, + canLoginWebUi: false, + createdAt: CREATED_AT, + dailyResetMode: "fixed", + dailyResetTime: "00:00", + id: 22, + isEnabled: true, + key: "sk-stream", + limit5hResetMode: "fixed", + limit5hUsd: null, + limitConcurrentSessions: 0, + limitDailyUsd: null, + limitMonthlyUsd: null, + limitWeeklyUsd: null, + name: "stream-key", + providerGroup: null, + updatedAt: CREATED_AT, + userId: USER.id, +} satisfies Key; +const MESSAGE = { + apiKey: KEY.key, + createdAt: CREATED_AT, + id: 51, + key: KEY, + user: USER, +} satisfies MessageContext; + +function createProvider(): Provider { + return { + activeTimeEnd: null, + activeTimeStart: null, + allowedClients: [], + allowedModels: null, + anthropicAdaptiveThinking: null, + anthropicMaxTokensPreference: null, + anthropicThinkingBudgetPreference: null, + blockedClients: [], + cacheTtlPreference: null, + cc: 0, + circuitBreakerFailureThreshold: 5, + circuitBreakerHalfOpenSuccessThreshold: 2, + circuitBreakerOpenDuration: 1_800_000, + codexImageGenerationPreference: null, + codexParallelToolCallsPreference: null, + codexReasoningEffortPreference: null, + codexReasoningSummaryPreference: null, + codexServiceTierPreference: null, + codexTextVerbosityPreference: null, + context1mPreference: null, + costMultiplier: 1, + createdAt: CREATED_AT, + customHeaders: null, + dailyResetMode: "fixed", + dailyResetTime: "00:00", + disableSessionReuse: false, + faviconUrl: null, + firstByteTimeoutStreamingMs: 0, + geminiGoogleSearchPreference: null, + groupPriorities: null, + groupTag: null, + id: 8, + isEnabled: true, + key: "provider-key", + limit5hResetMode: "fixed", + limit5hUsd: null, + limitConcurrentSessions: 0, + limitDailyUsd: null, + limitMonthlyUsd: null, + limitTotalUsd: null, + limitWeeklyUsd: null, + maxRetryAttempts: null, + mcpPassthroughType: "none", + mcpPassthroughUrl: null, + modelRedirects: null, + name: "stream-terminal-provider", + preserveClientIp: false, + priority: 1, + providerType: "claude", + providerVendorId: null, + proxyFallbackToDirect: false, + proxyUrl: null, + requestTimeoutNonStreamingMs: 0, + rpd: 0, + rpm: 0, + streamingIdleTimeoutMs: 0, + swapCacheTtlBilling: false, + totalCostResetAt: null, + tpm: 0, + updatedAt: CREATED_AT, + url: "https://provider.test", + websiteUrl: null, + weight: 1, + } satisfies Provider; +} + +async function createSession(options: { + readonly responseController?: AbortController; +}): Promise<{ readonly releaseAgent: ReturnType; readonly session: ProxySession }> { + const request = new Request("https://hub.test/v1/messages", { + body: JSON.stringify({ messages: [], stream: true }), + headers: { "content-type": "application/json" }, + method: "POST", + }); + const session = await ProxySession.fromContext(new Context(request)); + const releaseAgent = vi.fn(); + session.setProvider(createProvider()); + session.setMessageContext(MESSAGE); + Object.defineProperty(session, "releaseAgent", { value: releaseAgent, writable: true }); + if (options.responseController) { + Object.defineProperty(session, "responseController", { value: options.responseController }); + } + return { releaseAgent, session }; +} + +async function settleTasks(): Promise { + while (mocks.tasks.length > 0) { + await Promise.all(mocks.tasks.splice(0, mocks.tasks.length)); + } +} + +function sseResponse(body: BodyInit): Response { + return new Response(body, { headers: { "content-type": "text/event-stream" } }); +} + +describe("ProxyResponseHandler.dispatch stream terminal behavior", () => { + beforeEach(() => { + mocks.tasks.length = 0; + vi.clearAllMocks(); + mocks.durable.mockResolvedValue(true); + }); + + it("persists a naturally completed stream and releases its transport", async () => { + const { releaseAgent, session } = await createSession({}); + const returned = await ProxyResponseHandler.dispatch( + session, + sseResponse('event: message_stop\ndata: {"type":"message_stop"}\n\n') + ); + + await returned.text(); + await settleTasks(); + + expect(mocks.durable).toHaveBeenCalledWith( + 51, + expect.objectContaining({ statusCode: 200 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect(mocks.trackerEnd).toHaveBeenCalledWith(USER.id, MESSAGE.id); + expect(releaseAgent).toHaveBeenCalledOnce(); + }); + + it("persists a partial client-aborted stream as 499", async () => { + let abortSource = () => {}; + const source = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode('data: {"partial":true}\n\n')); + abortSource = () => controller.error(new DOMException("client aborted", "AbortError")); + }, + }); + const { releaseAgent, session } = await createSession({}); + const returned = await ProxyResponseHandler.dispatch(session, sseResponse(source)); + const reader = returned.body?.getReader(); + expect(reader).toBeDefined(); + await reader?.read(); + + await reader?.cancel(new Error("client disconnected")); + abortSource(); + await settleTasks(); + + expect(mocks.durable).toHaveBeenCalledWith( + 51, + expect.objectContaining({ statusCode: 499 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect(releaseAgent).toHaveBeenCalledOnce(); + }); + + it("persists a response-controller timeout as 502 and cancels the source", async () => { + const cancelSource = vi.fn(); + const responseController = new AbortController(); + const source = new ReadableStream({ cancel: cancelSource }); + const { releaseAgent, session } = await createSession({ responseController }); + const returned = await ProxyResponseHandler.dispatch(session, sseResponse(source)); + const bodyRead = returned.text(); + + responseController.abort(new Error("response deadline exceeded")); + await expect(bodyRead).rejects.toThrow("response deadline exceeded"); + await settleTasks(); + + expect(mocks.durable).toHaveBeenCalledWith( + 51, + expect.objectContaining({ statusCode: 502 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect(cancelSource).toHaveBeenCalledOnce(); + expect(releaseAgent).toHaveBeenCalledOnce(); + }); +}); diff --git a/tests/unit/proxy/session.test.ts b/tests/unit/proxy/session.test.ts index 306e5ba3e..5c77a0a29 100644 --- a/tests/unit/proxy/session.test.ts +++ b/tests/unit/proxy/session.test.ts @@ -117,19 +117,19 @@ function createSession({ } describe("ProxySession endpoint policy", () => { - it.each([ - V1_ENDPOINT_PATHS.MESSAGES_COUNT_TOKENS, - "/V1/RESPONSES/COMPACT/", - ])("应在创建时解析 raw passthrough policy: %s", (pathname) => { - const session = createSession({ - redirectedModel: null, - requestUrl: new URL(`http://localhost${pathname}`), - }); - - const policy = session.getEndpointPolicy(); - expect(isRawPassthroughEndpointPolicy(policy)).toBe(true); - expect(policy.trackConcurrentRequests).toBe(false); - }); + it.each([V1_ENDPOINT_PATHS.MESSAGES_COUNT_TOKENS, "/V1/RESPONSES/COMPACT/"])( + "应在创建时解析 raw passthrough policy: %s", + (pathname) => { + const session = createSession({ + redirectedModel: null, + requestUrl: new URL(`http://localhost${pathname}`), + }); + + const policy = session.getEndpointPolicy(); + expect(isRawPassthroughEndpointPolicy(policy)).toBe(true); + expect(policy.trackConcurrentRequests).toBe(false); + } + ); it("应在请求路径后续变更后保持创建时 policy 不变", () => { const session = createSession({ diff --git a/tests/unit/proxy/terminal-outcome-contract.test.ts b/tests/unit/proxy/terminal-outcome-contract.test.ts new file mode 100644 index 000000000..8a9d1f859 --- /dev/null +++ b/tests/unit/proxy/terminal-outcome-contract.test.ts @@ -0,0 +1,258 @@ +import type { ProxySession } from "@/app/v1/_lib/proxy/session"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +type SqlQuery = { + toQuery: (config: { + escapeName: (name: string) => string; + escapeParam: (index: number) => string; + escapeString: (value: string) => string; + paramStartIndex: { value: number }; + }) => { sql: string; params: unknown[] }; +}; + +function createDeferred() { + let resolve!: (value: T | PromiseLike) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + return { promise, resolve, reject }; +} + +function toSqlText(query: SqlQuery) { + return query.toQuery({ + escapeName: (name) => `"${name}"`, + escapeParam: (index) => `$${index}`, + escapeString: (value) => `'${value}'`, + paramStartIndex: { value: 1 }, + }); +} + +async function flushMicrotasks(): Promise { + for (let index = 0; index < 12; index++) { + await Promise.resolve(); + } +} + +describe("terminal outcome contract", () => { + afterEach(() => { + vi.doUnmock("@/app/v1/_lib/proxy/errors"); + vi.doUnmock("@/drizzle/db"); + vi.doUnmock("@/lib/config/env.schema"); + vi.doUnmock("@/lib/config/system-settings-cache"); + vi.doUnmock("@/lib/langfuse/emit-proxy-trace"); + vi.doUnmock("@/lib/logger"); + vi.doUnmock("@/lib/proxy-status-tracker"); + vi.doUnmock("@/lib/redis"); + }); + + it("waits for one committed top-level error outcome before returning and rolling up", async () => { + vi.resetModules(); + + const messageRequestId = 92_001; + const releaseCommit = createDeferred(); + const executedSql: Array<{ sql: string; params: unknown[] }> = []; + const rollupPipelines: Array> = []; + const execute = vi.fn(async (query: SqlQuery) => { + executedSql.push(toSqlText(query)); + await releaseCommit.promise; + return [{ id: messageRequestId }]; + }); + const createdAt = new Date("2026-07-15T00:00:00.000Z"); + const insert = vi.fn(() => ({ + values: vi.fn(() => ({ + returning: vi.fn(async () => [ + { + id: messageRequestId, + providerId: 7, + userId: 42, + key: "sk-terminal-outcome", + model: "gpt-4.1", + originalModel: "gpt-4.1", + durationMs: null, + costUsd: "0.125", + costMultiplier: "1", + sessionId: "terminal-outcome-session", + requestSequence: 1, + userAgent: "vitest", + clientIp: "127.0.0.1", + endpoint: "/v1/responses", + messagesCount: 1, + cacheTtlApplied: null, + cacheCreationInputTokens: null, + cacheCreation5mInputTokens: null, + cacheCreation1hInputTokens: null, + cacheReadInputTokens: null, + specialSettings: null, + createdAt, + updatedAt: createdAt, + deletedAt: null, + }, + ]), + })), + })); + + vi.doMock("@/drizzle/db", () => ({ + db: { insert, select: vi.fn(), update: vi.fn() }, + getMessageWriterDb: vi.fn(() => ({ execute, update: vi.fn() })), + })); + vi.doMock("@/lib/config/env.schema", () => ({ + getEnvConfig: () => ({ + MESSAGE_REQUEST_WRITE_MODE: "async", + MESSAGE_REQUEST_ASYNC_FLUSH_INTERVAL_MS: 60_000, + MESSAGE_REQUEST_ASYNC_BATCH_SIZE: 1_000, + MESSAGE_REQUEST_ASYNC_MAX_PENDING: 1_000, + }), + })); + vi.doMock("@/lib/config/system-settings-cache", () => ({ + getCachedSystemSettings: vi.fn(async () => ({ + passThroughUpstreamErrorMessage: false, + verboseProviderError: false, + })), + })); + vi.doMock("@/lib/langfuse/emit-proxy-trace", () => ({ emitProxyLangfuseTrace: vi.fn() })); + vi.doMock("@/lib/proxy-status-tracker", () => ({ + ProxyStatusTracker: { getInstance: () => ({ endRequest: vi.fn() }) }, + })); + vi.doMock("@/lib/logger", () => ({ + logger: { + trace: vi.fn(), + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + fatal: vi.fn(), + }, + })); + + const configSnapshot = JSON.stringify({ + configVersion: "cfg-terminal-outcome", + generatedAt: "2026-07-14T23:59:00.000Z", + siteTitle: "Status", + siteDescription: "Status", + timeZone: "UTC", + defaultIntervalMinutes: 5, + defaultRangeHours: 24, + groups: [ + { + sourceGroupId: 42, + sourceGroupName: "openai", + slug: "openai", + displayName: "OpenAI", + sortOrder: 1, + description: null, + models: [ + { + publicModelKey: "gpt-4.1", + label: "GPT-4.1", + vendorIconKey: "openai", + requestTypeBadge: "openaiCompatible", + }, + ], + }, + ], + }); + const redis = { + status: "ready", + hincrbyfloat: vi.fn(), + get: vi.fn(async (key: string) => { + if (key === "public-status:v2:config-version:current") return "cfg-terminal-outcome"; + if (key === "public-status:v2:config-internal:cfg-terminal-outcome") { + return configSnapshot; + } + return null; + }), + pipeline: vi.fn(() => { + const operations: Array<{ command: string; args: unknown[] }> = []; + return { + hincrbyfloat: (...args: unknown[]) => operations.push({ command: "hincrbyfloat", args }), + set: (...args: unknown[]) => operations.push({ command: "set", args }), + expire: (...args: unknown[]) => operations.push({ command: "expire", args }), + exec: async () => { + rollupPipelines.push(operations); + return operations.map(() => [null, 1] as [null, number]); + }, + }; + }), + }; + vi.doMock("@/lib/redis", () => ({ getRedisClient: vi.fn(() => redis) })); + vi.doMock("@/app/v1/_lib/proxy/errors", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, getErrorOverrideAsync: vi.fn(async () => undefined) }; + }); + + const { ProxyErrorHandler } = await import("@/app/v1/_lib/proxy/error-handler"); + const { createMessageRequest } = await import("@/repository/message"); + const { flushMessageRequestWriteBuffer, stopMessageRequestWriteBuffer } = await import( + "@/repository/message-write-buffer" + ); + + await createMessageRequest({ + provider_id: 7, + user_id: 42, + key: "sk-terminal-outcome", + model: "gpt-4.1", + original_model: "gpt-4.1", + cost_usd: 0.125, + }); + const session = { + sessionId: "terminal-outcome-session", + messageContext: { + id: messageRequestId, + user: { id: 42, name: "test-user" }, + key: { id: 2, name: "test-key" }, + }, + startTime: Date.now() - 250, + requestUrl: new URL("https://gateway.test/v1/responses"), + request: { message: { model: "gpt-4.1" }, model: "gpt-4.1", log: "{}" }, + provider: { id: 7, name: "provider-a", providerType: "openai", swapCacheTtlBilling: false }, + getProviderChain: () => [ + { + id: 7, + name: "provider-a", + groupTag: "openai", + reason: "retry_failed", + statusCode: 500, + }, + ], + getCurrentModel: () => "gpt-4.1", + getContext1mApplied: () => false, + getGroupCostMultiplier: () => 1, + getSpecialSettings: () => null, + } as ProxySession; + + const handlePromise = ProxyErrorHandler.handle(session, new Error("top-level failure")); + await flushMicrotasks(); + const flushPromise = flushMessageRequestWriteBuffer(); + await flushMicrotasks(); + + let settled = false; + void handlePromise.then( + () => { + settled = true; + }, + () => { + settled = true; + } + ); + await new Promise((resolve) => setImmediate(resolve)); + expect(settled).toBe(false); + expect(execute).toHaveBeenCalledTimes(1); + expect(rollupPipelines).toEqual([]); + expect(executedSql[0]?.sql).toContain("duration_ms"); + expect(executedSql[0]?.sql).toContain("status_code"); + expect(executedSql[0]?.sql).toContain("error_message"); + expect(executedSql[0]?.sql).toMatch(/"?status_code"? IS NULL/); + expect(executedSql[0]?.sql).toContain("RETURNING id"); + + releaseCommit.resolve(); + await flushPromise; + const response = await handlePromise; + + expect(response.status).toBe(500); + expect(execute).toHaveBeenCalledTimes(1); + await vi.waitFor(() => expect(rollupPipelines).toHaveLength(1)); + await stopMessageRequestWriteBuffer(); + }); +}); diff --git a/tests/unit/repository/message-aggregate-multiple-session-stats.test.ts b/tests/unit/repository/message-aggregate-multiple-session-stats.test.ts new file mode 100644 index 000000000..0410509af --- /dev/null +++ b/tests/unit/repository/message-aggregate-multiple-session-stats.test.ts @@ -0,0 +1,190 @@ +import { beforeEach, describe, expect, test, vi } from "vitest"; +import { providers, usageLedger } from "@/drizzle/schema"; +import { aggregateMultipleSessionStats } from "@/repository/message"; +import { createDrizzleQuery, sqlText } from "./message-query-test-support"; + +const boundary = vi.hoisted(() => { + const writerDb = { execute: vi.fn<(query: unknown) => Promise>() }; + return { + select: vi.fn<(selection?: unknown) => unknown>(), + selectDistinct: vi.fn<(selection?: unknown) => unknown>(), + execute: vi.fn<(query: unknown) => Promise>(), + ledgerOnly: vi.fn<() => Promise>(), + getWriterDb: vi.fn(() => writerDb), + }; +}); + +vi.mock("@/drizzle/db", () => ({ + db: { + select: boundary.select, + selectDistinct: boundary.selectDistinct, + execute: boundary.execute, + }, + getMessageWriterDb: boundary.getWriterDb, +})); +vi.mock("@/lib/config/env.schema", () => ({ + getEnvConfig: () => ({ MESSAGE_REQUEST_WRITE_MODE: "sync" }), + isDevelopment: () => false, +})); +vi.mock("@/lib/ledger-fallback", () => ({ isLedgerOnlyMode: boundary.ledgerOnly })); + +type StatsRow = { + readonly sessionId: string; + readonly requestCount: number; + readonly totalCostUsd: string; + readonly totalInputTokens: number; + readonly totalOutputTokens: number; + readonly totalCacheCreationTokens: number; + readonly totalCacheReadTokens: number; + readonly totalDurationMs: number; + readonly firstRequestAt: Date; + readonly lastRequestAt: Date; +}; + +const firstRequestAt = new Date("2026-05-02T08:00:00.000Z"); +const lastRequestAt = new Date("2026-05-02T09:00:00.000Z"); + +function statsRow(sessionId: string, ordinal: number): StatsRow { + return { + sessionId, + requestCount: ordinal, + totalCostUsd: `${ordinal}.000000000000`, + totalInputTokens: ordinal * 10, + totalOutputTokens: ordinal * 5, + totalCacheCreationTokens: ordinal * 3, + totalCacheReadTokens: ordinal * 7, + totalDurationMs: ordinal * 100, + firstRequestAt, + lastRequestAt, + }; +} + +describe("message repository aggregateMultipleSessionStats", () => { + beforeEach(() => { + boundary.select.mockReset(); + boundary.selectDistinct.mockReset(); + boundary.execute.mockReset(); + boundary.ledgerOnly.mockReset(); + }); + + test("returns an empty batch without querying the database", async () => { + const result = await aggregateMultipleSessionStats([]); + + expect(result).toEqual([]); + expect(boundary.select).not.toHaveBeenCalled(); + expect(boundary.selectDistinct).not.toHaveBeenCalled(); + expect(boundary.execute).not.toHaveBeenCalled(); + }); + + test("groups related rows and returns sessions in requested order", async () => { + const stats = createDrizzleQuery([ + statsRow("session-a", 1), + statsRow("session-b", 2), + statsRow("session-without-owner", 3), + ]); + const providerList = createDrizzleQuery([ + { sessionId: "session-b", providerId: 21, providerName: "Provider B1" }, + { sessionId: "session-a", providerId: 11, providerName: null }, + { sessionId: "session-b", providerId: 22, providerName: "Provider B2" }, + { sessionId: null, providerId: 99, providerName: "Ignored" }, + ]); + const modelList = createDrizzleQuery([ + { sessionId: "session-a", model: "model-a" }, + { sessionId: "session-b", model: "model-b1" }, + { sessionId: "session-b", model: "model-b2" }, + ]); + const cacheTtlList = createDrizzleQuery([ + { sessionId: "session-b", cacheTtl: "5m" }, + { sessionId: "session-b", cacheTtl: "1h" }, + { sessionId: "session-a", cacheTtl: null }, + ]); + + boundary.select.mockReturnValueOnce(stats); + boundary.selectDistinct + .mockReturnValueOnce(providerList) + .mockReturnValueOnce(modelList) + .mockReturnValueOnce(cacheTtlList); + boundary.execute.mockResolvedValueOnce([ + { + session_id: "session-a", + user_name: "Alice", + user_id: 1, + key_name: "Key A", + key_id: 101, + user_agent: null, + api_type: "claude", + }, + { + session_id: "session-b", + user_name: "Bob", + user_id: 2, + key_name: "Key B", + key_id: 202, + user_agent: "codex-cli/1.0", + api_type: "codex", + }, + ]); + + const result = await aggregateMultipleSessionStats([ + "session-b", + "session-without-owner", + "session-a", + ]); + + expect(result).toEqual([ + { + sessionId: "session-b", + requestCount: 2, + totalCostUsd: "2.000000000000", + totalInputTokens: 20, + totalOutputTokens: 10, + totalCacheCreationTokens: 6, + totalCacheReadTokens: 14, + totalDurationMs: 200, + firstRequestAt, + lastRequestAt, + providers: [ + { id: 21, name: "Provider B1" }, + { id: 22, name: "Provider B2" }, + ], + models: ["model-b1", "model-b2"], + userName: "Bob", + userId: 2, + keyName: "Key B", + keyId: 202, + userAgent: "codex-cli/1.0", + apiType: "codex", + cacheTtlApplied: "mixed", + }, + { + sessionId: "session-a", + requestCount: 1, + totalCostUsd: "1.000000000000", + totalInputTokens: 10, + totalOutputTokens: 5, + totalCacheCreationTokens: 3, + totalCacheReadTokens: 7, + totalDurationMs: 100, + firstRequestAt, + lastRequestAt, + providers: [{ id: 11, name: "Provider #11" }], + models: ["model-a"], + userName: "Alice", + userId: 1, + keyName: "Key A", + keyId: 101, + userAgent: null, + apiType: "claude", + cacheTtlApplied: null, + }, + ]); + expect(stats.trace.from).toEqual([usageLedger]); + expect(sqlText(stats.trace.where)).toContain("session-without-owner"); + expect(sqlText(stats.trace.groupBy)).toContain("session_id"); + expect(providerList.trace.leftJoins.map(({ source }) => source)).toEqual([providers]); + expect(modelList.trace.from).toEqual([usageLedger]); + expect(cacheTtlList.trace.from).toEqual([usageLedger]); + expect(sqlText(boundary.execute.mock.calls.at(0)?.at(0))).toContain("unnest"); + expect(sqlText(boundary.execute.mock.calls.at(0)?.at(0))).toContain("order by created_at"); + }); +}); diff --git a/tests/unit/repository/message-aggregate-session-stats.test.ts b/tests/unit/repository/message-aggregate-session-stats.test.ts new file mode 100644 index 000000000..6e7c1275c --- /dev/null +++ b/tests/unit/repository/message-aggregate-session-stats.test.ts @@ -0,0 +1,181 @@ +import { beforeEach, describe, expect, test, vi } from "vitest"; +import { keys as keysTable, messageRequest, providers, usageLedger, users } from "@/drizzle/schema"; +import { aggregateSessionStats } from "@/repository/message"; +import { createDrizzleQuery, sqlText } from "./message-query-test-support"; + +const boundary = vi.hoisted(() => { + const writerDb = { execute: vi.fn<(query: unknown) => Promise>() }; + return { + select: vi.fn<(selection?: unknown) => unknown>(), + selectDistinct: vi.fn<(selection?: unknown) => unknown>(), + execute: vi.fn<(query: unknown) => Promise>(), + ledgerOnly: vi.fn<() => Promise>(), + getWriterDb: vi.fn(() => writerDb), + }; +}); + +vi.mock("@/drizzle/db", () => ({ + db: { + select: boundary.select, + selectDistinct: boundary.selectDistinct, + execute: boundary.execute, + }, + getMessageWriterDb: boundary.getWriterDb, +})); +vi.mock("@/lib/config/env.schema", () => ({ + getEnvConfig: () => ({ MESSAGE_REQUEST_WRITE_MODE: "sync" }), + isDevelopment: () => false, +})); +vi.mock("@/lib/ledger-fallback", () => ({ isLedgerOnlyMode: boundary.ledgerOnly })); + +type StatsRow = { + readonly requestCount: number; + readonly totalCostUsd: string; + readonly totalInputTokens: number; + readonly totalOutputTokens: number; + readonly totalCacheCreationTokens: number; + readonly totalCacheReadTokens: number; + readonly totalDurationMs: number; + readonly firstRequestAt: Date; + readonly lastRequestAt: Date; +}; + +type UserInfoRow = { + readonly userName: string; + readonly userId: number; + readonly keyName: string; + readonly keyId: number; + readonly userAgent: string | null; + readonly apiType: string | null; +}; + +const firstRequestAt = new Date("2026-05-01T10:00:00.000Z"); +const lastRequestAt = new Date("2026-05-01T10:05:00.000Z"); +const statsRow = { + requestCount: 2, + totalCostUsd: "1.250000000000", + totalInputTokens: 120, + totalOutputTokens: 45, + totalCacheCreationTokens: 30, + totalCacheReadTokens: 70, + totalDurationMs: 900, + firstRequestAt, + lastRequestAt, +} satisfies StatsRow; +const userInfoRow = { + userName: "Ada", + userId: 17, + keyName: "analytics-key", + keyId: 23, + userAgent: "claude-cli/1.0", + apiType: "claude", +} satisfies UserInfoRow; + +function queuePopulatedAggregate( + cacheTtls: readonly (string | null)[], + userRows: readonly UserInfoRow[] = [userInfoRow] +) { + const stats = createDrizzleQuery([statsRow]); + const providerList = createDrizzleQuery([ + { providerId: 11, providerName: null }, + { providerId: 12, providerName: "Provider Twelve" }, + ]); + const modelList = createDrizzleQuery([{ model: "model-a" }, { model: "model-b" }]); + const cacheTtlList = createDrizzleQuery(cacheTtls.map((cacheTtl) => ({ cacheTtl }))); + const userInfo = createDrizzleQuery(userRows); + + boundary.select.mockReturnValueOnce(stats).mockReturnValueOnce(userInfo); + boundary.selectDistinct + .mockReturnValueOnce(providerList) + .mockReturnValueOnce(modelList) + .mockReturnValueOnce(cacheTtlList); + + return { stats, providerList, modelList, cacheTtlList, userInfo }; +} + +describe("message repository aggregateSessionStats", () => { + beforeEach(() => { + boundary.select.mockReset(); + boundary.selectDistinct.mockReset(); + boundary.execute.mockReset(); + boundary.ledgerOnly.mockReset(); + }); + + test("returns null when the session has no billable ledger rows", async () => { + const stats = createDrizzleQuery([]); + boundary.select.mockReturnValueOnce(stats); + + const result = await aggregateSessionStats("session-empty"); + + expect(result).toBeNull(); + expect(stats.trace.from).toEqual([usageLedger]); + expect(sqlText(stats.trace.where)).toContain("session-empty"); + expect(sqlText(stats.trace.where)).toContain("blocked_by"); + expect(boundary.selectDistinct).not.toHaveBeenCalled(); + }); + + test("returns populated statistics and preserves a single cache TTL", async () => { + const queries = queuePopulatedAggregate(["1h"]); + + const result = await aggregateSessionStats("session-populated"); + + expect(result).toEqual({ + sessionId: "session-populated", + requestCount: 2, + totalCostUsd: "1.250000000000", + totalInputTokens: 120, + totalOutputTokens: 45, + totalCacheCreationTokens: 30, + totalCacheReadTokens: 70, + totalDurationMs: 900, + firstRequestAt, + lastRequestAt, + providers: [ + { id: 11, name: "Provider #11" }, + { id: 12, name: "Provider Twelve" }, + ], + models: ["model-a", "model-b"], + userName: "Ada", + userId: 17, + keyName: "analytics-key", + keyId: 23, + userAgent: "claude-cli/1.0", + apiType: "claude", + cacheTtlApplied: "1h", + }); + expect(queries.providerList.trace.from).toEqual([usageLedger]); + expect(queries.providerList.trace.leftJoins.map(({ source }) => source)).toEqual([providers]); + expect(queries.modelList.trace.from).toEqual([usageLedger]); + expect(queries.cacheTtlList.trace.from).toEqual([usageLedger]); + expect(queries.userInfo.trace.from).toEqual([messageRequest]); + expect(queries.userInfo.trace.innerJoins.map(({ source }) => source)).toEqual([ + users, + keysTable, + ]); + expect(queries.userInfo.trace.limit).toEqual([1]); + }); + + test("returns a null cache TTL when the aggregate row contains only null TTL values", async () => { + queuePopulatedAggregate([null]); + + const result = await aggregateSessionStats("session-null-ttl"); + + expect(result).toMatchObject({ cacheTtlApplied: null }); + }); + + test("returns mixed when the session contains multiple cache TTL values", async () => { + queuePopulatedAggregate(["5m", "1h"]); + + const result = await aggregateSessionStats("session-mixed-ttl"); + + expect(result).toMatchObject({ cacheTtlApplied: "mixed" }); + }); + + test("returns null when billable stats exist without a corresponding session owner", async () => { + queuePopulatedAggregate(["5m"], []); + + const result = await aggregateSessionStats("session-without-owner"); + + expect(result).toBeNull(); + }); +}); diff --git a/tests/unit/repository/message-public-readback.test.ts b/tests/unit/repository/message-public-readback.test.ts new file mode 100644 index 000000000..8b2e49ea9 --- /dev/null +++ b/tests/unit/repository/message-public-readback.test.ts @@ -0,0 +1,206 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; + +function createLimitSelect(responses: readonly (readonly unknown[])[]) { + let callIndex = 0; + const events: string[] = []; + const select = vi.fn((_selection: unknown) => { + const rows = responses[callIndex] ?? []; + callIndex += 1; + const limit = vi.fn(async (_value: number) => { + events.push("limit"); + return rows; + }); + const orderBy = vi.fn((..._ordering: unknown[]) => { + events.push("orderBy"); + return { limit }; + }); + const where = vi.fn((_condition: unknown) => { + events.push("where"); + return { limit, orderBy }; + }); + const from = vi.fn((_table: unknown) => { + events.push("from"); + return { where }; + }); + return { from }; + }); + return { events, select }; +} + +function installReadBoundaries( + responses: readonly (readonly unknown[])[], + messageTableHasData = true +) { + const { events, select } = createLimitSelect(responses); + const execute = vi.fn(async (_query: unknown) => [{ has_data: messageTableHasData }]); + vi.doMock("@/drizzle/db", () => ({ + db: { select, execute, update: vi.fn() }, + getMessageWriterDb: vi.fn(() => ({ update: vi.fn(), execute: vi.fn() })), + })); + vi.doMock("@/lib/config/env.schema", () => ({ + getEnvConfig: vi.fn(() => ({ MESSAGE_REQUEST_WRITE_MODE: "sync" as const })), + isDevelopment: vi.fn(() => false), + })); + vi.doMock("@/lib/redis", () => ({ getRedisClient: vi.fn(() => null) })); + return { events, execute, select }; +} + +const CREATED_AT = new Date("2026-07-15T10:00:00.000Z"); +const LATEST_ROW = { + id: 1_001, + providerId: 7, + userId: 9, + key: "public-key", + durationMs: 800, + costUsd: "0.250000000000000", + createdAt: CREATED_AT, + updatedAt: CREATED_AT, + deletedAt: null, +}; +const MESSAGE_ROW = { + ...LATEST_ROW, + model: "gpt-4.1", + originalModel: "gpt-4.1-mini", + ttfbMs: 120, + costMultiplier: "1.5", + sessionId: "public-session", + userAgent: "vitest", + clientIp: "127.0.0.1", + endpoint: "/v1/responses", + messagesCount: 2, + statusCode: 200, + inputTokens: 40, + outputTokens: 12, + cacheCreationInputTokens: 0, + cacheReadInputTokens: 2, + cacheCreation5mInputTokens: 0, + cacheCreation1hInputTokens: 0, + cacheTtlApplied: null, + errorMessage: null, + providerChain: null, + blockedBy: null, + blockedReason: null, + context1mApplied: true, + swapCacheTtlApplied: false, + specialSettings: null, +}; +const LEDGER_ROW = { + requestId: 1_002, + finalProviderId: 17, + userId: 19, + key: "ledger-key", + model: "ledger-model", + originalModel: "requested-model", + endpoint: "/v1/messages", + statusCode: 201, + costUsd: "0.750000000000000", + costMultiplier: "2", + inputTokens: 70, + outputTokens: 20, + cacheCreationInputTokens: 3, + cacheReadInputTokens: 4, + cacheCreation5mInputTokens: 1, + cacheCreation1hInputTokens: 2, + cacheTtlApplied: "1h", + context1mApplied: false, + swapCacheTtlApplied: true, + durationMs: 1_500, + ttfbMs: 250, + sessionId: "ledger-session", + createdAt: CREATED_AT, +}; + +describe("message public readback", () => { + afterEach(() => { + vi.doUnmock("@/drizzle/db"); + vi.doUnmock("@/lib/config/env.schema"); + vi.doUnmock("@/lib/redis"); + }); + + it("returns the latest non-deleted request for a key in descending time order", async () => { + vi.resetModules(); + const boundary = installReadBoundaries([[LATEST_ROW]]); + const { findLatestMessageRequestByKey } = await import("@/repository/message"); + + const result = await findLatestMessageRequestByKey("public-key"); + + expect(result).toMatchObject({ id: 1_001, key: "public-key", costUsd: "0.250000000000000" }); + expect(boundary.events).toEqual(["from", "where", "orderBy", "limit"]); + }); + + it("returns a complete request by its public id without consulting the ledger", async () => { + vi.resetModules(); + const boundary = installReadBoundaries([[MESSAGE_ROW]]); + const { findMessageRequestById } = await import("@/repository/message"); + + const result = await findMessageRequestById(1_001); + + expect(result).toMatchObject({ + id: 1_001, + model: "gpt-4.1", + costMultiplier: 1.5, + context1mApplied: true, + }); + expect(boundary.execute).not.toHaveBeenCalled(); + }); + + it("falls back to the billing ledger when the message table is empty", async () => { + vi.resetModules(); + const boundary = installReadBoundaries([[], [LEDGER_ROW]], false); + const { findMessageRequestById } = await import("@/repository/message"); + + const result = await findMessageRequestById(1_002); + + expect(result).toMatchObject({ + id: 1_002, + providerId: 17, + model: "ledger-model", + costMultiplier: 2, + sessionId: "ledger-session", + userAgent: null, + }); + expect(boundary.select).toHaveBeenCalledTimes(2); + expect(boundary.execute).toHaveBeenCalledTimes(1); + }); + + it("returns null when neither direct data nor ledger-only mode applies", async () => { + vi.resetModules(); + const boundary = installReadBoundaries([[]], true); + const { findMessageRequestById } = await import("@/repository/message"); + + const result = await findMessageRequestById(1_099); + + expect(result).toBeNull(); + expect(boundary.select).toHaveBeenCalledTimes(1); + }); + + it("returns the public audit projection for a session sequence", async () => { + vi.resetModules(); + const auditRow = { + statusCode: 403, + blockedBy: "sensitive_words", + blockedReason: "policy", + cacheTtlApplied: "5m", + context1mApplied: false, + swapCacheTtlApplied: true, + specialSettings: [ + { + type: "guard_intercept", + scope: "guard", + hit: true, + guard: "sensitive_words", + action: "block_request", + statusCode: 403, + reason: "policy", + }, + ], + }; + const boundary = installReadBoundaries([[auditRow]]); + const { findMessageRequestAuditBySessionIdAndSequence } = await import("@/repository/message"); + + const result = await findMessageRequestAuditBySessionIdAndSequence("audit-session", 4); + + expect(result).toEqual(auditRow); + expect(boundary.events).toEqual(["from", "where", "limit"]); + }); +}); diff --git a/tests/unit/repository/message-public-status-rollup.test.ts b/tests/unit/repository/message-public-status-rollup.test.ts index 2046ec89e..864a1af78 100644 --- a/tests/unit/repository/message-public-status-rollup.test.ts +++ b/tests/unit/repository/message-public-status-rollup.test.ts @@ -1,13 +1,25 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { + DurableMessageRequestUpdateOptions, + MessageRequestUpdatePatch, +} from "@/repository/message-write-buffer"; const mockDbInsertValues = vi.hoisted(() => vi.fn()); const mockDbInsertReturning = vi.hoisted(() => vi.fn()); const mockDbUpdateSet = vi.hoisted(() => vi.fn()); const mockDbUpdateWhere = vi.hoisted(() => vi.fn()); +const mockDbUpdateReturning = vi.hoisted(() => vi.fn()); +const mockWriterDbUpdate = vi.hoisted(() => vi.fn()); +const mockWriterDbUpdateSet = vi.hoisted(() => vi.fn()); +const mockWriterDbUpdateWhere = vi.hoisted(() => vi.fn()); +const mockWriterDbUpdateReturning = vi.hoisted(() => vi.fn()); +const mockGetMessageWriterDb = vi.hoisted(() => vi.fn()); const mockDbSelectLimit = vi.hoisted(() => vi.fn()); const mockQueuePublicStatusRollupWrite = vi.hoisted(() => vi.fn()); const mockGetConfiguredPublicStatusGroupsForRollupResolution = vi.hoisted(() => vi.fn()); const mockGetEnvConfig = vi.hoisted(() => vi.fn()); +const mockEnqueueMessageRequestUpdate = vi.hoisted(() => vi.fn()); +const mockEnqueueMessageRequestUpdateDurably = vi.hoisted(() => vi.fn()); vi.mock("@/drizzle/schema", () => ({ keys: {}, @@ -34,6 +46,7 @@ vi.mock("@/drizzle/schema", () => ({ cacheCreation1hInputTokens: "cacheCreation1hInputTokens", cacheReadInputTokens: "cacheReadInputTokens", specialSettings: "specialSettings", + statusCode: "statusCode", createdAt: "createdAt", updatedAt: "updatedAt", deletedAt: "deletedAt", @@ -62,6 +75,7 @@ vi.mock("@/drizzle/db", () => ({ })), })), }, + getMessageWriterDb: mockGetMessageWriterDb, })); vi.mock("@/lib/config/env.schema", () => ({ @@ -81,21 +95,58 @@ vi.mock("@/lib/public-status/rollup-store", () => ({ })); vi.mock("@/repository/message-write-buffer", () => ({ - enqueueMessageRequestUpdate: vi.fn(), + enqueueMessageRequestUpdate: mockEnqueueMessageRequestUpdate, + enqueueMessageRequestUpdateDurably: mockEnqueueMessageRequestUpdateDurably, })); function flushMicrotasks(): Promise { return new Promise((resolve) => setTimeout(resolve, 0)); } +function createDeferred() { + let resolve!: (value: T | PromiseLike) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + return { promise, resolve, reject }; +} + +function sqlToString(value: unknown): string { + const seen = new Set(); + const visit = (node: unknown): string => { + if (node == null || seen.has(node)) return ""; + if (typeof node === "string") return node; + if (typeof node !== "object") return String(node); + seen.add(node); + if (Array.isArray(node)) return node.map(visit).join(" "); + + const record = node as Record; + if (typeof record.name === "string") return record.name; + if (Array.isArray(record.value)) return record.value.map(visit).join(" "); + if (record.value != null) return visit(record.value); + if (record.queryChunks != null) return visit(record.queryChunks); + return ""; + }; + return visit(value); +} + describe("repository/message public status rollup hook", () => { beforeEach(() => { vi.clearAllMocks(); vi.resetModules(); mockGetEnvConfig.mockReturnValue({ MESSAGE_REQUEST_WRITE_MODE: "sync" }); + mockEnqueueMessageRequestUpdateDurably.mockResolvedValue(undefined); mockDbInsertValues.mockReturnValue({ returning: mockDbInsertReturning }); mockDbUpdateSet.mockReturnValue({ where: mockDbUpdateWhere }); - mockDbUpdateWhere.mockResolvedValue(undefined); + mockDbUpdateWhere.mockReturnValue({ returning: mockDbUpdateReturning }); + mockDbUpdateReturning.mockResolvedValue([]); + mockWriterDbUpdate.mockReturnValue({ set: mockWriterDbUpdateSet }); + mockWriterDbUpdateSet.mockReturnValue({ where: mockWriterDbUpdateWhere }); + mockWriterDbUpdateWhere.mockReturnValue({ returning: mockWriterDbUpdateReturning }); + mockWriterDbUpdateReturning.mockResolvedValue([]); + mockGetMessageWriterDb.mockReturnValue({ update: mockWriterDbUpdate }); mockDbSelectLimit.mockResolvedValue([]); mockGetConfiguredPublicStatusGroupsForRollupResolution.mockResolvedValue({ retryable: false, @@ -126,6 +177,225 @@ describe("repository/message public status rollup hook", () => { }); }); + it("writes timeout fallback details only while the request is unfinalized", async () => { + mockDbUpdateReturning.mockResolvedValueOnce([{ id: 606 }]); + const { updateMessageRequestDetailsIfUnfinalized } = await import("@/repository/message"); + + await updateMessageRequestDetailsIfUnfinalized(606, { + statusCode: 500, + errorMessage: "Error: stream_finalization_timeout", + }); + + expect(mockGetMessageWriterDb).not.toHaveBeenCalled(); + expect(mockDbUpdateWhere).toHaveBeenCalledTimes(1); + expect(mockDbUpdateReturning).toHaveBeenCalledWith({ id: "id" }); + const whereSql = sqlToString(mockDbUpdateWhere.mock.calls[0]?.[0]).toLowerCase(); + expect(whereSql).toContain("statuscode"); + expect(whereSql).toContain("is null"); + }); + + it("keeps timeout fallback conditional in async write mode", async () => { + mockGetEnvConfig.mockReturnValue({ MESSAGE_REQUEST_WRITE_MODE: "async" }); + mockWriterDbUpdateReturning.mockResolvedValueOnce([{ id: 608 }]); + const { updateMessageRequestDetailsIfUnfinalized } = await import("@/repository/message"); + + await updateMessageRequestDetailsIfUnfinalized(608, { + statusCode: 500, + errorMessage: "Error: stream_finalization_timeout", + providerChain: [{ id: 1, name: "provider-a" }], + }); + + expect(mockEnqueueMessageRequestUpdate).not.toHaveBeenCalled(); + expect(mockGetMessageWriterDb).toHaveBeenCalledTimes(1); + expect(mockWriterDbUpdateWhere).toHaveBeenCalledTimes(1); + expect(mockWriterDbUpdateReturning).toHaveBeenCalledWith({ id: "id" }); + expect(mockDbUpdateSet).not.toHaveBeenCalled(); + const whereSql = sqlToString(mockWriterDbUpdateWhere.mock.calls[0]?.[0]).toLowerCase(); + expect(whereSql).toContain("statuscode"); + expect(whereSql).toContain("is null"); + }); + + it("does not queue terminal rollup when the timeout fallback loses the terminal CAS", async () => { + mockDbUpdateReturning.mockResolvedValueOnce([]); + const { updateMessageRequestDetailsIfUnfinalized } = await import("@/repository/message"); + + await updateMessageRequestDetailsIfUnfinalized(607, { + statusCode: 500, + errorMessage: "Error: stream_finalization_timeout", + providerChain: [{ id: 1, name: "provider-a" }], + }); + await flushMicrotasks(); + + expect(mockGetMessageWriterDb).not.toHaveBeenCalled(); + expect(mockQueuePublicStatusRollupWrite).not.toHaveBeenCalled(); + }); + + it("async durable details queue public-status rollup only after the batch commit ack", async () => { + mockGetEnvConfig.mockReturnValue({ MESSAGE_REQUEST_WRITE_MODE: "async" }); + const durableAck = createDeferred(); + mockEnqueueMessageRequestUpdateDurably.mockImplementationOnce( + async (_id, details, options: DurableMessageRequestUpdateOptions | undefined) => { + await durableAck.promise; + await options?.onCommitted?.(details); + return true; + } + ); + mockDbInsertReturning.mockResolvedValueOnce([ + { + id: 808, + providerId: 1, + userId: 2, + key: "sk-durable", + model: "gpt-4.1", + originalModel: "gpt-4.1", + durationMs: 100, + costUsd: null, + costMultiplier: null, + sessionId: "session-durable", + requestSequence: 1, + userAgent: null, + clientIp: null, + endpoint: "/v1/messages", + messagesCount: 1, + cacheTtlApplied: null, + cacheCreationInputTokens: null, + cacheCreation5mInputTokens: null, + cacheCreation1hInputTokens: null, + cacheReadInputTokens: null, + specialSettings: null, + createdAt: new Date("2026-04-21T10:02:00.000Z"), + updatedAt: new Date("2026-04-21T10:02:00.000Z"), + deletedAt: null, + }, + ]); + + const { createMessageRequest, updateMessageRequestDetailsDurably } = await import( + "@/repository/message" + ); + await createMessageRequest({ + provider_id: 1, + user_id: 2, + key: "sk-durable", + model: "gpt-4.1", + original_model: "gpt-4.1", + }); + + const updatePromise = updateMessageRequestDetailsDurably(808, { + statusCode: 200, + outputTokens: 10, + providerChain: [{ id: 1, name: "provider-a", groupTag: "openai" }], + model: "gpt-4.1", + }); + await flushMicrotasks(); + + expect(mockQueuePublicStatusRollupWrite).not.toHaveBeenCalled(); + durableAck.resolve(); + await updatePromise; + await flushMicrotasks(); + + expect(mockQueuePublicStatusRollupWrite).toHaveBeenCalledTimes(1); + }); + + it("publishes the public-status rollup when a timed-out durable waiter commits later", async () => { + mockGetEnvConfig.mockReturnValue({ MESSAGE_REQUEST_WRITE_MODE: "async" }); + let onCommitted: DurableMessageRequestUpdateOptions["onCommitted"]; + mockEnqueueMessageRequestUpdateDurably.mockImplementationOnce( + (_id, _details, options: DurableMessageRequestUpdateOptions | undefined) => { + onCommitted = options?.onCommitted; + return Promise.reject(new Error("durable acknowledgement timed out")); + } + ); + mockDbInsertReturning.mockResolvedValueOnce([ + { + id: 810, + providerId: 1, + userId: 2, + key: "sk-late-commit", + model: "gpt-4.1", + originalModel: "gpt-4.1", + durationMs: null, + costUsd: null, + costMultiplier: null, + sessionId: "session-late-commit", + requestSequence: 1, + userAgent: null, + clientIp: null, + endpoint: "/v1/messages", + messagesCount: 1, + cacheTtlApplied: null, + cacheCreationInputTokens: null, + cacheCreation5mInputTokens: null, + cacheCreation1hInputTokens: null, + cacheReadInputTokens: null, + specialSettings: null, + createdAt: new Date("2026-04-21T10:06:00.000Z"), + updatedAt: new Date("2026-04-21T10:06:00.000Z"), + deletedAt: null, + }, + ]); + + const { createMessageRequest, updateMessageRequestDetailsDurably } = await import( + "@/repository/message" + ); + await createMessageRequest({ + provider_id: 1, + user_id: 2, + key: "sk-late-commit", + model: "gpt-4.1", + original_model: "gpt-4.1", + }); + + const terminalPatch = { + durationMs: 1_500, + statusCode: 200, + outputTokens: 10, + providerChain: [{ id: 1, name: "provider-a", groupTag: "openai" }], + model: "gpt-4.1", + } satisfies Readonly; + await expect(updateMessageRequestDetailsDurably(810, terminalPatch)).rejects.toThrow( + "durable acknowledgement timed out" + ); + await flushMicrotasks(); + + expect(onCommitted).toBeTypeOf("function"); + expect(mockQueuePublicStatusRollupWrite).not.toHaveBeenCalled(); + + await onCommitted?.(terminalPatch); + await flushMicrotasks(); + + expect(mockQueuePublicStatusRollupWrite).toHaveBeenCalledTimes(1); + expect(mockQueuePublicStatusRollupWrite).toHaveBeenCalledWith( + expect.objectContaining({ + event: expect.objectContaining({ durationMs: 1_500 }), + }) + ); + }); + + it("sync durable details remain direct and queue rollup after the DB write", async () => { + mockGetEnvConfig.mockReturnValue({ MESSAGE_REQUEST_WRITE_MODE: "sync" }); + mockDbUpdateReturning.mockResolvedValueOnce([{ id: 809 }]); + mockDbSelectLimit.mockResolvedValueOnce([ + { + createdAt: new Date("2026-04-21T10:04:00.000Z"), + model: "gpt-4.1", + originalModel: "gpt-4.1", + durationMs: 100, + }, + ]); + const { updateMessageRequestDetailsDurably } = await import("@/repository/message"); + + await updateMessageRequestDetailsDurably(809, { + statusCode: 200, + providerChain: [{ id: 1, name: "provider-a", groupTag: "openai" }], + }); + await flushMicrotasks(); + + expect(mockDbUpdateSet).toHaveBeenCalledTimes(1); + expect(mockGetMessageWriterDb).not.toHaveBeenCalled(); + expect(mockEnqueueMessageRequestUpdateDurably).not.toHaveBeenCalled(); + expect(mockQueuePublicStatusRollupWrite).toHaveBeenCalledTimes(1); + }); + it("queues one rollup for duplicate terminal updates without double counting", async () => { mockDbInsertReturning.mockResolvedValue([ { diff --git a/tests/unit/repository/message-query-test-support.ts b/tests/unit/repository/message-query-test-support.ts new file mode 100644 index 000000000..428bba826 --- /dev/null +++ b/tests/unit/repository/message-query-test-support.ts @@ -0,0 +1,101 @@ +type JoinTrace = { + readonly source: unknown; + readonly predicate: unknown; +}; + +export type DrizzleQueryTrace = { + readonly from: unknown[]; + readonly where: unknown[]; + readonly leftJoins: JoinTrace[]; + readonly innerJoins: JoinTrace[]; + readonly groupBy: unknown[][]; + readonly orderBy: unknown[][]; + readonly limit: number[]; + readonly offset: number[]; +}; + +function isRecord(value: unknown): value is Readonly> { + return typeof value === "object" && value !== null; +} + +export type DrizzleQuery = PromiseLike & { + readonly trace: DrizzleQueryTrace; + readonly from: (source: unknown) => DrizzleQuery; + readonly where: (predicate: unknown) => DrizzleQuery; + readonly leftJoin: (source: unknown, predicate: unknown) => DrizzleQuery; + readonly innerJoin: (source: unknown, predicate: unknown) => DrizzleQuery; + readonly groupBy: (...expressions: unknown[]) => DrizzleQuery; + readonly orderBy: (...expressions: unknown[]) => DrizzleQuery; + readonly limit: (value: number) => DrizzleQuery; + readonly offset: (value: number) => DrizzleQuery; +}; + +export function createDrizzleQuery(result: TResult): DrizzleQuery { + const trace: DrizzleQueryTrace = { + from: [], + where: [], + leftJoins: [], + innerJoins: [], + groupBy: [], + orderBy: [], + limit: [], + offset: [], + }; + + const query = Object.assign(Promise.resolve(result), { + trace, + from: (source: unknown) => { + trace.from.push(source); + return query; + }, + where: (predicate: unknown) => { + trace.where.push(predicate); + return query; + }, + leftJoin: (source: unknown, predicate: unknown) => { + trace.leftJoins.push({ source, predicate }); + return query; + }, + innerJoin: (source: unknown, predicate: unknown) => { + trace.innerJoins.push({ source, predicate }); + return query; + }, + groupBy: (...expressions: unknown[]) => { + trace.groupBy.push(expressions); + return query; + }, + orderBy: (...expressions: unknown[]) => { + trace.orderBy.push(expressions); + return query; + }, + limit: (value: number) => { + trace.limit.push(value); + return query; + }, + offset: (value: number) => { + trace.offset.push(value); + return query; + }, + }); + + return query; +} + +export function sqlText(value: unknown): string { + const visited = new Set(); + + const visit = (node: unknown): string => { + if (node === null || node === undefined) return ""; + if (["string", "number", "boolean"].includes(typeof node)) return String(node); + if (Array.isArray(node)) return node.map(visit).join(" "); + if (!isRecord(node) || visited.has(node)) return ""; + + visited.add(node); + if ("queryChunks" in node) return visit(node.queryChunks); + if ("value" in node) return visit(node.value); + if (typeof node.name === "string") return node.name; + return Object.values(node).map(visit).join(" "); + }; + + return visit(value).replace(/\s+/g, " ").trim().toLowerCase(); +} diff --git a/tests/unit/repository/message-session-readback.test.ts b/tests/unit/repository/message-session-readback.test.ts new file mode 100644 index 000000000..215915026 --- /dev/null +++ b/tests/unit/repository/message-session-readback.test.ts @@ -0,0 +1,261 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; + +function installDbBoundary(db: object): void { + vi.doMock("@/drizzle/db", () => ({ + db, + getMessageWriterDb: vi.fn(() => ({ update: vi.fn(), execute: vi.fn() })), + })); + vi.doMock("@/lib/config/env.schema", () => ({ + getEnvConfig: vi.fn(() => ({ MESSAGE_REQUEST_WRITE_MODE: "sync" as const })), + isDevelopment: vi.fn(() => false), + })); + vi.doMock("@/lib/redis", () => ({ getRedisClient: vi.fn(() => null) })); +} + +function createLimitSelect(responses: readonly (readonly unknown[])[]) { + let callIndex = 0; + const events: string[] = []; + const select = vi.fn((_selection: unknown) => { + const rows = responses[callIndex] ?? []; + callIndex += 1; + const limit = vi.fn(async (_value: number) => { + events.push("limit"); + return rows; + }); + const orderBy = vi.fn((..._ordering: unknown[]) => { + events.push("orderBy"); + return { limit }; + }); + const where = vi.fn((_condition: unknown) => { + events.push("where"); + return { limit, orderBy }; + }); + const from = vi.fn((_table: unknown) => { + events.push("from"); + return { where }; + }); + return { from }; + }); + return { events, select }; +} + +function installLimitBoundaries( + responses: readonly (readonly unknown[])[], + messageTableHasData = true +) { + const { events, select } = createLimitSelect(responses); + const execute = vi.fn(async (_query: unknown) => [{ has_data: messageTableHasData }]); + installDbBoundary({ select, execute, update: vi.fn() }); + return { events, execute, select }; +} + +const CREATED_AT = new Date("2026-07-15T11:00:00.000Z"); +const SESSION_ROW = { + id: 1_101, + providerId: 31, + userId: 41, + key: "session-key", + model: "claude-sonnet-4", + originalModel: "claude-sonnet-4", + durationMs: 1_100, + costUsd: "0.330000000000000", + costMultiplier: "1.25", + sessionId: "session-readback", + userAgent: "vitest", + clientIp: "127.0.0.1", + messagesCount: 3, + statusCode: 200, + inputTokens: 90, + outputTokens: 30, + cacheCreationInputTokens: 0, + cacheReadInputTokens: 4, + cacheCreation5mInputTokens: 0, + cacheCreation1hInputTokens: 0, + cacheTtlApplied: null, + errorMessage: null, + providerChain: null, + blockedBy: null, + blockedReason: null, + createdAt: CREATED_AT, + updatedAt: CREATED_AT, + deletedAt: null, +}; +const LEDGER_ROW = { + requestId: 1_102, + finalProviderId: 32, + userId: 42, + key: "ledger-session-key", + model: "ledger-model", + originalModel: "requested-model", + endpoint: "/v1/responses", + statusCode: 200, + costUsd: "0.440000000000000", + costMultiplier: "1.5", + inputTokens: 100, + outputTokens: 40, + cacheCreationInputTokens: 2, + cacheReadInputTokens: 3, + cacheCreation5mInputTokens: 2, + cacheCreation1hInputTokens: 0, + cacheTtlApplied: "5m", + context1mApplied: true, + swapCacheTtlApplied: false, + durationMs: 1_200, + ttfbMs: 200, + sessionId: "ledger-session-readback", + createdAt: CREATED_AT, +}; + +describe("message session readback", () => { + afterEach(() => { + vi.doUnmock("@/drizzle/db"); + vi.doUnmock("@/lib/config/env.schema"); + vi.doUnmock("@/lib/redis"); + }); + + it("returns the newest direct request for a session", async () => { + vi.resetModules(); + const boundary = installLimitBoundaries([[SESSION_ROW]]); + const { findMessageRequestBySessionId } = await import("@/repository/message"); + + const result = await findMessageRequestBySessionId("session-readback"); + + expect(result).toMatchObject({ + id: 1_101, + sessionId: "session-readback", + costMultiplier: 1.25, + }); + expect(boundary.events).toEqual(["from", "where", "orderBy", "limit"]); + expect(boundary.execute).not.toHaveBeenCalled(); + }); + + it("falls back to the newest ledger request for a session in ledger-only mode", async () => { + vi.resetModules(); + const boundary = installLimitBoundaries([[], [LEDGER_ROW]], false); + const { findMessageRequestBySessionId } = await import("@/repository/message"); + + const result = await findMessageRequestBySessionId("ledger-session-readback"); + + expect(result).toMatchObject({ + id: 1_102, + providerId: 32, + sessionId: "ledger-session-readback", + endpoint: "/v1/responses", + userAgent: null, + }); + expect(boundary.select).toHaveBeenCalledTimes(2); + expect(boundary.execute).toHaveBeenCalledTimes(1); + }); + + it("returns null for a missing session when the message table remains authoritative", async () => { + vi.resetModules(); + const boundary = installLimitBoundaries([[]], true); + const { findMessageRequestBySessionId } = await import("@/repository/message"); + + const result = await findMessageRequestBySessionId("missing-session"); + + expect(result).toBeNull(); + expect(boundary.select).toHaveBeenCalledTimes(1); + }); + + it("returns the first initial-selection provider chain by request sequence", async () => { + vi.resetModules(); + const providerChain = [ + { + id: 31, + name: "origin-provider", + groupTag: "anthropic", + reason: "initial_selection" as const, + }, + ]; + const boundary = installLimitBoundaries([[{ providerChain }]]); + const { findSessionOriginChain } = await import("@/repository/message"); + + const result = await findSessionOriginChain("session-readback"); + + expect(result).toEqual(providerChain); + expect(boundary.events).toEqual(["from", "where", "orderBy", "limit"]); + }); + + it("returns paged requests in repository order with the legacy sequence fallback", async () => { + vi.resetModules(); + const events: string[] = []; + let queryIndex = 0; + const rows = [ + { + id: 1_104, + sequence: 4, + model: "model-four", + statusCode: 200, + costUsd: "0.04", + createdAt: CREATED_AT, + inputTokens: 4, + outputTokens: 2, + errorMessage: null, + }, + { + id: 1_103, + sequence: null, + model: "legacy-model", + statusCode: 500, + costUsd: null, + createdAt: CREATED_AT, + inputTokens: null, + outputTokens: null, + errorMessage: "failed", + }, + ]; + const select = vi.fn((_selection: unknown) => { + queryIndex += 1; + if (queryIndex === 1) { + return { from: vi.fn(() => ({ where: vi.fn(async () => [{ count: 2 }]) })) }; + } + return { + from: vi.fn(() => ({ + where: vi.fn(() => ({ + orderBy: vi.fn(() => ({ + limit: vi.fn((limit: number) => ({ + offset: vi.fn(async (offset: number) => { + events.push(`limit:${limit}`, `offset:${offset}`); + return rows; + }), + })), + })), + })), + })), + }; + }); + installDbBoundary({ select, execute: vi.fn(), update: vi.fn() }); + const { findRequestsBySessionId } = await import("@/repository/message"); + + const result = await findRequestsBySessionId("session-readback", { + limit: 2, + offset: 1, + order: "desc", + }); + + expect(result).toEqual({ + requests: [rows[0], { ...rows[1], sequence: 1 }], + total: 2, + }); + expect(events).toEqual(["limit:2", "offset:1"]); + }); + + it("returns neighboring request sequences with null-safe public results", async () => { + vi.resetModules(); + const responses = [[{ sequence: 3 }], []] as const; + let queryIndex = 0; + const select = vi.fn((_selection: unknown) => { + const rows = responses[queryIndex] ?? []; + queryIndex += 1; + return { from: vi.fn(() => ({ where: vi.fn(async () => rows) })) }; + }); + installDbBoundary({ select, execute: vi.fn(), update: vi.fn() }); + const { findAdjacentRequestSequences } = await import("@/repository/message"); + + const result = await findAdjacentRequestSequences("session-readback", 4); + + expect(result).toEqual({ prevSequence: 3, nextSequence: null }); + expect(select).toHaveBeenCalledTimes(2); + }); +}); diff --git a/tests/unit/repository/message-session-request-query.test.ts b/tests/unit/repository/message-session-request-query.test.ts new file mode 100644 index 000000000..0cea807d3 --- /dev/null +++ b/tests/unit/repository/message-session-request-query.test.ts @@ -0,0 +1,171 @@ +import { beforeEach, describe, expect, test, vi } from "vitest"; +import { messageRequest } from "@/drizzle/schema"; +import { findAdjacentRequestSequences, findRequestsBySessionId } from "@/repository/message"; +import { createDrizzleQuery, sqlText } from "./message-query-test-support"; + +const boundary = vi.hoisted(() => { + const writerDb = { execute: vi.fn<(query: unknown) => Promise>() }; + return { + select: vi.fn<(selection?: unknown) => unknown>(), + execute: vi.fn<(query: unknown) => Promise>(), + ledgerOnly: vi.fn<() => Promise>(), + getWriterDb: vi.fn(() => writerDb), + }; +}); + +vi.mock("@/drizzle/db", () => ({ + db: { select: boundary.select, execute: boundary.execute }, + getMessageWriterDb: boundary.getWriterDb, +})); +vi.mock("@/lib/config/env.schema", () => ({ + getEnvConfig: () => ({ MESSAGE_REQUEST_WRITE_MODE: "sync" }), + isDevelopment: () => false, +})); +vi.mock("@/lib/ledger-fallback", () => ({ isLedgerOnlyMode: boundary.ledgerOnly })); + +type MessageRow = typeof messageRequest.$inferSelect; +type RequestRow = Pick< + MessageRow, + | "id" + | "model" + | "statusCode" + | "costUsd" + | "createdAt" + | "inputTokens" + | "outputTokens" + | "errorMessage" +> & { readonly sequence: MessageRow["requestSequence"] }; + +const firstCreatedAt = new Date("2026-05-04T10:00:00.000Z"); +const secondCreatedAt = new Date("2026-05-04T10:01:00.000Z"); + +describe("message repository session request queries", () => { + beforeEach(() => { + boundary.select.mockReset(); + boundary.execute.mockReset(); + boundary.ledgerOnly.mockReset(); + }); + + test("returns the default page in ascending sequence order", async () => { + const count = createDrizzleQuery([{ count: 2 }]); + const rows = createDrizzleQuery([ + { + id: 31, + sequence: null, + model: "model-a", + statusCode: 200, + costUsd: "0.100000000000000", + createdAt: firstCreatedAt, + inputTokens: 10, + outputTokens: 5, + errorMessage: null, + }, + { + id: 32, + sequence: 3, + model: "model-b", + statusCode: 429, + costUsd: "0.200000000000000", + createdAt: secondCreatedAt, + inputTokens: 20, + outputTokens: 8, + errorMessage: "rate limited", + }, + ]); + boundary.select.mockReturnValueOnce(count).mockReturnValueOnce(rows); + + const result = await findRequestsBySessionId("session-requests"); + + expect(result).toEqual({ + total: 2, + requests: [ + { + id: 31, + sequence: 1, + model: "model-a", + statusCode: 200, + costUsd: "0.100000000000000", + createdAt: firstCreatedAt, + inputTokens: 10, + outputTokens: 5, + errorMessage: null, + }, + { + id: 32, + sequence: 3, + model: "model-b", + statusCode: 429, + costUsd: "0.200000000000000", + createdAt: secondCreatedAt, + inputTokens: 20, + outputTokens: 8, + errorMessage: "rate limited", + }, + ], + }); + expect(count.trace.from).toEqual([messageRequest]); + expect(rows.trace.from).toEqual([messageRequest]); + expect(sqlText(rows.trace.where)).toContain("session-requests"); + expect(sqlText(rows.trace.where)).toContain("deleted_at"); + expect(sqlText(rows.trace.orderBy)).toContain("request_sequence asc"); + expect(rows.trace.limit).toEqual([20]); + expect(rows.trace.offset).toEqual([0]); + }); + + test("applies descending order and explicit pagination", async () => { + const count = createDrizzleQuery([{ count: 5 }]); + const rows = createDrizzleQuery([ + { + id: 35, + sequence: 5, + model: null, + statusCode: null, + costUsd: null, + createdAt: secondCreatedAt, + inputTokens: null, + outputTokens: null, + errorMessage: null, + }, + ]); + boundary.select.mockReturnValueOnce(count).mockReturnValueOnce(rows); + + const result = await findRequestsBySessionId("session-desc", { + limit: 1, + offset: 2, + order: "desc", + }); + + expect(result.total).toBe(5); + expect(result.requests.map(({ sequence }) => sequence)).toEqual([5]); + expect(sqlText(rows.trace.orderBy)).toContain("request_sequence desc"); + expect(rows.trace.limit).toEqual([1]); + expect(rows.trace.offset).toEqual([2]); + }); + + test("returns adjacent neighbors using session-scoped sequence predicates", async () => { + const previous = createDrizzleQuery([{ sequence: 4 }]); + const next = createDrizzleQuery([{ sequence: 9 }]); + boundary.select.mockReturnValueOnce(previous).mockReturnValueOnce(next); + + const result = await findAdjacentRequestSequences("session-neighbors", 6); + + expect(result).toEqual({ prevSequence: 4, nextSequence: 9 }); + expect(previous.trace.from).toEqual([messageRequest]); + expect(next.trace.from).toEqual([messageRequest]); + expect(sqlText(previous.trace.where)).toContain("session-neighbors"); + expect(sqlText(previous.trace.where)).toContain("request_sequence < 6"); + expect(sqlText(next.trace.where)).toContain("request_sequence > 6"); + expect(sqlText(boundary.select.mock.calls.at(0)?.at(0))).toContain("max"); + expect(sqlText(boundary.select.mock.calls.at(1)?.at(0))).toContain("min"); + }); + + test("returns null neighbors when neither adjacent sequence exists", async () => { + boundary.select + .mockReturnValueOnce(createDrizzleQuery([{ sequence: null }])) + .mockReturnValueOnce(createDrizzleQuery([])); + + const result = await findAdjacentRequestSequences("session-isolated", 1); + + expect(result).toEqual({ prevSequence: null, nextSequence: null }); + }); +}); diff --git a/tests/unit/repository/message-terminal-cas-durable.test.ts b/tests/unit/repository/message-terminal-cas-durable.test.ts new file mode 100644 index 000000000..f4f26ce8e --- /dev/null +++ b/tests/unit/repository/message-terminal-cas-durable.test.ts @@ -0,0 +1,282 @@ +import { CasingCache } from "drizzle-orm/casing"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +type SqlQuery = { + toQuery: (config: { + escapeName: (name: string) => string; + escapeParam: (index: number) => string; + escapeString: (value: string) => string; + casing: CasingCache; + paramStartIndex: { value: number }; + }) => { sql: string; params: unknown[] }; +}; + +function createDeferred() { + return Promise.withResolvers(); +} + +function isSqlQuery(value: unknown): value is SqlQuery { + return typeof value === "object" && value !== null && "toQuery" in value; +} + +function renderSql(value: unknown) { + if (!isSqlQuery(value)) throw new TypeError("Expected a Drizzle SQL query"); + return value.toQuery({ + escapeName: (name) => `"${name}"`, + escapeParam: (index) => `$${index}`, + escapeString: (text) => `'${text}'`, + casing: new CasingCache(), + paramStartIndex: { value: 1 }, + }); +} + +async function flushMicrotasks(): Promise { + for (let index = 0; index < 12; index++) await Promise.resolve(); +} + +function installCasBoundaries(returnedRows: readonly { readonly id: number }[]) { + const writerReturning = vi.fn(async (_selection: unknown) => returnedRows); + const writerWhere = vi.fn((_condition: unknown) => ({ returning: writerReturning })); + const writerSet = vi.fn((_patch: Record) => ({ where: writerWhere })); + const writerUpdate = vi.fn((_table: unknown) => ({ set: writerSet })); + const defaultUpdate = vi.fn(); + const getRedisClient = vi.fn(() => null); + + vi.doMock("@/drizzle/db", () => ({ + db: { update: defaultUpdate, select: vi.fn(), execute: vi.fn() }, + getMessageWriterDb: vi.fn(() => ({ update: writerUpdate, execute: vi.fn() })), + })); + vi.doMock("@/lib/config/env.schema", () => ({ + getEnvConfig: vi.fn(() => ({ MESSAGE_REQUEST_WRITE_MODE: "async" as const })), + isDevelopment: vi.fn(() => false), + })); + vi.doMock("@/lib/redis", () => ({ getRedisClient })); + + return { defaultUpdate, getRedisClient, writerReturning, writerSet, writerWhere }; +} + +function installAsyncDurableBoundaries(execute: (query: SqlQuery) => Promise) { + const createdAt = new Date("2026-07-15T09:00:00.000Z"); + const select = vi.fn((_selection: unknown) => ({ + from: vi.fn((_table: unknown) => ({ + where: vi.fn((_condition: unknown) => ({ + limit: vi.fn(async (_limit: number) => [ + { createdAt, model: "gpt-4.1", originalModel: "gpt-4.1", durationMs: null }, + ]), + })), + })), + })); + const redisGet = vi.fn(async (key: string) => { + if (key === "public-status:v2:config-version:current") return "cfg-terminal"; + if (key !== "public-status:v2:config-internal:cfg-terminal") return null; + return JSON.stringify({ + configVersion: "cfg-terminal", + generatedAt: "2026-07-15T08:59:00.000Z", + siteTitle: "Status", + siteDescription: "Status", + timeZone: "UTC", + defaultIntervalMinutes: 5, + defaultRangeHours: 24, + groups: [ + { + sourceGroupId: 8, + sourceGroupName: "openai", + slug: "openai", + displayName: "OpenAI", + sortOrder: 1, + description: null, + models: [ + { + publicModelKey: "gpt-4.1", + label: "GPT-4.1", + vendorIconKey: "openai", + requestTypeBadge: "openaiCompatible", + }, + ], + }, + ], + }); + }); + const operations: string[] = []; + const pipelineExec = vi.fn(async () => operations.map(() => [null, 1] as const)); + const pipeline = { + hincrbyfloat: vi.fn((_key: string, field: string, _value: number) => operations.push(field)), + set: vi.fn((_key: string) => operations.push("coverage")), + expire: vi.fn((_key: string) => operations.push("expiry")), + exec: pipelineExec, + }; + const redis = { + status: "ready", + get: redisGet, + hincrbyfloat: vi.fn(), + pipeline: vi.fn(() => pipeline), + }; + + vi.doMock("@/drizzle/db", () => ({ + db: { select, update: vi.fn(), execute: vi.fn() }, + getMessageWriterDb: vi.fn(() => ({ execute, update: vi.fn() })), + })); + vi.doMock("@/lib/config/env.schema", () => ({ + getEnvConfig: vi.fn(() => ({ + MESSAGE_REQUEST_WRITE_MODE: "async" as const, + MESSAGE_REQUEST_ASYNC_FLUSH_INTERVAL_MS: 60_000, + MESSAGE_REQUEST_ASYNC_BATCH_SIZE: 1, + MESSAGE_REQUEST_ASYNC_MAX_PENDING: 100, + })), + isDevelopment: vi.fn(() => false), + })); + vi.doMock("@/lib/redis", () => ({ getRedisClient: vi.fn(() => redis) })); + + return { pipelineExec, redisGet }; +} + +function installSyncDurableBoundaries( + returnedRows: readonly { readonly id: number }[] = [{ id: 804 }] +) { + const returning = vi.fn(async (_selection: unknown) => returnedRows); + const where = vi.fn((_condition: unknown) => ({ returning })); + const set = vi.fn((_patch: Record) => ({ where })); + const update = vi.fn((_table: unknown) => ({ set })); + const writerExecute = vi.fn(); + vi.doMock("@/drizzle/db", () => ({ + db: { update, select: vi.fn(), execute: vi.fn() }, + getMessageWriterDb: vi.fn(() => ({ execute: writerExecute, update: vi.fn() })), + })); + vi.doMock("@/lib/config/env.schema", () => ({ + getEnvConfig: vi.fn(() => ({ MESSAGE_REQUEST_WRITE_MODE: "sync" as const })), + isDevelopment: vi.fn(() => false), + })); + vi.doMock("@/lib/redis", () => ({ getRedisClient: vi.fn(() => null) })); + return { set, writerExecute }; +} + +describe("message terminal CAS and durable acknowledgement", () => { + afterEach(() => { + vi.doUnmock("@/drizzle/db"); + vi.doUnmock("@/lib/config/env.schema"); + vi.doUnmock("@/lib/redis"); + }); + + it("claims an unfinalized request through the dedicated writer DB", async () => { + vi.resetModules(); + const boundary = installCasBoundaries([{ id: 801 }]); + const { updateMessageRequestDetailsIfUnfinalized } = await import("@/repository/message"); + + await updateMessageRequestDetailsIfUnfinalized(801, { + statusCode: 504, + errorMessage: "timeout", + }); + + expect(boundary.writerSet).toHaveBeenCalledWith( + expect.objectContaining({ statusCode: 504, errorMessage: "timeout" }) + ); + expect(renderSql(boundary.writerWhere.mock.calls[0]?.[0]).sql).toMatch(/status_code.*IS NULL/i); + expect(boundary.writerReturning).toHaveBeenCalledWith({ id: expect.anything() }); + expect(boundary.defaultUpdate).not.toHaveBeenCalled(); + }); + + it("leaves the existing terminal owner untouched when the CAS loses", async () => { + vi.resetModules(); + const boundary = installCasBoundaries([]); + const { updateMessageRequestDetailsIfUnfinalized } = await import("@/repository/message"); + + const result = await updateMessageRequestDetailsIfUnfinalized(802, { + statusCode: 500, + providerChain: [{ id: 4, name: "fallback", groupTag: "openai" }], + }); + + expect(result).toBe(false); + expect(boundary.writerReturning).toHaveBeenCalledTimes(1); + expect(boundary.getRedisClient).not.toHaveBeenCalled(); + expect(boundary.defaultUpdate).not.toHaveBeenCalled(); + }); + + it("keeps durable completion pending until SQL commits, then publishes the receipt", async () => { + vi.resetModules(); + const databaseCommit = createDeferred(); + const execute = vi.fn(async (_query: SqlQuery) => databaseCommit.promise); + const { pipelineExec, redisGet } = installAsyncDurableBoundaries(execute); + const { updateMessageRequestDetailsDurably } = await import("@/repository/message"); + + const completion = updateMessageRequestDetailsDurably(803, { + durationMs: 1_250, + statusCode: 200, + outputTokens: 50, + providerChain: [ + { id: 8, name: "winner", groupTag: "openai", reason: "request_success", statusCode: 200 }, + ], + model: "gpt-4.1", + }); + + let settled = false; + void completion.then( + () => { + settled = true; + }, + () => { + settled = true; + } + ); + await new Promise((resolve) => setImmediate(resolve)); + expect(settled).toBe(false); + expect(execute).toHaveBeenCalledTimes(1); + expect(redisGet).not.toHaveBeenCalled(); + + databaseCommit.resolve([{ id: 803 }]); + await completion; + await flushMicrotasks(); + + expect(redisGet).toHaveBeenCalled(); + expect(pipelineExec).toHaveBeenCalledTimes(1); + }); + + it("uses the direct DB update path when durable mode is synchronous", async () => { + vi.resetModules(); + const { set, writerExecute } = installSyncDurableBoundaries(); + const { updateMessageRequestDetailsDurably } = await import("@/repository/message"); + const onCommitted = vi.fn(); + + await updateMessageRequestDetailsDurably( + 804, + { durationMs: 900, statusCode: 201 }, + { onCommitted } + ); + + expect(set).toHaveBeenCalledWith( + expect.objectContaining({ durationMs: 900, statusCode: 201, updatedAt: expect.any(Date) }) + ); + expect(writerExecute).not.toHaveBeenCalled(); + expect(onCommitted).toHaveBeenCalledWith({ durationMs: 900, statusCode: 201 }); + }); + + it("does not wait for synchronous commit observers after SQL commits", async () => { + vi.resetModules(); + installSyncDurableBoundaries(); + const { updateMessageRequestDetailsDurably } = await import("@/repository/message"); + const observer = createDeferred(); + const onCommitted = vi.fn(() => observer.promise); + + await expect( + updateMessageRequestDetailsDurably(806, { durationMs: 900, statusCode: 201 }, { onCommitted }) + ).resolves.toBe(true); + expect(onCommitted).toHaveBeenCalledOnce(); + observer.resolve(); + }); + + it("does not publish a synchronous durable callback when the terminal CAS loses", async () => { + vi.resetModules(); + const { writerExecute } = installSyncDurableBoundaries([]); + const { updateMessageRequestDetailsDurably } = await import("@/repository/message"); + const onCommitted = vi.fn(); + + const committed = await updateMessageRequestDetailsDurably( + 805, + { durationMs: 1_100, statusCode: 502 }, + { onCommitted } + ); + + expect(committed).toBe(false); + expect(writerExecute).not.toHaveBeenCalled(); + expect(onCommitted).not.toHaveBeenCalled(); + }); +}); diff --git a/tests/unit/repository/message-terminal-cost-accounting.test.ts b/tests/unit/repository/message-terminal-cost-accounting.test.ts new file mode 100644 index 000000000..409c4afcf --- /dev/null +++ b/tests/unit/repository/message-terminal-cost-accounting.test.ts @@ -0,0 +1,133 @@ +import type { StoredCostBreakdown } from "@/types/cost-breakdown"; +import type { HedgeLoserBilling } from "@/types/cost-breakdown"; +import { CasingCache } from "drizzle-orm/casing"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +type SqlQuery = { + toQuery: (config: { + escapeName: (name: string) => string; + escapeParam: (index: number) => string; + escapeString: (value: string) => string; + casing: CasingCache; + paramStartIndex: { value: number }; + }) => { sql: string; params: unknown[] }; +}; + +function isSqlQuery(value: unknown): value is SqlQuery { + return typeof value === "object" && value !== null && "toQuery" in value; +} + +function renderSql(value: unknown) { + if (!isSqlQuery(value)) throw new TypeError("Expected a Drizzle SQL query"); + return value.toQuery({ + escapeName: (name) => `"${name}"`, + escapeParam: (index) => `$${index}`, + escapeString: (text) => `'${text}'`, + casing: new CasingCache(), + paramStartIndex: { value: 1 }, + }); +} + +function installCostBoundary(whereImplementation: (condition: unknown) => Promise) { + const where = vi.fn(whereImplementation); + const set = vi.fn((_patch: Record) => ({ where })); + const update = vi.fn((_table: unknown) => ({ set })); + vi.doMock("@/drizzle/db", () => ({ + db: { update, select: vi.fn(), execute: vi.fn() }, + getMessageWriterDb: vi.fn(() => ({ update: vi.fn(), execute: vi.fn() })), + })); + vi.doMock("@/lib/config/env.schema", () => ({ + getEnvConfig: vi.fn(() => ({ MESSAGE_REQUEST_WRITE_MODE: "sync" as const })), + isDevelopment: vi.fn(() => false), + })); + vi.doMock("@/lib/redis", () => ({ getRedisClient: vi.fn(() => null) })); + return { set, update, where }; +} + +const BREAKDOWN = { + input: "0.04", + output: "0.06", + cache_creation: "0", + cache_read: "0", + base_total: "0.10", + provider_multiplier: 1, + group_multiplier: 1, + total: "0.10", +} satisfies StoredCostBreakdown; + +const LOSER = { + providerId: 12, + providerName: "hedge-loser", + attemptNumber: 2, + costUsd: "0.015", + inputTokens: 40, + outputTokens: 5, +} satisfies HedgeLoserBilling; + +describe("message terminal cost accounting", () => { + afterEach(() => { + vi.useRealTimers(); + vi.doUnmock("@/drizzle/db"); + vi.doUnmock("@/lib/config/env.schema"); + vi.doUnmock("@/lib/redis"); + }); + + it("replaces winner cost with winner plus the authoritative loser sum", async () => { + vi.resetModules(); + const { set, update } = installCostBoundary(async () => []); + const { updateMessageRequestWinnerCost } = await import("@/repository/message"); + + await updateMessageRequestWinnerCost(901, "0.1", BREAKDOWN); + + const patch = set.mock.calls[0]?.[0]; + if (!patch) throw new Error("Winner update did not reach the DB boundary"); + const costSql = renderSql(patch.costUsd); + expect(costSql.sql).toMatch(/jsonb_array_elements.*hedge_losers/i); + expect(costSql.sql).toMatch(/SUM.*costUsd/i); + expect(costSql.params).toContain("0.100000000000000"); + expect(patch.costBreakdown).toEqual(BREAKDOWN); + expect(update).toHaveBeenCalledTimes(1); + }); + + it("guards a loser write by provider and attempt while atomically adding its cost", async () => { + vi.resetModules(); + const { set, update, where } = installCostBoundary(async () => []); + const { addMessageRequestHedgeLoserCost } = await import("@/repository/message"); + + await addMessageRequestHedgeLoserCost(902, "0.015", LOSER); + + const patch = set.mock.calls[0]?.[0]; + if (!patch) throw new Error("Loser update did not reach the DB boundary"); + const costSql = renderSql(patch.costUsd); + const losersSql = renderSql(patch.hedgeLosers); + const guardSql = renderSql(where.mock.calls[0]?.[0]); + expect(costSql.sql).toContain("COALESCE"); + expect(costSql.params).toContain("0.015000000000000"); + expect(losersSql.params).toContain(JSON.stringify([LOSER])); + expect(guardSql.sql).toContain("@>"); + expect(guardSql.params).toContain( + JSON.stringify([{ providerId: LOSER.providerId, attemptNumber: LOSER.attemptNumber }]) + ); + expect(update).toHaveBeenCalledTimes(1); + }); + + it("retries an ambiguous loser write without changing its idempotency key", async () => { + vi.resetModules(); + vi.useFakeTimers(); + let attempts = 0; + const boundary = installCostBoundary(async () => { + attempts += 1; + if (attempts < 3) throw new Error("transient writer failure"); + return []; + }); + const { addMessageRequestHedgeLoserCost } = await import("@/repository/message"); + + const completion = addMessageRequestHedgeLoserCost(903, "0.015", LOSER); + await vi.advanceTimersByTimeAsync(150); + await completion; + + expect(boundary.update).toHaveBeenCalledTimes(3); + const guards = boundary.where.mock.calls.map(([condition]) => renderSql(condition).params); + expect(new Set(guards.map((params) => JSON.stringify(params))).size).toBe(1); + }); +}); diff --git a/tests/unit/repository/message-terminal-public-status-seam.test.ts b/tests/unit/repository/message-terminal-public-status-seam.test.ts new file mode 100644 index 000000000..e684da8ef --- /dev/null +++ b/tests/unit/repository/message-terminal-public-status-seam.test.ts @@ -0,0 +1,548 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; + +type OwnerOrder = "primary-first" | "fallback-first"; + +type TerminalRow = { + id: number; + createdAt: Date; + model: string; + originalModel: string; + durationMs: number | null; + statusCode: number | null; +}; + +function createDeferred() { + let resolve!: (value: T | PromiseLike) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + return { promise, resolve, reject }; +} + +function toSqlText(query: { + toQuery: (config: { + escapeName: (name: string) => string; + escapeParam: (index: number) => string; + escapeString: (value: string) => string; + paramStartIndex: { value: number }; + }) => { sql: string; params: unknown[] }; +}) { + return query.toQuery({ + escapeName: (name) => `"${name}"`, + escapeParam: (index) => `$${index}`, + escapeString: (value) => `'${value}'`, + paramStartIndex: { value: 1 }, + }); +} + +async function flushMicrotasks(): Promise { + for (let index = 0; index < 12; index++) { + await Promise.resolve(); + } +} + +describe("message terminal public-status public seam", () => { + afterEach(() => { + vi.useRealTimers(); + vi.doUnmock("@/drizzle/db"); + vi.doUnmock("@/lib/config/env.schema"); + vi.doUnmock("@/lib/logger"); + vi.doUnmock("@/lib/redis"); + }); + + it.each(["primary-first", "fallback-first"])( + "%s publishes exactly one rollup from the terminal SQL owner", + async (ownerOrder) => { + vi.resetModules(); + vi.useFakeTimers(); + + const id = ownerOrder === "primary-first" ? 91_001 : 91_002; + const row: TerminalRow = { + id, + createdAt: new Date("2026-07-13T12:00:00.000Z"), + model: "gpt-4.1", + originalModel: "gpt-4.1", + durationMs: null, + statusCode: null, + }; + const releasePrimary = createDeferred(); + const primaryReceipts: number[][] = []; + const fallbackReceipts: number[][] = []; + const primarySql: Array<{ sql: string; params: unknown[] }> = []; + const rollupPipelines: Array> = []; + + const primaryDetails = { + durationMs: 1_200, + statusCode: 200, + outputTokens: 60, + providerChain: [ + { + id: 1, + name: "primary-provider", + groupTag: "openai", + reason: "request_success" as const, + statusCode: 200, + }, + ], + model: "gpt-4.1", + }; + const fallbackDetails = { + durationMs: 2_400, + statusCode: 504, + outputTokens: 0, + errorMessage: "Error: stream_finalization_timeout", + providerChain: [ + { + id: 2, + name: "fallback-provider", + groupTag: "openai", + reason: "retry_failed" as const, + statusCode: 504, + }, + ], + model: "gpt-4.1", + }; + + const execute = vi.fn(async (query: Parameters[0]) => { + const built = toSqlText(query); + primarySql.push(built); + await releasePrimary.promise; + if (row.statusCode !== null) { + primaryReceipts.push([]); + return []; + } + row.durationMs = primaryDetails.durationMs; + row.statusCode = primaryDetails.statusCode; + primaryReceipts.push([id]); + return [{ id }]; + }); + + const writerUpdate = vi.fn(() => ({ + set: vi.fn((patch: Record) => ({ + where: vi.fn(() => ({ + returning: vi.fn(async () => { + if (row.statusCode !== null) { + fallbackReceipts.push([]); + return []; + } + row.durationMs = patch.durationMs as number; + row.statusCode = patch.statusCode as number; + fallbackReceipts.push([id]); + return [{ id }]; + }), + })), + })), + })); + const writerDb = { execute, update: writerUpdate }; + + vi.doMock("@/drizzle/db", () => ({ + db: { + select: vi.fn(() => ({ + from: vi.fn(() => ({ + where: vi.fn(() => ({ + limit: vi.fn(async () => [ + { + createdAt: row.createdAt, + model: row.model, + originalModel: row.originalModel, + durationMs: row.durationMs, + }, + ]), + })), + })), + })), + update: vi.fn(), + }, + getMessageWriterDb: vi.fn(() => writerDb), + })); + vi.doMock("@/lib/config/env.schema", () => ({ + getEnvConfig: () => ({ + MESSAGE_REQUEST_WRITE_MODE: "async", + MESSAGE_REQUEST_ASYNC_FLUSH_INTERVAL_MS: 60_000, + MESSAGE_REQUEST_ASYNC_BATCH_SIZE: 1_000, + MESSAGE_REQUEST_ASYNC_MAX_PENDING: 1_000, + }), + })); + vi.doMock("@/lib/logger", () => ({ + logger: { + trace: vi.fn(), + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }, + })); + + const configSnapshot = JSON.stringify({ + configVersion: "cfg-r2-seam", + generatedAt: "2026-07-13T11:59:00.000Z", + siteTitle: "Status", + siteDescription: "Status", + timeZone: "UTC", + defaultIntervalMinutes: 5, + defaultRangeHours: 24, + groups: [ + { + sourceGroupId: 42, + sourceGroupName: "openai", + slug: "openai", + displayName: "OpenAI", + sortOrder: 1, + description: null, + models: [ + { + publicModelKey: "gpt-4.1", + label: "GPT-4.1", + vendorIconKey: "openai", + requestTypeBadge: "openaiCompatible", + }, + ], + }, + ], + }); + const redis = { + status: "ready", + hincrbyfloat: vi.fn(), + get: vi.fn(async (key: string) => { + if (key === "public-status:v2:config-version:current") { + return "cfg-r2-seam"; + } + if (key === "public-status:v2:config-internal:cfg-r2-seam") { + return configSnapshot; + } + return null; + }), + pipeline: vi.fn(() => { + const operations: Array<{ command: string; args: unknown[] }> = []; + return { + hincrbyfloat: (...args: unknown[]) => { + operations.push({ command: "hincrbyfloat", args }); + }, + set: (...args: unknown[]) => { + operations.push({ command: "set", args }); + }, + expire: (...args: unknown[]) => { + operations.push({ command: "expire", args }); + }, + exec: async () => { + rollupPipelines.push(operations); + return operations.map(() => [null, 1] as [null, number]); + }, + }; + }), + }; + vi.doMock("@/lib/redis", () => ({ + getRedisClient: vi.fn(() => redis), + })); + + const { updateMessageRequestDetailsDurably, updateMessageRequestDetailsIfUnfinalized } = + await import("@/repository/message"); + const { flushMessageRequestWriteBuffer, stopMessageRequestWriteBuffer } = await import( + "@/repository/message-write-buffer" + ); + + const primary = updateMessageRequestDetailsDurably(id, primaryDetails, { timeoutMs: 10 }); + const primaryResult = primary.catch((error: unknown) => error); + const flush = flushMessageRequestWriteBuffer(); + + await vi.advanceTimersByTimeAsync(10); + await expect(primaryResult).resolves.toEqual( + expect.objectContaining({ + message: "durable message_request acknowledgement timed out", + }) + ); + + if (ownerOrder === "fallback-first") { + await updateMessageRequestDetailsIfUnfinalized(id, fallbackDetails); + releasePrimary.resolve(); + await flush; + } else { + releasePrimary.resolve(); + await flush; + await updateMessageRequestDetailsIfUnfinalized(id, fallbackDetails); + } + await flushMicrotasks(); + + expect(primarySql).toHaveLength(1); + expect(primarySql[0]?.sql).toMatch(/"?status_code"? IS NULL/); + expect(primarySql[0]?.sql).toContain("RETURNING id"); + expect(primaryReceipts).toEqual(ownerOrder === "primary-first" ? [[id]] : [[]]); + expect(fallbackReceipts).toEqual(ownerOrder === "fallback-first" ? [[id]] : [[]]); + expect(row).toMatchObject( + ownerOrder === "primary-first" + ? { durationMs: primaryDetails.durationMs, statusCode: primaryDetails.statusCode } + : { durationMs: fallbackDetails.durationMs, statusCode: fallbackDetails.statusCode } + ); + expect(redis.get.mock.calls).toEqual([ + ["public-status:v2:config-version:current"], + ["public-status:v2:config-internal:cfg-r2-seam"], + ]); + expect(rollupPipelines).toHaveLength(1); + + const rollupFields = rollupPipelines[0]! + .filter((operation) => operation.command === "hincrbyfloat") + .map((operation) => String(operation.args[1])); + const expectedMetric = ownerOrder === "primary-first" ? "success" : "failure"; + const losingMetric = ownerOrder === "primary-first" ? "failure" : "success"; + expect(rollupFields).toContain(`42|gpt-4.1|${expectedMetric}`); + expect(rollupFields).not.toContain(`42|gpt-4.1|${losingMetric}`); + + await stopMessageRequestWriteBuffer(); + } + ); + + it("same-ID pending durable contention publishes one rollup from the first owner", async () => { + vi.resetModules(); + vi.useFakeTimers(); + + const id = 91_003; + const oldFailureDetails = { + durationMs: 4_200, + statusCode: 502, + inputTokens: 31, + outputTokens: 3, + ttfbMs: 900, + providerChain: [ + { + id: 11, + name: "old-failure-provider", + groupTag: "openai", + reason: "retry_failed" as const, + statusCode: 502, + }, + ], + providerId: 11, + errorMessage: "Error: old upstream failure", + model: "gpt-4.1", + }; + const latestSuccessDetails = { + durationMs: 1_500, + statusCode: 200, + outputTokens: 96, + ttfbMs: 300, + providerChain: [ + { + id: 22, + name: "latest-success-provider", + groupTag: "openai", + reason: "request_success" as const, + statusCode: 200, + }, + ], + providerId: 22, + model: "gpt-4.1", + }; + const row: TerminalRow & { + inputTokens: number | null; + outputTokens: number | null; + ttfbMs: number | null; + providerChain: unknown; + providerId: number | null; + } = { + id, + createdAt: new Date("2026-07-13T12:05:00.000Z"), + model: "gpt-4.1", + originalModel: "gpt-4.1", + durationMs: null, + statusCode: null, + inputTokens: null, + outputTokens: null, + ttfbMs: null, + providerChain: null, + providerId: null, + }; + const releaseCommit = createDeferred(); + const committedSql: Array<{ sql: string; params: unknown[] }> = []; + const rollupPipelines: Array> = []; + + const execute = vi.fn(async (query: Parameters[0]) => { + const built = toSqlText(query); + committedSql.push(built); + await releaseCommit.promise; + + const readCaseValue = (columnName: string): unknown => { + const column = `"${columnName}"`; + const clauseStart = built.sql.indexOf(`${column} = CASE id`); + const clauseEnd = built.sql.indexOf(`ELSE ${column} END`, clauseStart); + if (clauseStart === -1 || clauseEnd === -1) { + throw new Error(`Missing batch CASE clause for ${columnName}`); + } + + const parameterIndexes = Array.from( + built.sql.slice(clauseStart, clauseEnd).matchAll(/\$(\d+)/g), + (match) => Number(match[1]) + ); + const valueParameterIndex = parameterIndexes[1]; + if (valueParameterIndex === undefined) { + throw new Error(`Missing batch value parameter for ${columnName}`); + } + return built.params[valueParameterIndex - 1]; + }; + + row.durationMs = Number(readCaseValue("duration_ms")); + row.statusCode = Number(readCaseValue("status_code")); + row.inputTokens = Number(readCaseValue("input_tokens")); + row.outputTokens = Number(readCaseValue("output_tokens")); + row.ttfbMs = Number(readCaseValue("ttfb_ms")); + row.providerChain = JSON.parse(String(readCaseValue("provider_chain"))); + row.providerId = Number(readCaseValue("provider_id")); + return [{ id }]; + }); + const writerDb = { execute, update: vi.fn() }; + + vi.doMock("@/drizzle/db", () => ({ + db: { + select: vi.fn(() => ({ + from: vi.fn(() => ({ + where: vi.fn(() => ({ + limit: vi.fn(async () => [ + { + createdAt: row.createdAt, + model: row.model, + originalModel: row.originalModel, + durationMs: row.durationMs, + }, + ]), + })), + })), + })), + update: vi.fn(), + }, + getMessageWriterDb: vi.fn(() => writerDb), + })); + vi.doMock("@/lib/config/env.schema", () => ({ + getEnvConfig: () => ({ + MESSAGE_REQUEST_WRITE_MODE: "async", + MESSAGE_REQUEST_ASYNC_FLUSH_INTERVAL_MS: 60_000, + MESSAGE_REQUEST_ASYNC_BATCH_SIZE: 1_000, + MESSAGE_REQUEST_ASYNC_MAX_PENDING: 1_000, + }), + })); + vi.doMock("@/lib/logger", () => ({ + logger: { + trace: vi.fn(), + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }, + })); + + const configSnapshot = JSON.stringify({ + configVersion: "cfg-r2-same-id", + generatedAt: "2026-07-13T12:04:00.000Z", + siteTitle: "Status", + siteDescription: "Status", + timeZone: "UTC", + defaultIntervalMinutes: 5, + defaultRangeHours: 24, + groups: [ + { + sourceGroupId: 42, + sourceGroupName: "openai", + slug: "openai", + displayName: "OpenAI", + sortOrder: 1, + description: null, + models: [ + { + publicModelKey: "gpt-4.1", + label: "GPT-4.1", + vendorIconKey: "openai", + requestTypeBadge: "openaiCompatible", + }, + ], + }, + ], + }); + const redis = { + status: "ready", + hincrbyfloat: vi.fn(), + get: vi.fn(async (key: string) => { + if (key === "public-status:v2:config-version:current") { + return "cfg-r2-same-id"; + } + if (key === "public-status:v2:config-internal:cfg-r2-same-id") { + return configSnapshot; + } + return null; + }), + pipeline: vi.fn(() => { + const operations: Array<{ command: string; args: unknown[] }> = []; + return { + hincrbyfloat: (...args: unknown[]) => { + operations.push({ command: "hincrbyfloat", args }); + }, + set: (...args: unknown[]) => { + operations.push({ command: "set", args }); + }, + expire: (...args: unknown[]) => { + operations.push({ command: "expire", args }); + }, + exec: async () => { + rollupPipelines.push(operations); + return operations.map(() => [null, 1] as [null, number]); + }, + }; + }), + }; + vi.doMock("@/lib/redis", () => ({ + getRedisClient: vi.fn(() => redis), + })); + + const { updateMessageRequestDetailsDurably } = await import("@/repository/message"); + const { flushMessageRequestWriteBuffer, stopMessageRequestWriteBuffer } = await import( + "@/repository/message-write-buffer" + ); + + const oldFailure = updateMessageRequestDetailsDurably(id, oldFailureDetails); + const latestSuccess = updateMessageRequestDetailsDurably(id, latestSuccessDetails); + const flush = flushMessageRequestWriteBuffer(); + await flushMicrotasks(); + + expect(execute).toHaveBeenCalledTimes(1); + expect(rollupPipelines).toEqual([]); + expect(row.statusCode).toBeNull(); + + releaseCommit.resolve(); + const [oldFailureResult, latestSuccessResult] = await Promise.all([ + oldFailure, + latestSuccess, + flush, + ]); + await flushMicrotasks(); + + expect(oldFailureResult).toBe(true); + expect(latestSuccessResult).toBe(false); + + expect(committedSql).toHaveLength(1); + expect(committedSql[0]?.sql).toMatch(/"?status_code"? IS NULL/); + expect(committedSql[0]?.sql).toContain("RETURNING id"); + expect(row).toMatchObject({ + durationMs: oldFailureDetails.durationMs, + statusCode: oldFailureDetails.statusCode, + inputTokens: oldFailureDetails.inputTokens, + outputTokens: oldFailureDetails.outputTokens, + ttfbMs: oldFailureDetails.ttfbMs, + providerChain: oldFailureDetails.providerChain, + providerId: oldFailureDetails.providerId, + }); + expect(rollupPipelines).toHaveLength(1); + + const rollupIncrementOperations = rollupPipelines[0]!.filter( + (operation) => operation.command === "hincrbyfloat" + ); + expect(rollupIncrementOperations).toHaveLength(1); + const rollupIncrements = Object.fromEntries( + rollupIncrementOperations.map((operation) => [ + String(operation.args[1]), + Number(operation.args[2]), + ]) + ); + expect(rollupIncrements).toEqual({ "42|gpt-4.1|failure": 1 }); + + await stopMessageRequestWriteBuffer(); + }); +}); diff --git a/tests/unit/repository/message-terminal-write-apis.test.ts b/tests/unit/repository/message-terminal-write-apis.test.ts new file mode 100644 index 000000000..cf060c3a2 --- /dev/null +++ b/tests/unit/repository/message-terminal-write-apis.test.ts @@ -0,0 +1,182 @@ +import type { StoredCostBreakdown } from "@/types/cost-breakdown"; +import type { CreateMessageRequestData } from "@/types/message"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +function installSyncBoundaries(insertedRows: readonly Record[] = []) { + const insertReturning = vi.fn(async (_selection: unknown) => insertedRows); + const insertValues = vi.fn((_values: Record) => ({ + returning: insertReturning, + })); + const insert = vi.fn((_table: unknown) => ({ values: insertValues })); + const updateWhere = vi.fn(async (_condition: unknown) => []); + const updateSet = vi.fn((_values: Record) => ({ where: updateWhere })); + const update = vi.fn((_table: unknown) => ({ set: updateSet })); + const writerUpdate = vi.fn((_table: unknown) => ({ set: updateSet })); + + vi.doMock("@/drizzle/db", () => ({ + db: { insert, update, select: vi.fn(), execute: vi.fn() }, + getMessageWriterDb: vi.fn(() => ({ update: writerUpdate, execute: vi.fn() })), + })); + vi.doMock("@/lib/config/env.schema", () => ({ + getEnvConfig: vi.fn(() => ({ MESSAGE_REQUEST_WRITE_MODE: "sync" as const })), + isDevelopment: vi.fn(() => false), + })); + + return { insertValues, update, updateSet, updateWhere }; +} + +const BREAKDOWN = { + input: "0.01", + output: "0.02", + cache_creation: "0", + cache_read: "0", + base_total: "0.03", + provider_multiplier: 1.5, + group_multiplier: 2, + total: "0.09", +} satisfies StoredCostBreakdown; + +describe("message terminal write APIs", () => { + afterEach(() => { + vi.doUnmock("@/drizzle/db"); + vi.doUnmock("@/lib/config/env.schema"); + }); + + it("creates a request through the repository and returns its public row", async () => { + vi.resetModules(); + const createdAt = new Date("2026-07-15T08:00:00.000Z"); + const row = { + id: 701, + providerId: 11, + userId: 22, + key: "key-create", + model: "claude-sonnet-4", + originalModel: "claude-sonnet-4", + durationMs: 120, + costUsd: "0.125000000000000", + costMultiplier: "1.5", + sessionId: "session-create", + requestSequence: 3, + userAgent: "vitest", + clientIp: "127.0.0.1", + endpoint: "/v1/messages", + messagesCount: 2, + cacheTtlApplied: null, + cacheCreationInputTokens: 4, + cacheCreation5mInputTokens: 4, + cacheCreation1hInputTokens: 0, + cacheReadInputTokens: 5, + specialSettings: null, + createdAt, + updatedAt: createdAt, + deletedAt: null, + }; + const { insertValues } = installSyncBoundaries([row]); + const data = { + provider_id: 11, + user_id: 22, + key: "key-create", + model: "claude-sonnet-4", + original_model: "claude-sonnet-4", + duration_ms: 120, + cost_usd: "0.125", + cost_multiplier: 1.5, + group_cost_multiplier: 2, + session_id: "session-create", + request_sequence: 3, + user_agent: "vitest", + client_ip: "127.0.0.1", + endpoint: "/v1/messages", + messages_count: 2, + cache_creation_input_tokens: 4, + cache_creation_5m_input_tokens: 4, + cache_creation_1h_input_tokens: 0, + cache_read_input_tokens: 5, + } satisfies CreateMessageRequestData; + const { createMessageRequest } = await import("@/repository/message"); + + const result = await createMessageRequest(data); + + expect(result).toMatchObject({ + id: 701, + costUsd: "0.125000000000000", + costMultiplier: 1.5, + sessionId: "session-create", + createdAt, + }); + expect(insertValues).toHaveBeenCalledWith( + expect.objectContaining({ + costUsd: "0.125000000000000", + costMultiplier: "1.5", + groupCostMultiplier: "2", + requestSequence: 3, + }) + ); + }); + + it("writes duration through the synchronous database boundary", async () => { + vi.resetModules(); + const { updateSet, updateWhere } = installSyncBoundaries(); + const { updateMessageRequestDuration } = await import("@/repository/message"); + + const result = await updateMessageRequestDuration(702, 345); + + expect(result).toBeUndefined(); + expect(updateSet).toHaveBeenCalledWith({ durationMs: 345, updatedAt: expect.any(Date) }); + expect(updateWhere).toHaveBeenCalledTimes(1); + }); + + it("formats and writes the request cost", async () => { + vi.resetModules(); + const { updateSet } = installSyncBoundaries(); + const { updateMessageRequestCost } = await import("@/repository/message"); + + await updateMessageRequestCost(703, "0.123456789"); + + expect(updateSet).toHaveBeenCalledWith({ + costUsd: "0.123456789000000", + updatedAt: expect.any(Date), + }); + }); + + it("writes a formatted cost with its billing breakdown", async () => { + vi.resetModules(); + const { updateSet } = installSyncBoundaries(); + const { updateMessageRequestCostWithBreakdown } = await import("@/repository/message"); + + await updateMessageRequestCostWithBreakdown(704, "0.09", BREAKDOWN); + + expect(updateSet).toHaveBeenCalledWith({ + costUsd: "0.090000000000000", + costBreakdown: BREAKDOWN, + updatedAt: expect.any(Date), + }); + }); + + it("writes the supplied terminal detail fields without inventing omitted fields", async () => { + vi.resetModules(); + const { updateSet } = installSyncBoundaries(); + const { updateMessageRequestDetails } = await import("@/repository/message"); + const details = { + inputTokens: 101, + outputTokens: 23, + ttfbMs: null, + cacheCreationInputTokens: 7, + cacheReadInputTokens: 8, + cacheCreation5mInputTokens: 3, + cacheCreation1hInputTokens: 4, + cacheTtlApplied: "5m" as const, + errorMessage: "upstream closed", + model: "redirected-model", + actualResponseModel: null, + providerId: 91, + context1mApplied: true, + swapCacheTtlApplied: false, + }; + + await updateMessageRequestDetails(705, details); + + expect(updateSet).toHaveBeenCalledWith({ ...details, updatedAt: expect.any(Date) }); + expect(updateSet.mock.calls[0]?.[0]).not.toHaveProperty("statusCode"); + }); +}); diff --git a/tests/unit/repository/message-usage-logs-query.test.ts b/tests/unit/repository/message-usage-logs-query.test.ts new file mode 100644 index 000000000..1cc44667d --- /dev/null +++ b/tests/unit/repository/message-usage-logs-query.test.ts @@ -0,0 +1,203 @@ +import { beforeEach, describe, expect, test, vi } from "vitest"; +import { messageRequest, usageLedger } from "@/drizzle/schema"; +import { findUsageLogs } from "@/repository/message"; +import { createDrizzleQuery, sqlText } from "./message-query-test-support"; + +const boundary = vi.hoisted(() => { + const writerDb = { execute: vi.fn<(query: unknown) => Promise>() }; + return { + select: vi.fn<(selection?: unknown) => unknown>(), + execute: vi.fn<(query: unknown) => Promise>(), + ledgerOnly: vi.fn<() => Promise>(), + getWriterDb: vi.fn(() => writerDb), + }; +}); + +vi.mock("@/drizzle/db", () => ({ + db: { select: boundary.select, execute: boundary.execute }, + getMessageWriterDb: boundary.getWriterDb, +})); +vi.mock("@/lib/config/env.schema", () => ({ + getEnvConfig: () => ({ MESSAGE_REQUEST_WRITE_MODE: "sync" }), + isDevelopment: () => false, +})); +vi.mock("@/lib/ledger-fallback", () => ({ isLedgerOnlyMode: boundary.ledgerOnly })); + +type MessageRow = typeof messageRequest.$inferSelect; +type PrimaryRow = Pick< + MessageRow, + | "id" + | "providerId" + | "userId" + | "key" + | "model" + | "durationMs" + | "costUsd" + | "costMultiplier" + | "sessionId" + | "requestSequence" + | "statusCode" + | "inputTokens" + | "outputTokens" + | "cacheTtlApplied" + | "createdAt" + | "updatedAt" + | "deletedAt" +>; + +const createdAt = new Date("2026-05-03T12:00:00.000Z"); +const updatedAt = new Date("2026-05-03T12:00:01.000Z"); +const primaryRow = { + id: 41, + providerId: 7, + userId: 9, + key: "key-primary", + model: "model-primary", + durationMs: 120, + costUsd: "0.125000000000", + costMultiplier: "1.5", + sessionId: "session-primary", + requestSequence: 4, + statusCode: 200, + inputTokens: 80, + outputTokens: 20, + cacheTtlApplied: "1h", + createdAt, + updatedAt, + deletedAt: null, +} satisfies PrimaryRow; + +describe("message repository findUsageLogs", () => { + beforeEach(() => { + boundary.select.mockReset(); + boundary.execute.mockReset(); + boundary.ledgerOnly.mockReset(); + }); + + test("returns primary message logs with filters and offset pagination", async () => { + const count = createDrizzleQuery([{ count: 7 }]); + const rows = createDrizzleQuery([primaryRow]); + boundary.select.mockReturnValueOnce(count).mockReturnValueOnce(rows); + + const result = await findUsageLogs({ + userId: 9, + startDate: new Date("2026-05-01T00:00:00.000Z"), + endDate: new Date("2026-05-04T00:00:00.000Z"), + model: "model-primary", + page: 3, + pageSize: 25, + }); + + expect(result.total).toBe(7); + expect(result.logs).toEqual([ + expect.objectContaining({ + id: 41, + providerId: 7, + userId: 9, + model: "model-primary", + costUsd: "0.125000000000000", + costMultiplier: 1.5, + sessionId: "session-primary", + requestSequence: 4, + cacheTtlApplied: "1h", + createdAt, + updatedAt, + }), + ]); + expect(count.trace.from).toEqual([messageRequest]); + expect(rows.trace.from).toEqual([messageRequest]); + expect(sqlText(rows.trace.where)).toContain("deleted_at"); + expect(sqlText(rows.trace.where)).toContain("2026-05-01t00:00:00.000z"); + expect(sqlText(rows.trace.where)).toContain("model-primary"); + expect(sqlText(rows.trace.orderBy)).toContain("created_at desc"); + expect(rows.trace.limit).toEqual([25]); + expect(rows.trace.offset).toEqual([50]); + expect(boundary.ledgerOnly).not.toHaveBeenCalled(); + }); + + test("returns the empty primary page when ledger fallback mode is disabled", async () => { + const count = createDrizzleQuery([]); + const rows = createDrizzleQuery([]); + boundary.select.mockReturnValueOnce(count).mockReturnValueOnce(rows); + boundary.ledgerOnly.mockResolvedValueOnce(false); + + const result = await findUsageLogs({}); + + expect(result).toEqual({ logs: [], total: 0 }); + expect(boundary.select).toHaveBeenCalledTimes(2); + expect(boundary.ledgerOnly).toHaveBeenCalledOnce(); + }); + + test("falls back to ledger rows with equivalent filters and pagination", async () => { + const primaryCount = createDrizzleQuery([{ count: 0 }]); + const primaryRows = createDrizzleQuery([]); + const ledgerCount = createDrizzleQuery([{ count: 3 }]); + const ledgerRows = createDrizzleQuery([ + { + requestId: 88, + finalProviderId: 12, + userId: 9, + key: "key-ledger", + model: "model-ledger", + originalModel: "model-original", + endpoint: "/v1/messages", + statusCode: 201, + costUsd: "0.750000000000", + costMultiplier: "1.25", + inputTokens: 90, + outputTokens: 30, + cacheCreationInputTokens: 10, + cacheReadInputTokens: 15, + cacheCreation5mInputTokens: 4, + cacheCreation1hInputTokens: 6, + cacheTtlApplied: "mixed", + context1mApplied: true, + swapCacheTtlApplied: false, + durationMs: 250, + ttfbMs: 40, + sessionId: "session-ledger", + createdAt, + }, + ]); + boundary.select + .mockReturnValueOnce(primaryCount) + .mockReturnValueOnce(primaryRows) + .mockReturnValueOnce(ledgerCount) + .mockReturnValueOnce(ledgerRows); + boundary.ledgerOnly.mockResolvedValueOnce(true); + + const result = await findUsageLogs({ + userId: 9, + startDate: new Date("2026-05-01T00:00:00.000Z"), + endDate: new Date("2026-05-04T00:00:00.000Z"), + model: "model-ledger", + page: 2, + pageSize: 10, + }); + + expect(result.total).toBe(3); + expect(result.logs).toEqual([ + expect.objectContaining({ + id: 88, + providerId: 12, + key: "key-ledger", + model: "model-ledger", + originalModel: "model-original", + sessionId: "session-ledger", + userAgent: null, + costMultiplier: 1.25, + cacheTtlApplied: "mixed", + createdAt, + updatedAt: createdAt, + }), + ]); + expect(ledgerCount.trace.from).toEqual([usageLedger]); + expect(ledgerRows.trace.from).toEqual([usageLedger]); + expect(sqlText(ledgerRows.trace.where)).toContain("blocked_by"); + expect(sqlText(ledgerRows.trace.where)).toContain("model-ledger"); + expect(sqlText(ledgerRows.trace.orderBy)).toContain("created_at desc"); + expect(sqlText(ledgerRows.trace.orderBy)).toContain("request_id desc"); + expect(ledgerRows.trace.limit).toEqual([10]); + expect(ledgerRows.trace.offset).toEqual([10]); + }); +}); diff --git a/tests/unit/repository/message-write-buffer.test.ts b/tests/unit/repository/message-write-buffer.test.ts index 88c53d1aa..cdac4fc6c 100644 --- a/tests/unit/repository/message-write-buffer.test.ts +++ b/tests/unit/repository/message-write-buffer.test.ts @@ -29,6 +29,14 @@ function toSqlText(query: { toQuery: (config: any) => { sql: string; params: unk }); } +function successfulRowsForQuery(query: { + toQuery: (config: any) => { sql: string; params: unknown[] }; +}): Array<{ id: number }> { + const { params } = toSqlText(query); + const numericParams = params.filter((value): value is number => typeof value === "number"); + return Array.from(new Set(numericParams), (id) => ({ id })); +} + function createDeferred() { let resolve!: (value: T) => void; let reject!: (error: unknown) => void; @@ -50,11 +58,22 @@ describe("message_request 异步批量写入", () => { ]; const originalEnv = snapshotEnv(envKeys); - const executeMock = vi.fn(async () => []); + const executeMock = vi.fn(async () => [] as Array<{ id: number }>); + const defaultExecuteMock = vi.fn(async () => [] as Array<{ id: number }>); + const getMessageWriterDbMock = vi.fn(); + const loggerWarnMock = vi.fn(); + const loggerErrorMock = vi.fn(); beforeEach(() => { vi.resetModules(); - executeMock.mockClear(); + executeMock.mockReset(); + executeMock.mockImplementation(async (query) => successfulRowsForQuery(query)); + defaultExecuteMock.mockReset(); + defaultExecuteMock.mockImplementation(async (query) => successfulRowsForQuery(query)); + getMessageWriterDbMock.mockReset(); + getMessageWriterDbMock.mockReturnValue({ execute: executeMock }); + loggerWarnMock.mockReset(); + loggerErrorMock.mockReset(); process.env.NODE_ENV = "test"; process.env.DSN = "postgres://postgres:postgres@localhost:5432/claude_code_hub_test"; @@ -64,7 +83,7 @@ describe("message_request 异步批量写入", () => { vi.doMock("@/drizzle/db", () => ({ db: { - execute: executeMock, + execute: defaultExecuteMock, // 避免 tests/setup.ts 的 afterAll 清理逻辑因 mock 缺失 select 而报错 select: () => ({ from: () => ({ @@ -72,10 +91,21 @@ describe("message_request 异步批量写入", () => { }), }), }, + getMessageWriterDb: getMessageWriterDbMock, + })); + vi.doMock("@/lib/logger", () => ({ + logger: { + trace: vi.fn(), + debug: vi.fn(), + info: vi.fn(), + warn: loggerWarnMock, + error: loggerErrorMock, + }, })); }); afterEach(() => { + vi.useRealTimers(); restoreEnv(originalEnv); }); @@ -89,7 +119,9 @@ describe("message_request 异步批量写入", () => { enqueueMessageRequestUpdate(1, { durationMs: 123 }); await flushMessageRequestWriteBuffer(); + expect(getMessageWriterDbMock).not.toHaveBeenCalled(); expect(executeMock).not.toHaveBeenCalled(); + expect(defaultExecuteMock).not.toHaveBeenCalled(); }); it("async 模式下应合并同一 id 的多次更新并批量写入", async () => { @@ -118,6 +150,603 @@ describe("message_request 异步批量写入", () => { expect(built.sql).toContain("ttfb_ms"); expect(built.sql).toContain("updated_at"); expect(built.sql).toContain("deleted_at IS NULL"); + expect(built.sql).not.toContain("RETURNING id"); + }); + + it("batch SQL 应显式使用 writer DB handle,而不是默认 ALS DB", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + + const { + enqueueMessageRequestUpdate, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + enqueueMessageRequestUpdate(43, { durationMs: 101 }); + await flushMessageRequestWriteBuffer(); + await stopMessageRequestWriteBuffer(); + + expect(getMessageWriterDbMock).toHaveBeenCalledTimes(1); + expect(executeMock).toHaveBeenCalledTimes(1); + expect(defaultExecuteMock).not.toHaveBeenCalled(); + }); + + it("普通 enqueue 立即返回,但 durable enqueue 应等待批量 SQL 成功", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + + const deferred = createDeferred(); + executeMock.mockImplementationOnce(async () => deferred.promise); + + const { + enqueueMessageRequestUpdate, + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + expect(enqueueMessageRequestUpdate(1, { durationMs: 10 })).toBeUndefined(); + const durablePromise = enqueueMessageRequestUpdateDurably(2, { statusCode: 200 }); + const flushPromise = flushMessageRequestWriteBuffer(); + + expect(executeMock).toHaveBeenCalledTimes(1); + const built = toSqlText(executeMock.mock.calls[0]?.[0]); + expect(built.sql).toContain("RETURNING id"); + let settled = false; + void durablePromise.then( + () => { + settled = true; + }, + () => { + settled = true; + } + ); + await new Promise((resolve) => setImmediate(resolve)); + expect(settled).toBe(false); + + deferred.resolve([{ id: 2 }]); + await flushPromise; + await durablePromise; + await stopMessageRequestWriteBuffer(); + }); + + it("mixed batch 只对 durable id 应用 status_code 终态 fence", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + + const deferred = createDeferred(); + executeMock.mockImplementationOnce(async () => deferred.promise); + + const { + enqueueMessageRequestUpdate, + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + enqueueMessageRequestUpdate(51001, { durationMs: 17 }); + const durablePromise = enqueueMessageRequestUpdateDurably(52002, { statusCode: 503 }); + const flushPromise = flushMessageRequestWriteBuffer(); + + const built = toSqlText(executeMock.mock.calls[0]?.[0]); + const ordinaryIdOccurrences = built.params.filter((value) => value === 51001).length; + const durableIdOccurrences = built.params.filter((value) => value === 52002).length; + + expect(built.sql).toMatch(/"?status_code"? IS NULL/); + expect(built.sql).toContain("RETURNING id"); + expect(durableIdOccurrences).toBeGreaterThan(ordinaryIdOccurrences); + + deferred.resolve([{ id: 51001 }, { id: 52002 }]); + await flushPromise; + await expect(durablePromise).resolves.toBe(true); + await stopMessageRequestWriteBuffer(); + }); + + it("多个 durable 终态应由同一次 batch flush 共同确认", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + + const deferred = createDeferred(); + executeMock.mockImplementationOnce(async () => deferred.promise); + + const { + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + const first = enqueueMessageRequestUpdateDurably(11, { statusCode: 200 }); + const second = enqueueMessageRequestUpdateDurably(12, { statusCode: 500 }); + const flushPromise = flushMessageRequestWriteBuffer(); + + expect(executeMock).toHaveBeenCalledTimes(1); + const built = toSqlText(executeMock.mock.calls[0]?.[0]); + expect(built.params).toContain(11); + expect(built.params).toContain(12); + + deferred.resolve([{ id: 11 }, { id: 12 }]); + await flushPromise; + await expect(Promise.all([first, second])).resolves.toEqual([true, true]); + await stopMessageRequestWriteBuffer(); + }); + + it("同一 id 的后续 durable contender 不得覆盖首个 terminal owner", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + + const deferred = createDeferred(); + executeMock.mockImplementationOnce(async () => deferred.promise); + + const { + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + const ownerCallback = vi.fn(); + const contenderCallback = vi.fn(); + + const owner = enqueueMessageRequestUpdateDurably( + 13, + { statusCode: 200, errorMessage: "owner" }, + { onCommitted: ownerCallback } + ); + const contender = enqueueMessageRequestUpdateDurably( + 13, + { statusCode: 499, errorMessage: "contender" }, + { onCommitted: contenderCallback } + ); + const flushPromise = flushMessageRequestWriteBuffer(); + + deferred.resolve([{ id: 13 }]); + await flushPromise; + await expect(owner).resolves.toBe(true); + await expect(contender).resolves.toBe(false); + + const built = toSqlText(executeMock.mock.calls[0]?.[0]); + expect(built.params).toContain("owner"); + expect(built.params).not.toContain("contender"); + expect(ownerCallback).toHaveBeenCalledOnce(); + expect(contenderCallback).not.toHaveBeenCalled(); + expect(executeMock).toHaveBeenCalledTimes(1); + await stopMessageRequestWriteBuffer(); + }); + + it("DB flush 失败时不得确认 durable waiter,重试成功后才确认", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + + executeMock.mockRejectedValueOnce(new Error("db down")); + executeMock.mockResolvedValueOnce([{ id: 21 }]); + + const { + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + const durablePromise = enqueueMessageRequestUpdateDurably(21, { statusCode: 200 }); + await flushMessageRequestWriteBuffer(); + let settled = false; + void durablePromise.then( + () => { + settled = true; + }, + () => { + settled = true; + } + ); + await new Promise((resolve) => setImmediate(resolve)); + expect(settled).toBe(false); + + await flushMessageRequestWriteBuffer(); + await expect(durablePromise).resolves.toBe(true); + expect(executeMock).toHaveBeenCalledTimes(2); + await stopMessageRequestWriteBuffer(); + }); + + it("队列全部由 durable 终态保护时,应拒绝新的 durable id 而不丢旧终态", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + process.env.MESSAGE_REQUEST_ASYNC_MAX_PENDING = "100"; + + const { + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + const protectedPromises = Array.from({ length: 100 }, (_, index) => + enqueueMessageRequestUpdateDurably(1000 + index, { statusCode: 200 }) + ); + + await expect(enqueueMessageRequestUpdateDurably(9999, { statusCode: 500 })).rejects.toThrow( + "durable message_request queue is full" + ); + + await flushMessageRequestWriteBuffer(); + await expect(Promise.all(protectedPromises)).resolves.toHaveLength(100); + const built = toSqlText(executeMock.mock.calls[0]?.[0]); + expect(built.params).toContain(1000); + expect(built.params).toContain(1099); + expect(built.params).not.toContain(9999); + await stopMessageRequestWriteBuffer(); + }); + + it("durable ack timeout 后应清理 waiter,并允许同 id 后续重新提交", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + + const { + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + await expect( + enqueueMessageRequestUpdateDurably(31, { statusCode: 200 }, { timeoutMs: 10 }) + ).rejects.toThrow("durable message_request acknowledgement timed out"); + + const retry = enqueueMessageRequestUpdateDurably(31, { statusCode: 200 }); + await flushMessageRequestWriteBuffer(); + await expect(retry).resolves.toBe(true); + await stopMessageRequestWriteBuffer(); + }); + + it("pending durable ack 超时后应删除整代 patch,后续提交不得继承 stale 字段", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + + const { + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + vi.useFakeTimers(); + const staleGeneration = enqueueMessageRequestUpdateDurably( + 311, + { statusCode: 200, errorMessage: "stale-primary-generation" }, + { timeoutMs: 10 } + ); + const staleGenerationResult = staleGeneration.catch((error: unknown) => error); + await vi.advanceTimersByTimeAsync(10); + await expect(staleGenerationResult).resolves.toEqual( + expect.objectContaining({ + message: "durable message_request acknowledgement timed out", + }) + ); + + const retry = enqueueMessageRequestUpdateDurably(311, { statusCode: 502 }); + await flushMessageRequestWriteBuffer(); + await expect(retry).resolves.toBe(true); + + const built = toSqlText(executeMock.mock.calls[0]?.[0]); + expect(built.params).not.toContain("stale-primary-generation"); + await stopMessageRequestWriteBuffer(); + }); + + it("in-flight durable ack 超时后应允许同 id 重新提交且只确认新 batch", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + + const firstExecute = createDeferred(); + executeMock.mockImplementationOnce(async () => firstExecute.promise); + executeMock.mockResolvedValueOnce([{ id: 32 }]); + + const { + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + const first = enqueueMessageRequestUpdateDurably(32, { statusCode: 500 }, { timeoutMs: 10 }); + const flushPromise = flushMessageRequestWriteBuffer(); + await expect(first).rejects.toThrow("durable message_request acknowledgement timed out"); + + const retry = enqueueMessageRequestUpdateDurably(32, { statusCode: 200 }); + firstExecute.resolve([]); + + await flushPromise; + await expect(retry).resolves.toBe(true); + expect(executeMock).toHaveBeenCalledTimes(2); + await stopMessageRequestWriteBuffer(); + }); + + it("in-flight durable generation 超时且写入失败后不得作为普通 patch 重排", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + + const firstExecute = createDeferred(); + executeMock.mockImplementationOnce(async () => firstExecute.promise); + + const { + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + vi.useFakeTimers(); + const first = enqueueMessageRequestUpdateDurably( + 321, + { statusCode: 200, errorMessage: "stale-primary" }, + { timeoutMs: 10 } + ); + const firstResult = first.catch((error: unknown) => error); + const flushPromise = flushMessageRequestWriteBuffer(); + + await vi.advanceTimersByTimeAsync(10); + await expect(firstResult).resolves.toEqual( + expect.objectContaining({ + message: "durable message_request acknowledgement timed out", + }) + ); + firstExecute.reject(new Error("db down after timeout")); + await flushPromise; + + await flushMessageRequestWriteBuffer(); + expect(executeMock).toHaveBeenCalledTimes(1); + await stopMessageRequestWriteBuffer(); + }); + + it("durable ack 超时后 late primary 真正提交时仍只发布一次 commit receipt", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + + const releasePrimary = createDeferred(); + executeMock.mockImplementationOnce(async () => { + await releasePrimary.promise; + return [{ id: 323 }]; + }); + const onCommitted = vi.fn(); + + const { + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + vi.useFakeTimers(); + const primary = enqueueMessageRequestUpdateDurably( + 323, + { statusCode: 200 }, + { timeoutMs: 10, onCommitted } + ); + const primaryResult = primary.catch((error: unknown) => error); + const flushPromise = flushMessageRequestWriteBuffer(); + + await vi.advanceTimersByTimeAsync(10); + await expect(primaryResult).resolves.toEqual( + expect.objectContaining({ + message: "durable message_request acknowledgement timed out", + }) + ); + expect(onCommitted).not.toHaveBeenCalled(); + + releasePrimary.resolve(); + await flushPromise; + + expect(onCommitted).toHaveBeenCalledTimes(1); + await stopMessageRequestWriteBuffer(); + }); + + it("commit receipt 回调失败不得让已提交的 durable flush 失败", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + + const { + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + const durable = enqueueMessageRequestUpdateDurably( + 324, + { statusCode: 200 }, + { + onCommitted: () => { + throw new Error("rollup callback failed"); + }, + } + ); + + await expect(flushMessageRequestWriteBuffer()).resolves.toBeUndefined(); + await expect(durable).resolves.toBe(true); + expect(loggerErrorMock).toHaveBeenCalledWith( + "[MessageRequestWriteBuffer] Durable commit callback failed", + expect.objectContaining({ + error: "rollup callback failed", + messageRequestId: 324, + }) + ); + await stopMessageRequestWriteBuffer(); + }); + + it("stop 应等待已提交终态的异步 commit callback 完成", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + const callback = createDeferred(); + const { + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + const durable = enqueueMessageRequestUpdateDurably( + 325, + { statusCode: 200 }, + { onCommitted: () => callback.promise } + ); + await flushMessageRequestWriteBuffer(); + await durable; + + let stopped = false; + const stop = stopMessageRequestWriteBuffer().then(() => { + stopped = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(stopped).toBe(false); + + callback.resolve(); + await stop; + expect(stopped).toBe(true); + }); + + it("fallback CAS 先写入后,late durable primary 不得覆盖既有终态", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + + let persistedStatus: number | null = null; + const releasePrimary = createDeferred(); + const onCommitted = vi.fn(); + executeMock.mockImplementationOnce(async (query) => { + await releasePrimary.promise; + const built = toSqlText(query); + const hasTerminalFence = /"?status_code"? IS NULL/.test(built.sql); + if (hasTerminalFence && persistedStatus !== null) { + return []; + } + persistedStatus = 200; + return [{ id: 322 }]; + }); + + const { + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + vi.useFakeTimers(); + const primary = enqueueMessageRequestUpdateDurably( + 322, + { statusCode: 200 }, + { timeoutMs: 10, onCommitted } + ); + const primaryResult = primary.catch((error: unknown) => error); + const flushPromise = flushMessageRequestWriteBuffer(); + + await vi.advanceTimersByTimeAsync(10); + await expect(primaryResult).resolves.toEqual( + expect.objectContaining({ + message: "durable message_request acknowledgement timed out", + }) + ); + persistedStatus = 502; + releasePrimary.resolve(); + + await flushPromise; + expect(persistedStatus).toBe(502); + expect(onCommitted).not.toHaveBeenCalled(); + await stopMessageRequestWriteBuffer(); + }); + + it("durable batch 成功但目标行未更新时不得虚假确认", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + executeMock.mockResolvedValueOnce([]); + + const { + enqueueMessageRequestUpdateDurably, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + const durablePromise = enqueueMessageRequestUpdateDurably(33, { statusCode: 200 }); + await flushMessageRequestWriteBuffer(); + + await expect(durablePromise).rejects.toThrow( + "durable message_request update did not persist id 33" + ); + await stopMessageRequestWriteBuffer(); + }); + + it.each([ + { databaseOutcome: "成功", shouldReject: false }, + { databaseOutcome: "失败", shouldReject: true }, + ])( + "executor 首次同步重入 stop 时应共享同一 Promise, 并等待 DB $databaseOutcome", + async ({ shouldReject }) => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + + const databaseBarrier = createDeferred>(); + const databaseError = new Error("db unavailable"); + let reentrantStopPromise: Promise | undefined; + let stopMessageRequestWriteBuffer!: () => Promise; + + executeMock.mockImplementation((query) => { + if (!reentrantStopPromise) { + reentrantStopPromise = stopMessageRequestWriteBuffer(); + return databaseBarrier.promise; + } + return shouldReject + ? Promise.reject(databaseError) + : Promise.resolve(successfulRowsForQuery(query)); + }); + + const messageWriteBuffer = await import("@/repository/message-write-buffer"); + stopMessageRequestWriteBuffer = messageWriteBuffer.stopMessageRequestWriteBuffer; + messageWriteBuffer.enqueueMessageRequestUpdate(42, { durationMs: 100 }); + + const outerStopPromise = stopMessageRequestWriteBuffer(); + const reentrantPromise = reentrantStopPromise; + if (!reentrantPromise) { + throw new Error("executor did not synchronously re-enter stop"); + } + const samePromise = outerStopPromise === reentrantPromise; + let outerSettled = false; + let reentrantSettled = false; + void outerStopPromise.then( + () => { + outerSettled = true; + }, + () => { + outerSettled = true; + } + ); + void reentrantPromise.then( + () => { + reentrantSettled = true; + }, + () => { + reentrantSettled = true; + } + ); + await new Promise((resolve) => setImmediate(resolve)); + const settlementsBeforeRelease = [outerSettled, reentrantSettled]; + + if (shouldReject) { + databaseBarrier.reject(databaseError); + } else { + databaseBarrier.resolve([]); + } + const stopResults = await Promise.allSettled([outerStopPromise, reentrantPromise]); + + expect(settlementsBeforeRelease).toEqual([false, false]); + if (shouldReject) { + const shutdownError = "message_request writer shutdown persistence failed"; + expect(stopResults).toEqual([ + { status: "rejected", reason: expect.objectContaining({ message: shutdownError }) }, + { status: "rejected", reason: expect.objectContaining({ message: shutdownError }) }, + ]); + } else { + expect(stopResults).toEqual([ + { status: "fulfilled", value: undefined }, + { status: "fulfilled", value: undefined }, + ]); + } + expect(samePromise).toBe(true); + } + ); + + it("stop 无法刷写剩余终态时所有调用都应持续拒绝同一错误", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + executeMock.mockRejectedValue(new Error("db unavailable")); + + const { enqueueMessageRequestUpdateDurably, stopMessageRequestWriteBuffer } = await import( + "@/repository/message-write-buffer" + ); + + const durablePromise = enqueueMessageRequestUpdateDurably(41, { statusCode: 500 }); + const durableResult = durablePromise.catch((error: unknown) => error); + const shutdownError = "message_request writer shutdown persistence failed"; + const stopResults = await Promise.allSettled([ + stopMessageRequestWriteBuffer(), + stopMessageRequestWriteBuffer(), + ]); + + expect(stopResults).toEqual([ + { status: "rejected", reason: expect.objectContaining({ message: shutdownError }) }, + { status: "rejected", reason: expect.objectContaining({ message: shutdownError }) }, + ]); + + await expect(durableResult).resolves.toEqual( + expect.objectContaining({ message: shutdownError }) + ); + await expect(stopMessageRequestWriteBuffer()).rejects.toThrow(shutdownError); }); it("应对 costUsd/providerChain 做显式类型转换(numeric/jsonb)", async () => { @@ -159,11 +788,17 @@ describe("message_request 异步批量写入", () => { expect(executeMock).toHaveBeenCalledTimes(1); - const raced = await Promise.race([ - stopPromise.then(() => "stopped"), - Promise.resolve("pending"), - ]); - expect(raced).toBe("pending"); + let settled = false; + void stopPromise.then( + () => { + settled = true; + }, + () => { + settled = true; + } + ); + await new Promise((resolve) => setImmediate(resolve)); + expect(settled).toBe(false); deferred.resolve([]); await stopPromise; @@ -264,6 +899,87 @@ describe("message_request 异步批量写入", () => { expect(built.params).toContain(2099); }); + it("同 id patch 升级为终态后,overflow 索引应保留升级后的高优先级记录", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + process.env.MESSAGE_REQUEST_ASYNC_MAX_PENDING = "100"; + + const { enqueueMessageRequestUpdate, stopMessageRequestWriteBuffer } = await import( + "@/repository/message-write-buffer" + ); + + enqueueMessageRequestUpdate(3001, { model: "metadata-only" }); + enqueueMessageRequestUpdate(3002, { model: "evict-me" }); + for (let i = 0; i < 98; i++) { + enqueueMessageRequestUpdate(4000 + i, { durationMs: i }); + } + + enqueueMessageRequestUpdate(3001, { statusCode: 200 }); + enqueueMessageRequestUpdate(4999, { durationMs: 999 }); + await stopMessageRequestWriteBuffer(); + + const built = toSqlText(executeMock.mock.calls[0]?.[0]); + expect(built.params).toContain(3001); + expect(built.params).not.toContain(3002); + expect(built.params).toContain(4999); + }); + + it("DB 失败重排后,overflow 索引仍应淘汰最低优先级 ordinary patch", async () => { + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + process.env.MESSAGE_REQUEST_ASYNC_MAX_PENDING = "100"; + + const firstExecute = createDeferred(); + executeMock.mockImplementationOnce(async () => firstExecute.promise); + + const { + enqueueMessageRequestUpdate, + flushMessageRequestWriteBuffer, + stopMessageRequestWriteBuffer, + } = await import("@/repository/message-write-buffer"); + + enqueueMessageRequestUpdate(5001, { model: "old-low-priority" }); + const flushPromise = flushMessageRequestWriteBuffer(); + + for (let i = 0; i < 100; i++) { + enqueueMessageRequestUpdate(6000 + i, { durationMs: i }); + } + firstExecute.reject(new Error("db down")); + await flushPromise; + await stopMessageRequestWriteBuffer(); + + const retried = toSqlText(executeMock.mock.calls[1]?.[0]); + expect(retried.params).not.toContain(5001); + expect(retried.params).toContain(6000); + expect(retried.params).toContain(6099); + }); + + it("burst overflow 应限频为单条聚合告警", async () => { + vi.useFakeTimers(); + process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; + process.env.MESSAGE_REQUEST_ASYNC_MAX_PENDING = "100"; + + const { enqueueMessageRequestUpdate, stopMessageRequestWriteBuffer } = await import( + "@/repository/message-write-buffer" + ); + + for (let i = 0; i < 250; i++) { + enqueueMessageRequestUpdate(7000 + i, { durationMs: i }); + } + + expect(loggerWarnMock).not.toHaveBeenCalled(); + await vi.advanceTimersByTimeAsync(1000); + expect(loggerWarnMock).toHaveBeenCalledTimes(1); + expect(loggerWarnMock).toHaveBeenCalledWith( + "[MessageRequestWriteBuffer] Pending queue overflow, dropping updates", + expect.objectContaining({ + maxPending: 100, + droppedCount: 150, + currentPending: 100, + }) + ); + + await stopMessageRequestWriteBuffer(); + }); + it("costUsd 走纯替换语义(CASE id ... ::numeric,不累加)", async () => { process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; diff --git a/tests/unit/server-response-write-backpressure.test.ts b/tests/unit/server-response-write-backpressure.test.ts new file mode 100644 index 000000000..006c36aca --- /dev/null +++ b/tests/unit/server-response-write-backpressure.test.ts @@ -0,0 +1,491 @@ +import { EventEmitter } from "node:events"; +import http from "node:http"; +import { createRequire } from "node:module"; +import { Socket } from "node:net"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +const requireFromHere = createRequire(import.meta.url); + +type WebSocketLike = { + readyState: number; + send: (payload: string, callback?: (error?: Error) => void) => void; + close: (code: number, reason: string) => void; +}; + +type ServerModule = { + handleWebSocketConnection: ( + ws: WebSocketLike & EventEmitter, + request: { headers: Record; url: string } + ) => Promise; + forwardToInternalHttp: ( + ws: WebSocketLike, + request: { headers: Record; url: string }, + body: Record, + sessionId: string, + registerRequest?: ( + request: http.ClientRequest, + response?: http.IncomingMessage | null, + settleTurn?: () => boolean + ) => boolean | void, + close?: (code: number, reason: string) => void + ) => Promise; +}; + +const serverModule: ServerModule = requireFromHere("../../server.js"); + +function createClientRequest(writeResult: boolean, events: string[]): http.ClientRequest { + const request: http.ClientRequest = Object.create(http.ClientRequest.prototype); + EventEmitter.call(request); + Object.assign(request, { + destroyed: false, + write: () => { + events.push("write"); + return writeResult; + }, + end: () => { + events.push("end"); + return request; + }, + destroy: () => { + events.push("destroy"); + request.destroyed = true; + return request; + }, + }); + return request; +} + +function createIncomingResponse(): http.IncomingMessage { + const response = new http.IncomingMessage(new Socket()); + response.headers = { "content-type": "text/event-stream" }; + vi.spyOn(response, "pause"); + vi.spyOn(response, "resume"); + return response; +} + +function requestInput() { + return { + ws: { + readyState: 1, + send: vi.fn((_payload: string, callback?: (error?: Error) => void) => callback?.()), + close: vi.fn(), + }, + request: { headers: { authorization: "Bearer test" }, url: "/v1/responses" }, + body: { model: "gpt-5.5", input: "hello" }, + }; +} + +function createWebSocket(send: WebSocketLike["send"]) { + return Object.assign(new EventEmitter(), { readyState: 1, send, close: vi.fn() }); +} + +function forwardRequest( + request: http.ClientRequest, + close?: (code: number, reason: string) => void +) { + vi.spyOn(http, "request").mockImplementation(() => request); + const input = requestInput(); + return serverModule.forwardToInternalHttp( + input.ws, + input.request, + input.body, + "test-session", + undefined, + close + ); +} + +async function startSseBridge(send: WebSocketLike["send"]) { + const events: string[] = []; + const request = createClientRequest(true, events); + const response = createIncomingResponse(); + let respond: ((response: http.IncomingMessage) => void) | undefined; + vi.spyOn(http, "request").mockImplementation((_options, callback) => { + if (callback) respond = callback; + return request; + }); + const ws = createWebSocket(send); + await serverModule.handleWebSocketConnection(ws, { + headers: { host: "localhost" }, + url: "/v1/responses", + }); + ws.emit( + "message", + Buffer.from('{"type":"response.create","model":"gpt-5.5","input":"hello"}'), + false + ); + await Promise.resolve(); + respond?.(response); + return { events, request, response, ws }; +} + +afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); +}); + +describe("server response write backpressure", () => { + it("waits for request drain before ending a backpressured payload", async () => { + const events: string[] = []; + const request = createClientRequest(false, events); + const forwarding = forwardRequest(request); + + expect(events).toEqual(["write"]); + request.emit("drain"); + expect(events).toEqual(["write", "end"]); + request.emit("error", Object.assign(new Error("closed"), { code: "ECONNRESET" })); + await forwarding; + }); + + it.each(["ECONNREFUSED", "ECONNRESET"])( + "sends one fatal frame and waits for its acknowledgement on active request error %s", + async (code) => { + const events: string[] = []; + const request = createClientRequest(false, events); + vi.spyOn(http, "request").mockImplementation(() => request); + const input = requestInput(); + const sent: string[] = []; + let sendCallback: ((error?: Error) => void) | undefined; + input.ws.send = (payload, callback) => { + sent.push(payload); + sendCallback = callback; + }; + const close = vi.fn(); + + const forwarding = serverModule.forwardToInternalHttp( + input.ws, + input.request, + input.body, + "request-error-session", + undefined, + close + ); + let settled = false; + void forwarding.then(() => { + settled = true; + }); + + request.emit("error", Object.assign(new Error(code), { code })); + await new Promise((resolve) => setImmediate(resolve)); + + expect(sent).toHaveLength(1); + expect(JSON.parse(sent[0]).error.code).toBe("internal_request_error"); + expect(settled).toBe(false); + expect(close).not.toHaveBeenCalled(); + + sendCallback?.(); + await forwarding; + expect(close).toHaveBeenCalledWith(1011, "internal_request_error"); + + expect(() => request.emit("error", new Error("late request error"))).not.toThrow(); + expect(sent).toHaveLength(1); + } + ); + + it("force-settles an active turn without relying on request destroy events", async () => { + const events: string[] = []; + const request = createClientRequest(false, events); + vi.spyOn(http, "request").mockImplementation(() => request); + const input = requestInput(); + let settleTurn: (() => boolean) | undefined; + + const forwarding = serverModule.forwardToInternalHttp( + input.ws, + input.request, + input.body, + "force-settle-session", + (_request, _response, settle) => { + settleTurn = settle; + return true; + } + ); + + expect(settleTurn?.()).toBe(true); + await forwarding; + expect(events).toContain("destroy"); + expect(input.ws.send).not.toHaveBeenCalled(); + expect(() => request.emit("error", new Error("late request error"))).not.toThrow(); + }); + + it("lets close win before drain without ending twice or stranding completion", async () => { + const events: string[] = []; + const request = createClientRequest(false, events); + const forwarding = forwardRequest(request); + + request.emit("close"); + await forwarding; + + request.emit("drain"); + expect(events).toEqual(["write"]); + }); + + it("terminates a request body when drain never arrives", async () => { + vi.useFakeTimers(); + const events: string[] = []; + const request = createClientRequest(false, events); + const close = vi.fn(); + const forwarding = forwardRequest(request, close); + await vi.advanceTimersByTimeAsync(60_000); + await forwarding; + + expect(close.mock.calls).toEqual([[1011, "internal_request_drain_timeout"]]); + request.emit("drain"); + request.emit("error", new Error("late error")); + expect(events).toEqual(["write", "destroy"]); + }); + + it("serializes outbound sends and pauses SSE until callbacks release pressure", async () => { + const events: string[] = []; + const request = createClientRequest(true, events); + const response = createIncomingResponse(); + let respond: ((response: http.IncomingMessage) => void) | undefined; + vi.spyOn(http, "request").mockImplementation((_options, callback) => { + if (callback) respond = callback; + return request; + }); + const callbacks: Array<(error?: Error) => void> = []; + const sent: string[] = []; + const input = requestInput(); + input.ws.send = (payload, callback) => { + sent.push(payload); + if (callback) callbacks.push(callback); + }; + + const forwarding = serverModule.forwardToInternalHttp( + input.ws, + input.request, + input.body, + "session-2" + ); + respond?.(response); + response.emit( + "data", + 'data: {"type":"response.output_text.delta","delta":"a"}\n\n' + + 'data: {"type":"response.output_text.delta","delta":"b"}\n\n' + + 'data: {"type":"response.completed","response":{"id":"r1"}}\n\n' + ); + + expect(sent).toHaveLength(1); + expect(response.pause).toHaveBeenCalled(); + callbacks.shift()?.(); + expect(sent).toHaveLength(2); + callbacks.shift()?.(); + expect(sent).toHaveLength(3); + callbacks.shift()?.(); + expect(response.resume).toHaveBeenCalled(); + + response.emit("end"); + await forwarding; + response.destroy(); + }); + + it("invalidates late send callbacks and destroys both internal transports on close", async () => { + const callbacks: Array<(error?: Error) => void> = []; + const sent: string[] = []; + const bridge = await startSseBridge((payload, callback) => { + sent.push(payload); + if (callback) callbacks.push(callback); + }); + const destroyResponse = vi.spyOn(bridge.response, "destroy"); + bridge.response.emit( + "data", + 'data: {"type":"response.output_text.delta","delta":"a"}\n\n' + + 'data: {"type":"response.completed","response":{"id":"r1"}}\n\n' + ); + + expect(sent).toHaveLength(1); + bridge.ws.emit("close"); + expect(bridge.events).toContain("destroy"); + expect(destroyResponse).toHaveBeenCalledOnce(); + callbacks.shift()?.(); + expect(sent).toHaveLength(1); + }); + + it.each(["error", "timeout"] as const)("cleans up on send %s", async (failure) => { + if (failure === "timeout") vi.useFakeTimers(); + let callback: ((error?: Error) => void) | undefined; + const bridge = await startSseBridge((_payload, done) => { + callback = done; + }); + const destroyResponse = vi.spyOn(bridge.response, "destroy"); + bridge.response.emit("data", 'data: {"type":"response.output_text.delta","delta":"a"}\n\n'); + + if (failure === "error") callback?.(new Error("send failed")); + else await vi.advanceTimersByTimeAsync(60_000); + + expect(bridge.events).toContain("destroy"); + expect(destroyResponse).toHaveBeenCalledOnce(); + callback?.(); + expect(bridge.ws.close.mock.calls).toEqual([[1011, `outbound_send_${failure}`]]); + }); + + it("destroys upstream and closes once when outbound pending bytes overflow", async () => { + const bridge = await startSseBridge(vi.fn()); + const destroyResponse = vi.spyOn(bridge.response, "destroy"); + const delta = "x".repeat(1024 * 1024 + 1); + + bridge.response.emit( + "data", + `data: ${JSON.stringify({ type: "response.output_text.delta", delta })}\n\n` + ); + + expect(bridge.events).toContain("destroy"); + expect(destroyResponse).toHaveBeenCalledOnce(); + expect(bridge.ws.close.mock.calls).toEqual([[1011, "outbound_backpressure"]]); + }); + + it("sends a fatal terminal frame before initiating close", async () => { + let callback: ((error?: Error) => void) | undefined; + let payload = ""; + const bridge = await startSseBridge((sent, done) => { + payload = sent; + callback = done; + }); + + bridge.response.emit("end"); + + expect(JSON.parse(payload).error.code).toBe("stream_ended_without_terminal"); + expect(bridge.ws.close).not.toHaveBeenCalled(); + callback?.(); + expect(bridge.ws.close.mock.calls).toEqual([[1011, "stream_ended_without_terminal"]]); + }); + + it("sends a protocol error frame before closing the client", async () => { + const send = vi.fn(); + const ws = createWebSocket(send); + await serverModule.handleWebSocketConnection(ws, { headers: {}, url: "/v1/responses" }); + + ws.emit("message", Buffer.from("binary"), true); + + const payload = send.mock.calls[0]?.[0] ?? ""; + const callback = send.mock.calls[0]?.[1]; + expect(JSON.parse(payload)).toMatchObject({ error: { code: "invalid_frame_type" } }); + expect(ws.close).not.toHaveBeenCalled(); + callback?.(); + expect(ws.close).toHaveBeenCalledWith(1003, "binary_not_supported"); + }); + + it("aborts an active internal turn before a fatal protocol frame is acknowledged", async () => { + const callbacks: Array<(error?: Error) => void> = []; + const ws = createWebSocket((_payload, callback) => { + if (callback) callbacks.push(callback); + }); + const events: string[] = []; + const request = createClientRequest(true, events); + vi.spyOn(http, "request").mockImplementation(() => request); + await serverModule.handleWebSocketConnection(ws, { headers: {}, url: "/v1/responses" }); + + ws.emit( + "message", + Buffer.from('{"type":"response.create","model":"gpt-5.5","input":"hello"}'), + false + ); + await new Promise((resolve) => setImmediate(resolve)); + ws.emit("message", Buffer.from("binary"), true); + + expect(request.destroyed).toBe(true); + expect(ws.close).not.toHaveBeenCalled(); + callbacks.at(-1)?.(); + expect(ws.close).toHaveBeenCalledWith(1003, "binary_not_supported"); + }); + + it("settles an SSE turn when terminal acknowledgement precedes a premature close", async () => { + const events: string[] = []; + const request = createClientRequest(true, events); + const response = createIncomingResponse(); + response.complete = false; + let respond: ((response: http.IncomingMessage) => void) | undefined; + vi.spyOn(http, "request").mockImplementation((_options, callback) => { + if (callback) respond = callback; + return request; + }); + const close = vi.fn(); + const input = requestInput(); + const forwarding = serverModule.forwardToInternalHttp( + input.ws, + input.request, + input.body, + "terminal-close-session", + undefined, + close + ); + respond?.(response); + response.emit( + "data", + `data: ${JSON.stringify({ type: "response.completed", response: { id: "r1" } })}\n\n` + ); + response.emit("close"); + + await forwarding; + expect(close).not.toHaveBeenCalled(); + }); + + it("waits for the JSON terminal send acknowledgement across end and close", async () => { + const events: string[] = []; + const request = createClientRequest(true, events); + const response = new http.IncomingMessage(new Socket()); + response.headers = { "content-type": "application/json" }; + let respond: ((response: http.IncomingMessage) => void) | undefined; + vi.spyOn(http, "request").mockImplementation((_options, callback) => { + if (callback) respond = callback; + return request; + }); + let sendCallback: ((error?: Error) => void) | undefined; + const sent: string[] = []; + const close = vi.fn(); + const input = requestInput(); + input.ws.send = (payload, callback) => { + sent.push(payload); + sendCallback = callback; + }; + + const forwarding = serverModule.forwardToInternalHttp( + input.ws, + input.request, + input.body, + "json-session", + undefined, + close + ); + respond?.(response); + response.emit("data", Buffer.from('{"id":"response-1"}')); + response.emit("end"); + response.emit("close"); + + let settled = false; + void forwarding.then(() => { + settled = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(settled).toBe(false); + expect(sent.map((payload) => JSON.parse(payload).type)).toEqual(["response.completed"]); + expect(close).not.toHaveBeenCalled(); + + sendCallback?.(); + await forwarding; + expect(close).not.toHaveBeenCalled(); + }); + + it("does not start a request while a fatal protocol frame awaits acknowledgement", async () => { + const callbacks: Array<(error?: Error) => void> = []; + const ws = createWebSocket((_payload, callback) => { + if (callback) callbacks.push(callback); + }); + const requestSpy = vi + .spyOn(http, "request") + .mockImplementation(() => createClientRequest(true, [])); + await serverModule.handleWebSocketConnection(ws, { headers: {}, url: "/v1/responses" }); + + ws.emit("message", Buffer.from("binary"), true); + ws.emit( + "message", + Buffer.from('{"type":"response.create","model":"gpt-5.5","input":"hello"}'), + false + ); + await new Promise((resolve) => setImmediate(resolve)); + + expect(requestSpy).not.toHaveBeenCalled(); + callbacks.shift()?.(); + expect(ws.close).toHaveBeenCalledWith(1003, "binary_not_supported"); + }); +}); diff --git a/tests/unit/server-shutdown.test.ts b/tests/unit/server-shutdown.test.ts index 748715779..caa628d97 100644 --- a/tests/unit/server-shutdown.test.ts +++ b/tests/unit/server-shutdown.test.ts @@ -9,6 +9,7 @@ * - drain timeout fires when server.close never finishes */ +import { spawnSync } from "node:child_process"; import { createRequire } from "node:module"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; @@ -17,7 +18,7 @@ const requireFromHere = createRequire(import.meta.url); type ServerJsModule = { registerOrchestratedShutdown: ( server: { close: (cb: (err?: Error) => void) => void; on?: unknown }, - wss: { close: () => void } | null + wss: { close: (cb?: (err?: Error) => void) => void } | null ) => void; }; @@ -28,22 +29,27 @@ function loadServerModule(): ServerJsModule { describe.sequential("registerOrchestratedShutdown", () => { let prevExit: typeof process.exit; let originalSigterm: typeof process.on; + let prevStdoutWrite: typeof process.stdout.write; beforeEach(() => { + vi.resetModules(); prevExit = process.exit; originalSigterm = process.on; + prevStdoutWrite = process.stdout.write; delete (globalThis as unknown as { __CCH_LIFECYCLE__?: unknown }).__CCH_LIFECYCLE__; }); afterEach(() => { process.exit = prevExit; process.on = originalSigterm; + process.stdout.write = prevStdoutWrite; process.removeAllListeners("SIGTERM"); process.removeAllListeners("SIGINT"); delete (globalThis as unknown as { __CCH_LIFECYCLE__?: unknown }).__CCH_LIFECYCLE__; delete process.env.SHUTDOWN_DRAIN_MS; delete process.env.SHUTDOWN_CLEANUP_MS; delete process.env.SHUTDOWN_HARD_EXIT_MS; + vi.restoreAllMocks(); }); it("runs the full sequence: markShuttingDown -> server.close -> runApplicationCleanup -> exit(0)", async () => { @@ -89,6 +95,57 @@ describe.sequential("registerOrchestratedShutdown", () => { expect(exitSpy).toHaveBeenCalledWith(0); }); + it("waits for WSS close callback before cleanup and successful exit on SIGTERM/SIGINT", async () => { + process.env.SHUTDOWN_DRAIN_MS = "500"; + process.env.SHUTDOWN_CLEANUP_MS = "500"; + process.env.SHUTDOWN_HARD_EXIT_MS = "5000"; + + const { registerOrchestratedShutdown } = loadServerModule(); + + for (const signal of ["SIGTERM", "SIGINT"] as const) { + process.removeAllListeners("SIGTERM"); + process.removeAllListeners("SIGINT"); + + const closeServer = vi.fn((callback: (err?: Error) => void) => callback()); + let finishWssClose: ((err?: Error) => void) | undefined; + const closeWss = vi.fn((callback?: (err?: Error) => void) => { + finishWssClose = callback; + }); + const runApplicationCleanup = vi.fn(async () => {}); + (globalThis as unknown as { __CCH_LIFECYCLE__?: unknown }).__CCH_LIFECYCLE__ = { + markShuttingDown: vi.fn(), + isShuttingDown: vi.fn(() => true), + runApplicationCleanup, + }; + + const output: string[] = []; + process.stdout.write = ((chunk: string | Uint8Array) => { + output.push(String(chunk)); + return true; + }) as typeof process.stdout.write; + const exitSpy = vi.fn() as unknown as typeof process.exit; + process.exit = exitSpy; + + registerOrchestratedShutdown({ close: closeServer }, { close: closeWss }); + process.emit(signal); + await new Promise((resolve) => setImmediate(resolve)); + + expect(closeServer).toHaveBeenCalledTimes(1); + expect(closeWss).toHaveBeenCalledTimes(1); + expect(runApplicationCleanup).not.toHaveBeenCalled(); + expect(output.join("")).not.toContain('"msg":"shutdown_complete"'); + expect(exitSpy).not.toHaveBeenCalled(); + expect(finishWssClose).toBeTypeOf("function"); + + finishWssClose?.(); + await vi.waitFor(() => { + expect(runApplicationCleanup).toHaveBeenCalledWith(signal, { totalTimeoutMs: 500 }); + expect(output.join("")).toContain('"msg":"shutdown_complete"'); + expect(exitSpy).toHaveBeenCalledWith(0); + }); + } + }); + it("drain timeout fires when server.close never resolves", async () => { process.env.SHUTDOWN_DRAIN_MS = "100"; process.env.SHUTDOWN_CLEANUP_MS = "100"; @@ -156,7 +213,7 @@ describe.sequential("registerOrchestratedShutdown", () => { expect(exitSpy).toHaveBeenCalledWith(0); }); - it("survives missing lifecycle globals (logs warning, still exits)", async () => { + it("lifecycle globals 缺失时不得记录 shutdown_complete 或 exit(0)", async () => { process.env.SHUTDOWN_DRAIN_MS = "50"; process.env.SHUTDOWN_CLEANUP_MS = "50"; process.env.SHUTDOWN_HARD_EXIT_MS = "5000"; @@ -166,6 +223,12 @@ describe.sequential("registerOrchestratedShutdown", () => { delete (globalThis as unknown as { __CCH_LIFECYCLE__?: unknown }).__CCH_LIFECYCLE__; + const output: string[] = []; + process.stdout.write = ((chunk: string | Uint8Array) => { + output.push(String(chunk)); + return true; + }) as typeof process.stdout.write; + const exitSpy = vi.fn() as unknown as typeof process.exit; process.exit = exitSpy; @@ -176,6 +239,97 @@ describe.sequential("registerOrchestratedShutdown", () => { await new Promise((resolve) => setTimeout(resolve, 200)); expect(closeServer).toHaveBeenCalled(); - expect(exitSpy).toHaveBeenCalledWith(0); + expect(exitSpy).toHaveBeenCalledTimes(1); + expect(exitSpy).toHaveBeenCalledWith(1); + expect(exitSpy).not.toHaveBeenCalledWith(0); + expect(output.join("")).toContain('"msg":"shutdown_cleanup_unavailable"'); + expect(output.join("")).not.toContain('"msg":"shutdown_complete"'); + }); + + it("writer rejection 时不得记录 shutdown_complete 或 exit(0)", async () => { + process.env.SHUTDOWN_DRAIN_MS = "50"; + process.env.SHUTDOWN_CLEANUP_MS = "500"; + process.env.SHUTDOWN_HARD_EXIT_MS = "1000"; + + vi.doMock("@/lib/cache/session-cache", () => ({ stopCacheCleanup: () => {} })); + vi.doMock("@/lib/provider-endpoints/probe-scheduler", () => ({ + stopEndpointProbeScheduler: () => {}, + })); + vi.doMock("@/lib/public-status/scheduler", () => ({ + stopPublicStatusRebuildScheduler: async () => {}, + })); + vi.doMock("@/lib/provider-endpoints/probe-log-cleanup", () => ({ + stopEndpointProbeLogCleanup: () => {}, + })); + vi.doMock("@/lib/async-task-manager", () => ({ shutdownAllAsyncTasks: async () => {} })); + vi.doMock("@/repository/message-write-buffer", () => ({ + stopMessageRequestWriteBuffer: async () => { + throw new Error("writer rejected"); + }, + })); + vi.doMock("@/drizzle/db", () => ({ closeDbPools: async () => {} })); + vi.doMock("@/lib/langfuse", () => ({ shutdownLangfuse: async () => {} })); + vi.doMock("@/lib/redis", () => ({ closeRedis: async () => {} })); + + const lifecycle = await import("@/lib/lifecycle/shutdown"); + lifecycle.__resetShutdownStateForTests(); + (globalThis as unknown as { __CCH_LIFECYCLE__?: unknown }).__CCH_LIFECYCLE__ = lifecycle; + + const output: string[] = []; + process.stdout.write = ((chunk: string | Uint8Array) => { + output.push(String(chunk)); + return true; + }) as typeof process.stdout.write; + + const exitSpy = vi.fn() as unknown as typeof process.exit; + process.exit = exitSpy; + + const setTimeoutSpy = vi.spyOn(globalThis, "setTimeout"); + const clearTimeoutSpy = vi.spyOn(globalThis, "clearTimeout"); + + const { registerOrchestratedShutdown } = loadServerModule(); + registerOrchestratedShutdown({ close: (callback) => callback() }, null); + process.emit("SIGTERM"); + const hardExit = setTimeoutSpy.mock.results[0]?.value; + + try { + await new Promise((resolve) => setTimeout(resolve, 250)); + + expect(exitSpy).toHaveBeenCalledTimes(1); + expect(exitSpy).toHaveBeenCalledWith(1); + expect(exitSpy).not.toHaveBeenCalledWith(0); + expect(output.join("")).toContain('"msg":"shutdown_cleanup_error"'); + expect(output.join("")).not.toContain('"msg":"shutdown_complete"'); + expect(clearTimeoutSpy).not.toHaveBeenCalledWith(hardExit); + } finally { + clearTimeout(hardExit); + } + }); + + it("无其他 ref handle 时 hard watchdog 仍以非零状态退出", () => { + const serverPath = requireFromHere.resolve("../../server.js"); + const script = ` + process.env.SHUTDOWN_DRAIN_MS = "10"; + process.env.SHUTDOWN_CLEANUP_MS = "10"; + process.env.SHUTDOWN_HARD_EXIT_MS = "50"; + const { registerOrchestratedShutdown } = require(${JSON.stringify(serverPath)}); + globalThis.__CCH_LIFECYCLE__ = { + markShuttingDown() {}, + isShuttingDown() { return true; }, + runApplicationCleanup() { return new Promise(() => {}); }, + }; + registerOrchestratedShutdown({ close(callback) { callback(); } }, null); + process.emit("SIGTERM"); + `; + + const result = spawnSync(process.execPath, ["-e", script], { + encoding: "utf8", + timeout: 2_000, + }); + + expect(result.error).toBeUndefined(); + expect(result.status).toBe(1); + expect(result.stdout).toContain('"msg":"shutdown_hard_exit_watchdog"'); + expect(result.stdout).not.toContain('"msg":"shutdown_complete"'); }); }); diff --git a/tests/unit/usage-ledger/backfill.test.ts b/tests/unit/usage-ledger/backfill.test.ts index 7c8a7650d..e78efc6e7 100644 --- a/tests/unit/usage-ledger/backfill.test.ts +++ b/tests/unit/usage-ledger/backfill.test.ts @@ -1,12 +1,15 @@ import { readFileSync } from "node:fs"; import { resolve } from "node:path"; -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; process.env.DSN = ""; +const mockTransaction = vi.hoisted(() => vi.fn()); + vi.mock("@/drizzle/db", () => ({ db: { execute: vi.fn(), + transaction: mockTransaction, }, })); @@ -23,6 +26,10 @@ const serviceSource = readFileSync( ); describe("backfillUsageLedger", () => { + beforeEach(() => { + mockTransaction.mockReset(); + }); + it("exports backfillUsageLedger function", () => { expect(typeof backfillUsageLedger).toBe("function"); }); @@ -39,4 +46,36 @@ describe("backfillUsageLedger", () => { expect(serviceSource).toContain("success_rate_outcome"); expect(serviceSource).toContain("fn_compute_message_request_success_rate_outcome"); }); + + it("rejects before opening a transaction when already aborted", async () => { + const controller = new AbortController(); + controller.abort(); + + await expect(backfillUsageLedger(controller.signal)).rejects.toMatchObject({ + name: "AbortError", + }); + expect(mockTransaction).not.toHaveBeenCalled(); + }); + + it("observes abort after a batch and does not start another batch", async () => { + const controller = new AbortController(); + let resolveBatch!: (value: unknown[]) => void; + const batch = new Promise((resolve) => { + resolveBatch = resolve; + }); + const execute = vi + .fn() + .mockResolvedValueOnce([{ acquired: true }]) + .mockReturnValueOnce(batch) + .mockResolvedValueOnce([{ processed: 0, inserted: 0, updated: 0, max_id: 0 }]); + mockTransaction.mockImplementation(async (callback) => callback({ execute })); + + const backfill = backfillUsageLedger(controller.signal); + await new Promise((resolve) => setImmediate(resolve)); + controller.abort(); + resolveBatch([{ processed: 1, inserted: 1, updated: 0, max_id: 1 }]); + + await expect(backfill).rejects.toMatchObject({ name: "AbortError" }); + expect(execute).toHaveBeenCalledTimes(2); + }); });