From 3ac77b382f47c3beee06778cbbffe83b0df20593 Mon Sep 17 00:00:00 2001 From: ding113 Date: Tue, 14 Jul 2026 19:24:50 +0800 Subject: [PATCH 01/20] feat(db): split pool into data/control/writer lanes with bounded admission control Replace the single postgres.js pool with three isolated lanes (data, control, writer) that share a configurable total connection budget. Each lane gets a per-pool outstanding-operation admission wrapper that fast-fails with DB_POOL_ADMISSION_EXCEEDED before saturating the underlying driver queue, preventing unbounded memory growth under load. Route handlers now wrap with withDataDbScope so request-path queries use the data lane via AsyncLocalStorage, while control-plane and writer traffic remain isolated. Add statement_timeout and lock_timeout at the connection level so slow SQL cannot block the streaming finalization deadline. Introduce LOCAL_OVERLOAD error category so admission rejections are never retried, failovered, or counted against Provider/endpoint circuits. --- .env.example | 13 +- deploy/k8s/app/deployment.yaml | 2 +- src/app/v1/[...route]/route.ts | 19 +- src/app/v1/_lib/proxy/errors.ts | 10 +- src/app/v1/_lib/proxy/forwarder.ts | 68 +++++ src/app/v1beta/[...route]/route.ts | 19 +- src/drizzle/admitted-client.ts | 260 ++++++++++++++++ src/drizzle/db.ts | 164 ++++++++-- src/lib/config/env.schema.ts | 25 +- tests/configs/integration.config.ts | 4 + .../db-pool-isolation-postgres.test.ts | 126 ++++++++ .../db-pool-slow-close-postgres.test.ts | 282 ++++++++++++++++++ tests/unit/drizzle/db-admission.test.ts | 164 ++++++++++ tests/unit/drizzle/db-pool-config.test.ts | 268 +++++++++++++++-- tests/unit/drizzle/db-scope.test.ts | 93 ++++++ tests/unit/drizzle/db-shutdown.test.ts | 282 ++++++++++++++++++ .../client-abort-vs-upstream-499.test.ts | 12 + .../proxy-forwarder-hedge-first-byte.test.ts | 43 +++ .../proxy/proxy-forwarder-retry-limit.test.ts | 46 +++ 19 files changed, 1826 insertions(+), 74 deletions(-) create mode 100644 src/drizzle/admitted-client.ts create mode 100644 tests/integration/db-pool-isolation-postgres.test.ts create mode 100644 tests/integration/db-pool-slow-close-postgres.test.ts create mode 100644 tests/unit/drizzle/db-admission.test.ts create mode 100644 tests/unit/drizzle/db-scope.test.ts create mode 100644 tests/unit/drizzle/db-shutdown.test.ts 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/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/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/errors.ts b/src/app/v1/_lib/proxy/errors.ts index 91e220263..1731b63df 100644 --- a/src/app/v1/_lib/proxy/errors.ts +++ b/src/app/v1/_lib/proxy/errors.ts @@ -6,6 +6,7 @@ * 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"; @@ -554,6 +555,7 @@ export enum ErrorCategory { CLIENT_ABORT, // 客户端主动中断 → 不计入熔断器 + 不重试 + 直接返回 NON_RETRYABLE_CLIENT_ERROR, // 客户端输入错误（Prompt 超限、内容过滤、PDF 限制、Thinking 格式、参数缺失/额外参数、非法请求）→ 不计入熔断器 + 不重试 + 直接返回 RESOURCE_NOT_FOUND, // 上游 404 错误 → 不计入熔断器 + 直接切换供应商 + LOCAL_OVERLOAD, // 本地数据库等 admission 过载 → 不计入任何熔断器 + 不重试/切换供应商 } /** @@ -925,7 +927,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 +935,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)) { diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 4da548a5f..40f0b5ef8 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 } from "@/drizzle/admitted-client"; import { applyAnthropicProviderOverridesWithAudit } from "@/lib/anthropic/provider-overrides"; import { getCircuitState, @@ -1783,6 +1784,37 @@ export class ProxyForwarder { throw lastError; } + if (errorCategory === ErrorCategory.LOCAL_OVERLOAD) { + const admission = findDbPoolAdmissionError(lastError); + logger.warn("ProxyForwarder: Local database admission rejected", { + 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, + errorDetails: { + system: { + errorType: "DbPoolAdmissionError", + errorName: "DbPoolAdmissionError", + errorMessage: admission?.message ?? errorMessage, + errorCode: admission?.code, + }, + request: buildRequestDetails(session), + }, + }); + + throw lastError; + } + // 2.5 Reactive rectifier：命中后对同供应商“整流 + 重试一次” const reactiveRectifierResult = await tryApplyReactiveRectifier({ provider: currentProvider, @@ -3852,6 +3884,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 +3956,7 @@ export class ProxyForwarder { await finalizeHedgeLoserBilling({ messageRequestId, + messageRequestCreatedAtMs: messageRequestCreatedAtMs ?? Date.now(), loserSession: attempt.session, provider: attempt.provider, attemptNumber: attempt.sequence, @@ -4263,6 +4297,40 @@ export class ProxyForwarder { return; } + if (errorCategory === ErrorCategory.LOCAL_OVERLOAD) { + const admission = findDbPoolAdmissionError(error); + 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, + circuitState: getCircuitState(attempt.provider.id), + errorDetails: { + system: { + errorType: "DbPoolAdmissionError", + errorName: "DbPoolAdmissionError", + errorMessage: admission?.message ?? errorMessage, + 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, 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..b55eaa98d --- /dev/null +++ b/src/drizzle/admitted-client.ts @@ -0,0 +1,260 @@ +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 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; +} + +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 (result === target) { + if (property === "execute" || property === "forEach" || property === "cancel") { + 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, { + 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/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/tests/configs/integration.config.ts b/tests/configs/integration.config.ts index 064a31fab..513a74658 100644 --- a/tests/configs/integration.config.ts +++ b/tests/configs/integration.config.ts @@ -8,6 +8,10 @@ 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", ], api: { host: process.env.VITEST_API_HOST || "127.0.0.1", 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..821527cf1 --- /dev/null +++ b/tests/integration/db-pool-isolation-postgres.test.ts @@ -0,0 +1,126 @@ +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); + +if (!process.env.DSN && process.env.DATABASE_URL) { + process.env.DSN = process.env.DATABASE_URL; +} + +const previousPoolMax = process.env.DB_POOL_MAX; +process.env.DB_POOL_MAX = "6"; +vi.resetModules(); + +run("PostgreSQL pool isolation integration", () => { + let dbModule: typeof import("@/drizzle/db"); + + 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 () => { + dbModule = await import("@/drizzle/db"); + }); + + afterAll(async () => { + await dbModule.closeDbPools(); + if (previousPoolMax === undefined) { + delete process.env.DB_POOL_MAX; + } else { + process.env.DB_POOL_MAX = previousPoolMax; + } + }); + + 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/unit/drizzle/db-admission.test.ts b/tests/unit/drizzle/db-admission.test.ts new file mode 100644 index 000000000..f3fae9b24 --- /dev/null +++ b/tests/unit/drizzle/db-admission.test.ts @@ -0,0 +1,164 @@ +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; + }); +}); 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/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/proxy-forwarder-hedge-first-byte.test.ts b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts index 9a6a0bc3a..c4e83c99f 100644 --- a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts +++ b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts @@ -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"); + 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 b0f27dc6d..454b3a7d6 100644 --- a/tests/unit/proxy/proxy-forwarder-retry-limit.test.ts +++ b/tests/unit/proxy/proxy-forwarder-retry-limit.test.ts @@ -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(); From 3673105a484dc099925e4b8c6995a142c734983b Mon Sep 17 00:00:00 2001 From: ding113 Date: Tue, 14 Jul 2026 19:24:50 +0800 Subject: [PATCH 02/20] feat(proxy): add demand-driven response pump with bounded backpressure Introduce createDemandDrivenResponsePump to replace the TransformStream + tee() pipeline with a single high-water-mark-0 ReadableStream that primes exactly one upstream chunk and never reads ahead past unconsumed downstream demand. This eliminates unbounded buffering for slow clients and removes the generic streaming stale watchdog. Rewrite nodeStreamToWebStreamSafe to pause/resume the Node source on controller.desiredSize backpressure, handle pre-existing errored/destroyed/closed states before listener registration, and guard delayed asynchronous destroy(error) events so they cannot surface as uncaught exceptions after the Web stream has settled. Defer all circuit-breaker, session-binding, and Codex cache side effects into post-terminal tasks so a slow Redis call cannot turn an already-completed billable request into a fallback 500 or block lease settlement. --- .../proxy/demand-driven-response-pump.test.ts | 406 +++ .../_lib/proxy/demand-driven-response-pump.ts | 235 ++ .../v1/_lib/proxy/node-stream-to-web.test.ts | 235 +- src/app/v1/_lib/proxy/node-stream-to-web.ts | 249 +- src/app/v1/_lib/proxy/response-handler.ts | 1796 ++++++++---- .../integration/billing-model-source.test.ts | 259 +- tests/unit/proxy/pricing-no-price.test.ts | 9 +- ...nse-handler-abort-listener-cleanup.test.ts | 28 +- .../response-handler-bill-non-success.test.ts | 13 +- ...esponse-handler-client-abort-drain.test.ts | 2587 +++++++++++++++-- ...handler-endpoint-circuit-isolation.test.ts | 97 +- ...gemini-stream-passthrough-timeouts.test.ts | 143 +- ...ponse-handler-hedge-loser-priority.test.ts | 71 +- .../response-handler-lease-decrement.test.ts | 351 ++- .../proxy/response-handler-non200.test.ts | 18 +- 15 files changed, 5352 insertions(+), 1145 deletions(-) create mode 100644 src/app/v1/_lib/proxy/demand-driven-response-pump.test.ts create mode 100644 src/app/v1/_lib/proxy/demand-driven-response-pump.ts 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..438a732f9 --- /dev/null +++ b/src/app/v1/_lib/proxy/demand-driven-response-pump.ts @@ -0,0 +1,235 @@ +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; +} + +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 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 settle = (streamEndedNormally: boolean, error: Error | null) => { + if (settled) return; + settled = true; + state = "finalizing"; + pendingChunk = null; + releaseReader(); + clientController = null; + state = "closed"; + resolveCompletion({ streamEndedNormally, clientAborted, error }); + }; + + 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 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); + pendingChunk = result.value; + }, + (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) { + pendingChunk = null; + continue; + } + await ensureRead(); + } + })().finally(() => { + drainPromise = null; + if (!settled && state === "draining") { + scheduleDrain(); + } + }); + }; + + const startDrain = (_reason?: unknown) => { + if (settled || state === "finalizing" || state === "closed") return; + if (state === "draining") { + scheduleDrain(); + return; + } + clientAborted = true; + state = "draining"; + try { + clientController?.error( + _reason == null ? new Error("Client disconnected") : toError(_reason) + ); + } catch { + // The ReadableStream cancel algorithm may have already detached the controller. + } + 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; + pendingChunk = null; + try { + clientController?.enqueue(chunk); + } catch (error) { + finishWithError(error); + return; + } + + void ensureRead(); + }, + cancel(reason) { + options.onClientCancel?.(reason); + startDrain(reason); + }, + }, + { highWaterMark: 0 } + ); + + void ensureRead(); + + return { + stream, + completion, + startDrain, + cancelSource(reason) { + if (settled) return; + const normalized = reason == null ? new Error("Source cancelled") : toError(reason); + const cancelPromise = reader.cancel(normalized); + settle(false, normalized); + void cancelPromise.catch(() => { + // The pump has already recorded the hard-cancel cause. + }); + }, + 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/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/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index 1a8406b6d..f4114683b 100644 --- a/src/app/v1/_lib/proxy/response-handler.ts +++ b/src/app/v1/_lib/proxy/response-handler.ts @@ -38,6 +38,8 @@ import { addMessageRequestHedgeLoserCost, updateMessageRequestCostWithBreakdown, updateMessageRequestDetails, + updateMessageRequestDetailsDurably, + updateMessageRequestDetailsIfUnfinalized, updateMessageRequestDuration, updateMessageRequestWinnerCost, } from "@/repository/message"; @@ -49,6 +51,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 +87,150 @@ 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; +}): Promise { + const completeTerminalDetails = { + ...options.details, + durationMs: options.durationMs, + }; + try { + await updateMessageRequestDetailsDurably(options.messageRequestId, completeTerminalDetails); + return; + } catch (primaryError) { + logger.error("ResponseHandler: Durable non-stream terminal persistence failed", { + taskId: options.taskId, + messageId: options.messageRequestId, + statusCode: options.details.statusCode, + error: primaryError, + }); } - if (provider.firstByteTimeoutStreamingMs > 0) { - return Math.max(provider.firstByteTimeoutStreamingMs, provider.streamingIdleTimeoutMs); + try { + await updateMessageRequestDetailsIfUnfinalized( + options.messageRequestId, + completeTerminalDetails + ); + } catch (fallbackError) { + logger.error("ResponseHandler: Conditional non-stream terminal fallback failed", { + taskId: options.taskId, + messageId: options.messageRequestId, + statusCode: options.details.statusCode, + fallbackError, + }); + 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); +} + +function schedulePostTerminalSideEffects(options: { + taskId: string; + providerId: number; + sessionId: string | null; + commit: () => Promise; +}): void { + const effectTaskId = `${options.taskId}-post-terminal-effects`; + AsyncTaskManager.register( + effectTaskId, + async () => { + let commitPromise: Promise; + try { + commitPromise = Promise.resolve(options.commit()); + } catch (error) { + commitPromise = Promise.reject(error); + } + + await raceWithTimeout( + commitPromise, + STREAM_FINALIZATION_MAX_MS, + "post_terminal_side_effect_timeout" + ).catch((error) => { + logger.warn("[ResponseHandler] Post-terminal side effects did not complete", { + taskId: options.taskId, + providerId: options.providerId, + sessionId: options.sessionId, + error: error instanceof Error ? error.message : String(error), + }); + }); + }, + { + taskType: "post-terminal-side-effects", + staleTimeoutMs: STREAM_FINALIZATION_MAX_MS, + } + ); } // 流式统计只需要头部元信息和尾部 usage/final event。按字节保存窗口，避免 @@ -866,6 +1006,8 @@ type FinalizeDeferredStreamingResult = { usageMetrics: UsageMetrics | null; providerType: Provider["providerType"] | undefined; }; + /** Circuit and Session side effects, committed after durable terminal details. */ + commitSideEffects?: () => Promise; }; /** @@ -885,14 +1027,14 @@ 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 clearSessionBinding = async () => { @@ -987,10 +1129,6 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( detected.isError || (upstreamStatusCode >= 400 && errorMessage !== null); - if ((!meta || !provider) && shouldClearSessionBindingOnFailure) { - await clearSessionBinding(); - } - // 未启用延迟结算 / provider 缺失： // - 只返回“内部状态码 + 错误原因”，由调用方写入统计； // - 不在这里更新熔断/绑定（meta 缺失意味着 Forwarder 没有启用延迟结算；provider 缺失意味着无法归因）。 @@ -1002,6 +1140,7 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( isHedgeWinner, billHedgeLosers, clientAbortGateUsage, + commitSideEffects: shouldClearSessionBindingOnFailure ? clearSessionBinding : undefined, }; } @@ -1019,27 +1158,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。 @@ -1048,27 +1174,6 @@ 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. - } - session.addProviderToChain(providerForChain, { endpointId: meta.endpointId, endpointUrl: meta.endpointUrl, @@ -1078,6 +1183,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 +1210,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 +1228,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 +1246,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 +1274,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 +1290,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 +1304,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 +1331,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 +1407,6 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( }); } - // 统一更新两个数据源（确保监控数据一致） if (session.shouldTrackSessionObservability()) { void SessionManager.updateSessionProvider(session.sessionId, { providerId: meta.providerId, @@ -1270,30 +1414,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 +1436,7 @@ async function finalizeDeferredStreamingFinalizationIfNeeded( isHedgeWinner, billHedgeLosers, clientAbortGateUsage, + commitSideEffects, }; } @@ -1412,7 +1547,7 @@ export class ProxyResponseHandler { statsAbortController, taskId ); - const statsPromise = (async () => { + const runStatsTask = async () => { try { const responseText = await readResponseTextWithTaskActivity(responseForStats, taskId); @@ -1457,29 +1592,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 +1629,9 @@ export class ProxyResponseHandler { // 使用共享的统计处理方法 const duration = Date.now() - session.startTime; + if (messageContext) { + await updateMessageRequestDuration(messageContext.id, duration); + } const finalizedUsage = await finalizeRequestStats( session, responseText, @@ -1498,6 +1642,15 @@ export class ProxyResponseHandler { false // Gemini 非流式透传 ); + if (commitProviderFailure) { + schedulePostTerminalSideEffects({ + taskId, + providerId: provider.id, + sessionId: session.sessionId, + commit: commitProviderFailure, + }); + } + emitProxyLangfuseTrace(session, { responseHeaders: response.headers, responseText, @@ -1512,32 +1665,111 @@ 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, + }); + } + + if (messageContext) { + const duration = Date.now() - session.startTime; + await updateMessageRequestDuration(messageContext.id, duration); + await updateMessageRequestDetailsDurably(messageContext.id, { + 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, + }); + const tracker = ProxyStatusTracker.getInstance(); + tracker.endRequest(messageContext.user.id, messageContext.id); + } + + const postTerminalSideEffects: Array<() => Promise> = []; + if (session.sessionId) { + const sessionId = session.sessionId; + postTerminalSideEffects.push(async () => { + await SessionManager.clearSessionProvider(sessionId); + }); + } + 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, + } + ); + } + }); + } + if (postTerminalSideEffects.length > 0) { + schedulePostTerminalSideEffects({ + taskId, + providerId: provider.id, + sessionId: session.sessionId, + commit: async () => { + await Promise.all(postTerminalSideEffects.map((effect) => effect())); + }, + }); + } } 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,40 +1837,77 @@ export class ProxyResponseHandler { abortController.abort(); }); - const processingPromise = (async () => { - const finalizeNonStreamAbort = async (): Promise => { - const finalizedStatusCode = session.clientAbortSignal?.aborted ? 499 : statusCode; + 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); if (messageContext) { const duration = Date.now() - session.startTime; - await updateMessageRequestDuration(messageContext.id, duration); - await updateMessageRequestDetails(messageContext.id, { + 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(); - tracker.endRequest(messageContext.user.id, messageContext.id); + try { + await persistNonStreamTerminalDetails({ + taskId, + messageRequestId: messageContext.id, + durationMs: duration, + details: terminalDetails, + }); + } finally { + tracker.endRequest(messageContext.user.id, messageContext.id); + } } + 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); + + 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); } - ); - } + } + }); + } + + if (postTerminalSideEffects.length > 0) { + schedulePostTerminalSideEffects({ + taskId, + providerId: provider.id, + sessionId: session.sessionId, + commit: async () => { + await Promise.all(postTerminalSideEffects.map((effect) => effect())); + }, + }); } }; @@ -1657,6 +1926,9 @@ export class ProxyResponseHandler { providerId: provider.id, finalizeError, }); + if (isNonStreamTerminalPersistenceError(finalizeError)) { + throw finalizeError; + } } return; } @@ -1672,6 +1944,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 +1982,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,27 +2148,31 @@ 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, @@ -1892,10 +2181,7 @@ export class ProxyResponseHandler { 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 +2192,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 +2203,29 @@ 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, + }); + } finally { + tracker.endRequest(messageContext.user.id, messageContext.id); + } + } + + if (postTerminalSideEffects.length > 0) { + schedulePostTerminalSideEffects({ + taskId, + providerId: provider.id, + sessionId: session.sessionId, + commit: async () => { + await Promise.all(postTerminalSideEffects.map((effect) => effect())); + }, + }); } logger.debug("ResponseHandler: Non-stream response processed", { @@ -1941,6 +2246,9 @@ export class ProxyResponseHandler { isStreaming: false, }); } catch (error) { + if (isNonStreamTerminalPersistenceError(error)) { + throw error; + } if (session.sessionId && session.shouldPersistSessionDebugArtifacts()) { await discardBeforeResponseBodySnapshot(session); } @@ -1958,7 +2266,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 +2273,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 +2336,9 @@ export class ProxyResponseHandler { providerId: provider.id, finalizeError, }); + if (isNonStreamTerminalPersistenceError(finalizeError)) { + throw finalizeError; + } } } } else { @@ -2045,30 +2359,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; } @@ -2114,14 +2439,14 @@ export class ProxyResponseHandler { 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; }; @@ -2138,6 +2463,7 @@ export class ProxyResponseHandler { let streamEndedNormally = false; let responseTimeoutCleared = false; let abortReason: string | undefined; + let transportReleased = false; // 静默期 Watchdog：透传也需要支持中途卡住（无新数据推送） const idleTimeoutMs = @@ -2203,6 +2529,66 @@ export class ProxyResponseHandler { return flushAndSnapshot().text; }; + const releaseTransportResources = () => { + if (transportReleased) return; + transportReleased = true; + 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), + } + ); + } + try { + const cancelPromise = reader?.cancel(); + cancelPromise?.catch((error) => { + logger.warn("[ResponseHandler] Gemini passthrough: Failed to cancel stats reader", { + taskId, + providerId: provider.id, + providerName: provider.name, + error: error instanceof Error ? error.message : String(error), + }); + }); + } catch (error) { + logger.warn("[ResponseHandler] Gemini passthrough: Failed to cancel stats reader", { + taskId, + providerId: provider.id, + providerName: provider.name, + error: error instanceof Error ? error.message : String(error), + }); + } + try { + reader?.releaseLock(); + } catch (error) { + logger.warn("[ResponseHandler] Gemini passthrough: Failed to release reader lock", { + taskId, + providerId: provider.id, + providerName: provider.name, + error: error instanceof Error ? error.message : String(error), + }); + } + reader = null; + releaseSessionAgent(session); + }; + try { const body = responseForStats.body; if (!body) return; @@ -2255,6 +2641,7 @@ export class ProxyResponseHandler { const streamSnapshot = flushAndSnapshot(); const allContent = streamSnapshot.text; const clientAborted = session.clientAbortSignal?.aborted ?? false; + releaseTransportResources(); // 存储响应体到 Redis（5分钟过期） if ( @@ -2319,7 +2706,6 @@ export class ProxyResponseHandler { finalized.providerIdForPersistence ?? undefined, true // Gemini 流式透传(NDJSON 无 data:/event: 前缀,必须显式告知) ); - emitProxyLangfuseTrace(session, { responseHeaders: response.headers, responseText: allContent, @@ -2330,6 +2716,14 @@ export class ProxyResponseHandler { isStreaming: true, errorMessage: finalized.errorMessage ?? undefined, }); + if (finalized.commitSideEffects) { + schedulePostTerminalSideEffects({ + taskId, + providerId: provider.id, + sessionId: session.sessionId, + commit: finalized.commitSideEffects, + }); + } } catch (error) { const err = error instanceof Error ? error : new Error(String(error)); const clientAborted = session.clientAbortSignal?.aborted ?? false; @@ -2384,6 +2778,14 @@ export class ProxyResponseHandler { finalized.providerIdForPersistence ?? undefined, true // 流式透传错误兜底也是流式上下文 ); + if (finalized.commitSideEffects) { + schedulePostTerminalSideEffects({ + taskId, + providerId: provider.id, + sessionId: session.sessionId, + commit: finalized.commitSideEffects, + }); + } } catch (finalizeError) { await persistRequestFailure({ session, @@ -2395,83 +2797,31 @@ export class ProxyResponseHandler { }); } } 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?.( @@ -2544,21 +2894,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 +2908,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 +2944,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 +2986,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 +3026,97 @@ export class ProxyResponseHandler { }); } - abortController.abort(new Error("client_abort_drain_timeout")); + const drainTimeoutError = new Error("client_abort_drain_timeout"); + abortController.abort(drainTimeoutError); + responsePump?.cancelSource(drainTimeoutError); }, clientAbortDrainTimeoutMs); - }); + }; - 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; + // 统计/结算只保留有界的“头 + 尾”文本快照，避免长流式响应把进程堆撑满。 + 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 finalizeStream = async ( - allContent: string, - streamEndedNormally: boolean, - clientAborted: boolean, - abortReason?: string - ): Promise => { - const finalized = await finalizeDeferredStreamingFinalizationIfNeeded( + 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 error = new Error( + signal.reason instanceof Error ? signal.reason.message : "Response timeout" + ); + error.name = "AbortError"; + return error; + }; + + let terminalDetailsPersisted = false; + let streamFailurePersistencePromise: Promise | null = null; + const persistStreamFailureOnce = ( + options: Parameters[0] + ): Promise => { + if (terminalDetailsPersisted) return Promise.resolve(); + if (!streamFailurePersistencePromise) { + const persistenceDeadlineAtMs = Date.now() + STREAM_FAILURE_PERSISTENCE_MAX_MS; + streamFailurePersistencePromise = persistRequestFailure({ + ...options, + detailsWriter: updateMessageRequestDetailsIfUnfinalized, + awaitPersistence: (promise: Promise) => + raceWithDeadline( + promise, + persistenceDeadlineAtMs, + "stream_failure_persistence_timeout" + ), + }).catch((error) => { + logger.error("ResponseHandler: Stream failure fallback threw", { + taskId, + messageId: messageContext.id, + error, + }); + }); + } + return streamFailurePersistencePromise; + }; + + 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, @@ -2729,7 +3136,8 @@ export class ProxyResponseHandler { session.shouldPersistSessionDebugArtifacts() && !streamSnapshot?.truncated ) { - const beforeBody = (await consumeBeforeResponseBodySnapshot(session)) ?? allContent; + const beforeBody = + (await awaitFinalization(consumeBeforeResponseBodySnapshot(session))) ?? allContent; void SessionManager.storeSessionResponse( session.sessionId, allContent, @@ -2770,7 +3178,7 @@ export class ProxyResponseHandler { } const duration = Date.now() - session.startTime; - await updateMessageRequestDuration(messageContext.id, duration); + await awaitFinalization(updateMessageRequestDuration(messageContext.id, duration)); const tracker = ProxyStatusTracker.getInstance(); tracker.endRequest(messageContext.user.id, messageContext.id); @@ -2783,9 +3191,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 +3209,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 +3232,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 +3247,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 +3291,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 = @@ -2981,25 +3407,59 @@ export class ProxyResponseHandler { : extractActualResponseModelForProvider(provider.providerType, true, allContent); // 保存扩展信息（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, - }); + await awaitFinalization( + updateMessageRequestDetailsDurably(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 = true; + + const postTerminalSideEffects: Array<() => Promise> = []; + const commitDeferredSideEffects = finalized.commitSideEffects; + if (commitDeferredSideEffects) { + postTerminalSideEffects.push(commitDeferredSideEffects); + } + 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); + } + }); + } + if (postTerminalSideEffects.length > 0) { + schedulePostTerminalSideEffects({ + taskId, + providerId: providerIdForPersistence ?? provider.id, + sessionId: session.sessionId, + commit: async () => { + await Promise.all(postTerminalSideEffects.map((effect) => effect())); + }, + }); + } emitProxyLangfuseTrace(session, { responseHeaders: response.headers, @@ -3013,65 +3473,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); - - // 每次收到数据后重置静默期计时器（首次收到数据时启动） - startIdleTimer(); - logger.trace("ResponseHandler: Idle timer reset (data received)", { - taskId, - providerId: provider.id, - chunksCollected: getCollectedChunkCount(), - lastChunkSize: chunkSize, - idleTimeoutMs: idleTimeoutMs === Infinity ? "disabled" : idleTimeoutMs, - }); + logger.trace("ResponseHandler: Upstream stream chunk received", { + taskId, + providerId: provider.id, + chunksCollected: getCollectedChunkCount(), + lastChunkSize: chunkSize, + idleTimeoutMs: idleTimeoutMs === Infinity ? "disabled" : idleTimeoutMs, + }); - // 流式：读到第一块数据后立即清除响应超时定时器 - 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, - }); - } - } - } + if (isFirstChunk) { + session.recordTtfb(); + isFirstChunk = false; + if (clearResponseTimeoutOnce()) { + 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 +3574,7 @@ export class ProxyResponseHandler { }); // 回退：避免 finalizeStream 失败导致 request record 未被更新 - await persistRequestFailure({ + await persistStreamFailureOnce({ session, messageContext, statusCode: statusCode && statusCode >= 400 ? statusCode : 500, @@ -3101,7 +3589,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 +3627,7 @@ export class ProxyResponseHandler { }); // 回退：至少保证 DB 记录能落下，避免 orphan record - await persistRequestFailure({ + await persistStreamFailureOnce({ session, messageContext, statusCode: statusCode && statusCode >= 400 ? statusCode : 502, @@ -3175,7 +3666,7 @@ export class ProxyResponseHandler { }); // 回退：至少保证 DB 记录能落下，避免 orphan record - await persistRequestFailure({ + await persistStreamFailureOnce({ session, messageContext, statusCode: statusCode && statusCode >= 400 ? statusCode : 502, @@ -3208,7 +3699,7 @@ export class ProxyResponseHandler { }); // 回退：至少保证 DB 记录能落下，避免 orphan record - await persistRequestFailure({ + await persistStreamFailureOnce({ session, messageContext, statusCode: 502, @@ -3240,9 +3731,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 +3797,7 @@ export class ProxyResponseHandler { finalizeError, }); - await persistRequestFailure({ + await persistStreamFailureOnce({ session, messageContext, statusCode: 502, @@ -3290,7 +3821,7 @@ export class ProxyResponseHandler { }); // 回退：至少保证 DB 记录能落下，避免 orphan record - await persistRequestFailure({ + await persistStreamFailureOnce({ session, messageContext, statusCode: statusCode && statusCode >= 400 ? statusCode : 500, @@ -3303,44 +3834,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 +3894,7 @@ export class ProxyResponseHandler { }); } - return new Response(clientStream, { + return new Response(activeResponsePump.stream, { status: response.status, statusText: response.statusText, headers: finalStreamHeaders, @@ -4161,6 +4692,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 +4726,7 @@ export async function finalizeHedgeLoserBilling(params: { }): Promise { const { messageRequestId, + messageRequestCreatedAtMs, loserSession, provider, attemptNumber, @@ -4310,7 +4844,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", { @@ -4427,7 +4965,7 @@ export async function finalizeRequestStats( }); } - await updateMessageRequestDetails(messageContext.id, { + await updateMessageRequestDetailsDurably(messageContext.id, { statusCode: statusCode, ...(errorMessage ? { errorMessage } : {}), ttfbMs: session.ttfbMs ?? duration, @@ -4530,7 +5068,7 @@ export async function finalizeRequestStats( } // 7. 更新请求详情 - await updateMessageRequestDetails(messageContext.id, { + await updateMessageRequestDetailsDurably(messageContext.id, { statusCode: statusCode, inputTokens: normalizedUsage.input_tokens, outputTokens: normalizedUsage.output_tokens, @@ -4579,6 +5117,11 @@ type BillingComputeInputs = { groupCostMultiplier: number; }; +type CostTrackingEventContext = { + eventId: string | number; + createdAtMs: number; +}; + function sessionBillingInputs( session: ProxySession, provider: Provider, @@ -4600,9 +5143,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 +5155,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 +5192,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 +5241,36 @@ async function trackCostToRedis( } } +function buildProcessingErrorDetails(error: unknown): { + errorMessage: string; + errorStack?: string; + errorCause?: string; +} { + 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,8 +5283,12 @@ async function persistRequestFailure(options: { error: unknown; taskId: string; phase: "stream" | "non-stream"; + awaitPersistence?: (promise: Promise) => Promise; + detailsWriter?: typeof updateMessageRequestDetails; }): Promise { const { session, messageContext, statusCode, error, taskId, phase } = options; + const awaitPersistence = options.awaitPersistence ?? ((promise: Promise) => promise); + const detailsWriter = options.detailsWriter ?? updateMessageRequestDetails; if (!messageContext) { logger.warn("ResponseHandler: Cannot persist failure without messageContext", { @@ -4746,52 +5299,29 @@ async function persistRequestFailure(options: { } 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]`; - } - } - try { // 更新请求持续时间 - await updateMessageRequestDuration(messageContext.id, duration); + await awaitPersistence(updateMessageRequestDuration(messageContext.id, duration)); // 更新错误详情和 provider chain - await updateMessageRequestDetails(messageContext.id, { - statusCode, - errorMessage, - errorStack, - errorCause, - ttfbMs: phase === "non-stream" ? (session.ttfbMs ?? duration) : session.ttfbMs, - providerChain: session.getProviderChain(), - model: session.getCurrentModel() ?? undefined, - providerId: session.provider?.id, // 更新最终供应商ID（重试切换后） - context1mApplied: session.getContext1mApplied(), - swapCacheTtlApplied: session.provider?.swapCacheTtlBilling ?? false, - specialSettings: session.getSpecialSettings() ?? undefined, - }); + await awaitPersistence( + detailsWriter(messageContext.id, { + statusCode, + errorMessage, + errorStack, + errorCause, + ttfbMs: phase === "non-stream" ? (session.ttfbMs ?? duration) : session.ttfbMs, + providerChain: session.getProviderChain(), + model: session.getCurrentModel() ?? undefined, + providerId: session.provider?.id, // 更新最终供应商ID（重试切换后） + context1mApplied: session.getContext1mApplied(), + swapCacheTtlApplied: session.provider?.swapCacheTtlBilling ?? false, + specialSettings: session.getSpecialSettings() ?? undefined, + }) + ); if (session.sessionId && session.requestSequence != null) { if (session.shouldTrackSessionObservability()) { 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/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/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..53391b2cd 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,20 +882,228 @@ 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 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(); }); + 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(updateMessageRequestDuration).mockImplementationOnce(async () => { + recoveryStarted.resolve(); + await releaseRecovery.promise; + }); + 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(); + + releaseRecovery.resolve(); + await shutdownPromise; + await expect(processingTask).rejects.toThrow("error classification failed"); + + expect(updateMessageRequestDetails).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + 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(updateMessageRequestDuration).mockImplementationOnce(async () => { + recoveryStarted.resolve(); + await releaseRecovery.promise; + }); + 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).not.toHaveBeenCalled(); + + releaseRecovery.resolve(); + await shutdownPromise; + await downstreamRead; + await expect(processingTask).rejects.toThrow("error classification failed"); + + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 500, + errorMessage: "Error: error classification failed", + }) + ); + } finally { + releaseRecovery.resolve(); + await shutdownPromise?.catch(() => {}); + await downstreamRead?.catch(() => {}); + await drainAsyncTasks(); + } + }); + it("copies Buffer-backed stream windows before retaining stats snapshots", () => { const accumulator = new BoundedStreamTextAccumulator(); const headMarker = "head-copy-marker"; @@ -681,10 +1124,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 +1139,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 +1167,100 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); + const tracked = createPullTrackedResponsesSse(); - await ProxyResponseHandler.dispatch(session, createResponsesSse()); - await drainAsyncTasks(); - - expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); - expect(updateMessageRequestDuration).toHaveBeenCalledWith(123, expect.any(Number)); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( - 123, - expect.objectContaining({ - statusCode: 200, - inputTokens: 463, - outputTokens: 11, - }) - ); - }); + 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)); - 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, - }); + expect(tracked.getPullCount()).toBeLessThanOrEqual(3); + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); - await ProxyResponseHandler.dispatch(session, createOversizedResponsesSse()); + await reader?.cancel("test cleanup"); await drainAsyncTasks(); + expect(updateMessageRequestDuration).toHaveBeenCalledWith(123, expect.any(Number)); + }); - 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(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 +1274,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 +1309,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 +1339,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 +1362,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).toHaveBeenCalledTimes(1); + 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 +1396,59 @@ 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(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 +1461,66 @@ 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( + 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; + }); + 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 +1533,58 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); + vi.mocked(updateMessageRequestDuration).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(updateMessageRequestDuration).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( - 123, - expect.objectContaining({ - statusCode: 499, - errorMessage: "CLIENT_ABORTED", - }) - ); - } finally { + 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).toHaveBeenCalledTimes(1); + } finally { + 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 +1597,40 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); - - await ProxyResponseHandler.dispatch( - session, - createPreBodyHangingResponsesSse(upstreamController.signal) + vi.mocked(updateMessageRequestDuration).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", + statusCode: 500, + errorMessage: "Error: stream_finalization_timeout", }) ); } finally { + vi.clearAllTimers(); vi.useRealTimers(); } }); - it("preserves an existing idle deadline when the client aborts after a chunk", async () => { + it("does not continue a timed-out failure fallback in the background", async () => { vi.useFakeTimers(); + const flushMicrotasks = (remaining = 32): Promise => + remaining === 0 + ? Promise.resolve() + : Promise.resolve().then(() => flushMicrotasks(remaining - 1)); + let resolveFallbackDuration!: () => void; 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 +1643,1014 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); - - await ProxyResponseHandler.dispatch( - session, - createHangingResponsesSse(upstreamController.signal) - ); - await vi.advanceTimersByTimeAsync(0); - await vi.advanceTimersByTimeAsync(4_999); - expect(upstreamController.signal.aborted).toBe(false); - - clientController.abort(); - await vi.advanceTimersByTimeAsync(1); + vi.mocked(updateMessageRequestDuration) + .mockImplementationOnce(() => new Promise(() => {})) + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveFallbackDuration = resolve; + }) + ); + + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); + await vi.advanceTimersByTimeAsync(120_000); + await flushMicrotasks(); + + expect(updateMessageRequestDuration).toHaveBeenCalledTimes(2); + + await vi.advanceTimersByTimeAsync(5_000); + await flushMicrotasks(); const tasks = asyncTasks.splice(0, asyncTasks.length); - await Promise.allSettled(tasks); + await expectAllFulfilled(tasks); + expect(updateMessageRequestDetailsDurably).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsIfUnfinalized).not.toHaveBeenCalled(); - expect(upstreamController.signal.aborted).toBe(true); - expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( - 123, - expect.objectContaining({ - statusCode: 499, - errorMessage: "CLIENT_ABORTED", - }) - ); + resolveFallbackDuration(); + await flushMicrotasks(); + + expect(updateMessageRequestDetailsDurably).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsIfUnfinalized).not.toHaveBeenCalled(); } finally { + resolveFallbackDuration?.(); + 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 () => { + it("shares one deadline across sequential stream finalization steps", async () => { + vi.useFakeTimers(); + const flushMicrotasks = (remaining = 32): Promise => + remaining === 0 + ? Promise.resolve() + : Promise.resolve().then(() => flushMicrotasks(remaining - 1)); + let resolveDuration!: () => void; + let resolveTerminalDetails!: () => void; + let markDurationStarted!: () => void; + const durationStarted = new Promise((resolve) => { + markDurationStarted = resolve; + }); + let markTerminalDetailsStarted!: () => void; + const terminalDetailsStarted = new Promise((resolve) => { + markTerminalDetailsStarted = 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(updateMessageRequestDuration).mockImplementationOnce(() => { + markDurationStarted(); + return new Promise((resolve) => { + resolveDuration = resolve; + }); + }); + vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce(() => { + markTerminalDetailsStarted(); + return new Promise((resolve) => { + resolveTerminalDetails = resolve; + }); + }); + + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); + await durationStarted; + await vi.advanceTimersByTimeAsync(119_000); + resolveDuration(); + await terminalDetailsStarted; + await flushMicrotasks(); + + expect(updateMessageRequestDuration).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDetailsIfUnfinalized).not.toHaveBeenCalled(); + + await vi.advanceTimersByTimeAsync(1_000); + await flushMicrotasks(); + + expect(updateMessageRequestDuration).toHaveBeenCalledTimes(2); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); + expect(vi.mocked(updateMessageRequestDetailsIfUnfinalized).mock.calls[0]).toEqual([ + 123, + expect.objectContaining({ + statusCode: 500, + errorMessage: "Error: stream_finalization_timeout", + }), + ]); + } finally { + resolveDuration?.(); + resolveTerminalDetails?.(); + await flushMicrotasks(); + vi.clearAllTimers(); + vi.useRealTimers(); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await Promise.race([ + expectAllFulfilled(tasks), + new Promise((resolve) => setTimeout(resolve, 100)), + ]); + } + }); + + 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 resolveFallbackDuration!: () => void; + let markFallbackDurationStarted!: () => void; + const fallbackDurationStarted = new Promise((resolve) => { + markFallbackDurationStarted = 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(updateMessageRequestDuration) + .mockImplementationOnce(async () => undefined) + .mockImplementationOnce(() => { + markFallbackDurationStarted(); + return new Promise((resolve) => { + resolveFallbackDuration = resolve; + }); + }); + vi.mocked(updateMessageRequestDetailsDurably) + .mockImplementationOnce((_id, details) => { + markPrimaryDetailsStarted(); + return new Promise((resolve) => { + resolvePrimaryDetails = () => { + terminalStatusCode = details.statusCode ?? null; + resolve(); + }; + }); + }) + .mockImplementation(async (_id, details) => { + terminalStatusCode = details.statusCode ?? null; + }); + vi.mocked(updateMessageRequestDetailsIfUnfinalized).mockImplementation( + async (_id, details) => { + if (terminalStatusCode === null) { + terminalStatusCode = details.statusCode ?? null; + } + } + ); + + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); + await primaryDetailsStarted; + await vi.advanceTimersByTimeAsync(120_000); + await fallbackDurationStarted; + + resolvePrimaryDetails(); + await flushMicrotasks(); + expect(terminalStatusCode).toBe(200); + + resolveFallbackDuration(); + await flushMicrotasks(); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); + expect(terminalStatusCode).toBe(200); + } finally { + resolvePrimaryDetails?.(); + resolveFallbackDuration?.(); + 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).toHaveBeenCalledWith(123, expect.any(Number)); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 200, + inputTokens: 463, + outputTokens: 11, + }) + ); + }); + + 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(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, + }) + ); + + 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(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, + }), + ], + }) + ); + }); + + 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", + }) + ); + }); + + 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", + }) + ); + // 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, + }) + ); + + 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", + }) + ); + } 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", + }) + ); + } 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", + }) + ); + } 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(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 }); + 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", + }) + ); + } 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", + }) + ); + } 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", + }) + ); + }); + + 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", + }) + ); + }); + + 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).toHaveBeenCalledTimes(2); + expect(vi.mocked(updateMessageRequestDuration).mock.calls[1]).toEqual([ + 123, + expect.any(Number), + ]); + expect(vi.mocked(updateMessageRequestDetailsIfUnfinalized).mock.calls[0]).toEqual([ + 123, + expect.objectContaining({ + statusCode: 499, + errorMessage: "CLIENT_ABORTED", + providerId: 1, + providerChain: [ + expect.objectContaining({ + id: 1, + name: "avemujica-responses", + statusCode: 499, + errorMessage: "CLIENT_ABORTED", + }), + ], + }), + ]); + }); + + 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 }); setDeferredStreamingFinalization(session, { providerId: 1, providerName: "avemujica-responses", @@ -1189,27 +2664,593 @@ 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", }) ); } 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", + }) + ); + 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(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 () => await durableAck.promise + ); + 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(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", + }), + ], + }) + ); + }); + + 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"), + }), + ] + : [], + }) + ); + }); + + 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 () => await durableAck.promise + ); + 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 before completing 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).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDuration).toHaveBeenCalledWith(123, expect.any(Number)); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ statusCode: 200 }) + ); + }); + + it("persists one durable 502 before Provider circuit mutation on non-stream response timeout", async () => { + const durableAck = createDeferred(); + vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce( + async () => await durableAck.promise + ); + 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 () => await durableAck.promise + ); + 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 () => await durableAck.promise + ); + 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(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).toHaveBeenCalledWith(123, expect.any(Number)); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 502, + errorMessage: expect.stringContaining("Gemini non-stream body read failed"), + }) + ); + 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..ea5030e99 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,7 @@ vi.mock("@/repository/message", () => ({ updateMessageRequestCost: vi.fn(), updateMessageRequestCostWithBreakdown: vi.fn(), updateMessageRequestDetails: vi.fn(), + updateMessageRequestDetailsDurably: vi.fn(), updateMessageRequestDuration: vi.fn(), })); @@ -62,6 +77,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 +88,7 @@ vi.mock("@/lib/rate-limit", () => ({ trackCost: vi.fn(), trackUserDailyCost: vi.fn(), decrementLeaseBudget: vi.fn(), + settleLeaseBudgets: vi.fn(), }, })); @@ -89,16 +107,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 +135,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 +335,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 +356,30 @@ function setupCommonMocks() { updatedAt: new Date(), }); vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); 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 +399,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( @@ -385,7 +428,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( @@ -403,7 +447,8 @@ 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(); @@ -439,7 +484,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 +497,25 @@ 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 }) + ); }); it("streaming success without endpointId should NOT call any endpoint circuit breaker function", async () => { @@ -463,7 +523,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-gemini-stream-passthrough-timeouts.test.ts b/tests/unit/proxy/response-handler-gemini-stream-passthrough-timeouts.test.ts index c786c0e27..7011abacd 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), @@ -39,10 +53,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 +89,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 +121,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), @@ -395,7 +438,7 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { .detailSnapshotResponseBeforeSource ).toBeNull(); - await Promise.allSettled(asyncTasks); + await expectAllFulfilled(asyncTasks); }); test("Gemini 流式透传禁用 idle timeout 时不应回落到默认 stale cleanup", async () => { @@ -425,7 +468,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 +547,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 +613,7 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { } finally { clientAbortController.abort(new Error("test_cleanup")); await close(); - await Promise.allSettled(asyncTasks); + await expectAllFulfilled(asyncTasks); } }); @@ -635,7 +678,7 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { } finally { clientAbortController.abort(new Error("test_cleanup")); await close(); - await Promise.allSettled(asyncTasks); + await expectAllFulfilled(asyncTasks); } }); @@ -697,7 +740,7 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { expect(first.done).toBe(false); clientAbortController.abort(new Error("client_cancelled")); - await Promise.allSettled(asyncTasks); + await expectAllFulfilled(asyncTasks); expect(vi.mocked(SessionManager.clearSessionProvider)).toHaveBeenCalledWith( "gemini-abort-session" @@ -705,14 +748,14 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { } 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,9 +799,9 @@ 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, @@ -768,4 +811,74 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { ); expect(SessionManager.storeSessionResponse).not.toHaveBeenCalled(); }); + + test("Gemini 终态副作用挂起时应先释放传输资源并在共享 deadline 后收敛", async () => { + vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); + asyncTasks.length = 0; + vi.mocked(updateMessageRequestDetailsDurably).mockClear(); + vi.mocked(SessionManager.clearSessionProvider).mockClear(); + const releaseAgent = vi.fn(); + vi.mocked(SessionManager.clearSessionProvider).mockImplementationOnce(() => { + return new Promise(() => {}); + }); + + 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); + + await vi.advanceTimersByTimeAsync(120_000); + await expectAllFulfilled(asyncTasks.splice(0, asyncTasks.length)); + + expect(releaseAgent).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); }); 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..4affc7bcd 100644 --- a/tests/unit/proxy/response-handler-hedge-loser-priority.test.ts +++ b/tests/unit/proxy/response-handler-hedge-loser-priority.test.ts @@ -7,6 +7,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 +35,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 +63,7 @@ vi.mock("@/lib/rate-limit", () => ({ trackCost: mocks.trackCost, trackUserDailyCost: mocks.trackUserDailyCost, decrementLeaseBudget: mocks.decrementLeaseBudget, + settleLeaseBudgets: mocks.settleLeaseBudgets, }, })); @@ -188,6 +205,7 @@ 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 () => { @@ -202,6 +220,7 @@ describe("finalizeHedgeLoserBilling Codex priority snapshot", () => { const billed = await finalizeHedgeLoserBilling({ messageRequestId: 123, + messageRequestCreatedAtMs: new Date("2026-06-08T00:00:00.000Z").getTime(), loserSession: loserSession as any, provider, attemptNumber: 1, @@ -246,6 +265,7 @@ describe("finalizeHedgeLoserBilling Codex priority snapshot", () => { const billed = await finalizeHedgeLoserBilling({ messageRequestId: 123, + messageRequestCreatedAtMs: new Date("2026-06-08T00:00:00.000Z").getTime(), loserSession: loserSession as any, provider: loserProvider, attemptNumber: 1, @@ -270,17 +290,52 @@ 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 = createLoserSession(loserProvider, { sessionId: null }) as ReturnType< + typeof createLoserSession + > & { messageContext: null; 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: loserSession as any, + 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..5cf34b13d 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,7 @@ vi.mock("@/repository/message", () => ({ updateMessageRequestCost: vi.fn(), updateMessageRequestCostWithBreakdown: vi.fn(), updateMessageRequestDetails: vi.fn(), + updateMessageRequestDetailsDurably: vi.fn(), updateMessageRequestDuration: vi.fn(), })); @@ -71,6 +91,7 @@ vi.mock("@/lib/rate-limit", () => ({ trackCost: vi.fn(), trackUserDailyCost: vi.fn(), decrementLeaseBudget: vi.fn(), + settleLeaseBudgets: vi.fn(), }, })); @@ -97,6 +118,7 @@ import { SessionTracker } from "@/lib/session-tracker"; import { updateMessageRequestCost, updateMessageRequestDetails, + updateMessageRequestDetailsDurably, updateMessageRequestDuration, } from "@/repository/message"; import { findLatestPriceByModel } from "@/repository/model-price"; @@ -317,13 +339,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 +548,7 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { vi.mocked(findLatestPriceByModel).mockResolvedValue( makePriceRecord(originalModel, testPriceData) ); - vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponsePhaseSnapshot).mockResolvedValue(undefined); @@ -345,10 +558,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 +581,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,7 +624,7 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { }), session.requestSequence ); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( messageId, expect.objectContaining({ statusCode: 200, @@ -432,7 +634,7 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { ); }); - 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 +647,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 +689,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,8 +711,9 @@ 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, @@ -522,7 +723,7 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { ); }); - 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 +735,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 +802,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 +828,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: () => {}, From 9a4c66521526098ba19daaff726d98218d6c557b Mon Sep 17 00:00:00 2001 From: ding113 Date: Tue, 14 Jul 2026 19:24:51 +0800 Subject: [PATCH 03/20] feat(repository): add durable terminal write with CAS fallback for message requests Introduce enqueueMessageRequestUpdateDurably which returns a Promise that resolves only after the batch SQL commits, using RETURNING id and a status_code IS NULL fence to guarantee exactly-once terminal persistence. If the durable write fails or the row is already finalized, a conditional fallback (updateMessageRequestDetailsIfUnfinalized) prevents overwriting an existing terminal status. The write buffer now uses a dedicated writer-lane DB handle, an evictable min-heap for priority-aware overflow dropping, and aggregated overflow logging. Public-status rollup callbacks fire only after the batch commit acknowledgement, preventing premature or duplicate rollup emission when a timeout fallback races with a late primary write. --- src/repository/message-write-buffer.ts | 647 ++++++++++++++++-- src/repository/message.ts | 131 +++- .../message-public-status-rollup.test.ts | 263 ++++++- ...essage-terminal-public-status-seam.test.ts | 547 +++++++++++++++ .../repository/message-write-buffer.test.ts | 626 ++++++++++++++++- 5 files changed, 2113 insertions(+), 101 deletions(-) create mode 100644 tests/unit/repository/message-terminal-public-status-seam.test.ts diff --git a/src/repository/message-write-buffer.ts b/src/repository/message-write-buffer.ts index cab71d102..8ff18ead0 100644 --- a/src/repository/message-write-buffer.ts +++ b/src/repository/message-write-buffer.ts @@ -2,7 +2,7 @@ import "server-only"; import type { SQL } from "drizzle-orm"; import { sql } from "drizzle-orm"; -import { db } from "@/drizzle/db"; +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 +43,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 +251,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 +275,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 +335,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,8 +392,16 @@ 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 stopping = false; @@ -206,51 +411,267 @@ 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) { + if (activeAcknowledgement.state === "pending") { + const existing = this.pending.get(id); + if (existing?.durableAcknowledgement === activeAcknowledgement) { + if (options.onCommitted) { + activeAcknowledgement.onCommittedCallbacks.add(options.onCommitted); + } + this.setPending(id, mergePatch(existing.patch, patch), existing.durableAcknowledgement); + this.scheduleFlushIfNeeded(); + return activeAcknowledgement.promise; } } - if (droppedId === undefined) { - const first = this.pending.entries().next().value as - | [number, MessageRequestUpdatePatch] - | undefined; - if (first) { - droppedId = first[0]; - droppedPatch = first[1]; - } + return Promise.reject( + new Error(`durable message_request update already in flight for id ${id}`) + ); + } + + 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; + } + + this.scheduleFlushIfNeeded(); + return acknowledgement.promise; + } + + 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) { - 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 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") { + void result.catch((error: unknown) => { + logger.error("[MessageRequestWriteBuffer] Durable commit callback failed", { + error: error instanceof Error ? error.message : String(error), + messageRequestId: acknowledgement.id, + }); + }); + } + } catch (error) { + logger.error("[MessageRequestWriteBuffer] Durable commit callback failed", { + error: error instanceof Error ? error.message : String(error), + messageRequestId: acknowledgement.id, }); } } + acknowledgement.onCommittedCallbacks.clear(); + } + + 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,21 +721,73 @@ 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(); logger.error("[MessageRequestWriteBuffer] Flush failed, will retry later", { error: error instanceof Error ? error.message : String(error), @@ -346,17 +819,36 @@ 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") + ); + } + 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 +866,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 +890,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..22b457ca3 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, + type MessageRequestUpdatePatch, + enqueueMessageRequestUpdate, + enqueueMessageRequestUpdateDurably, +} 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; @@ -201,6 +207,18 @@ function queuePublicStatusRollupForFinalDetails( })(); } +function publishCommittedMessageRequestDetails( + id: number, + details: PublicStatusFinalDetails +): void { + if (details.durationMs !== undefined) { + updatePublicStatusRequestSeed(id, { durationMs: details.durationMs }); + } + if (details.providerChain !== undefined && details.statusCode !== undefined) { + queuePublicStatusRollupForFinalDetails(id, details); + } +} + /** * 创建消息请求记录 */ @@ -458,41 +476,40 @@ 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"]; // 特殊设置（审计/展示） - } + details: MessageRequestDetailsUpdate, + options: { onlyIfUnfinalized?: boolean } = {} ): Promise { - const shouldQueuePublicStatusRollup = - details.providerChain !== undefined && details.statusCode !== undefined; - - if (getEnvConfig().MESSAGE_REQUEST_WRITE_MODE === "async") { + if (getEnvConfig().MESSAGE_REQUEST_WRITE_MODE === "async" && !options.onlyIfUnfinalized) { enqueueMessageRequestUpdate(id, details); - if (shouldQueuePublicStatusRollup) { - queuePublicStatusRollupForFinalDetails(id, details); - } + publishCommittedMessageRequestDetails(id, details); return; } @@ -500,6 +517,9 @@ export async function updateMessageRequestDetails( updatedAt: new Date(), }; + if (details.durationMs !== undefined) { + updateData.durationMs = details.durationMs; + } if (details.statusCode !== undefined) { updateData.statusCode = details.statusCode; } @@ -558,10 +578,55 @@ 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 updated = await getMessageWriterDb() + .update(messageRequest) + .set(updateData) + .where(and(eq(messageRequest.id, id), isNull(messageRequest.statusCode))) + .returning({ id: messageRequest.id }); + if (updated.length === 0) { + return; + } + } else { + await db.update(messageRequest).set(updateData).where(eq(messageRequest.id, id)); + } + publishCommittedMessageRequestDetails(id, details); +} + +export async function updateMessageRequestDetailsIfUnfinalized( + id: number, + details: MessageRequestDetailsUpdate +): Promise { + await updateMessageRequestDetails(id, details, { onlyIfUnfinalized: true }); +} + +/** + * 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") { + await updateMessageRequestDetails(id, details); + return; } + + let commitPublished = false; + const publishCommit = (committedPatch: Readonly) => { + if (commitPublished) return; + commitPublished = true; + publishCommittedMessageRequestDetails(id, committedPatch); + return options?.onCommitted?.(committedPatch); + }; + + await enqueueMessageRequestUpdateDurably(id, details, { + ...options, + onCommitted: publishCommit, + }); + publishCommit(details); } /** diff --git a/tests/unit/repository/message-public-status-rollup.test.ts b/tests/unit/repository/message-public-status-rollup.test.ts index 2046ec89e..0e2e3633d 100644 --- a/tests/unit/repository/message-public-status-rollup.test.ts +++ b/tests/unit/repository/message-public-status-rollup.test.ts @@ -4,10 +4,18 @@ 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 +42,7 @@ vi.mock("@/drizzle/schema", () => ({ cacheCreation1hInputTokens: "cacheCreation1hInputTokens", cacheReadInputTokens: "cacheReadInputTokens", specialSettings: "specialSettings", + statusCode: "statusCode", createdAt: "createdAt", updatedAt: "updatedAt", deletedAt: "deletedAt", @@ -62,6 +71,7 @@ vi.mock("@/drizzle/db", () => ({ })), })), }, + getMessageWriterDb: mockGetMessageWriterDb, })); vi.mock("@/lib/config/env.schema", () => ({ @@ -81,21 +91,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 +173,218 @@ describe("repository/message public status rollup hook", () => { }); }); + it("writes timeout fallback details only while the request is unfinalized", async () => { + mockWriterDbUpdateReturning.mockResolvedValueOnce([{ id: 606 }]); + const { updateMessageRequestDetailsIfUnfinalized } = await import("@/repository/message"); + + await updateMessageRequestDetailsIfUnfinalized(606, { + statusCode: 500, + errorMessage: "Error: stream_finalization_timeout", + }); + + 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("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 () => { + mockWriterDbUpdateReturning.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).toHaveBeenCalledTimes(1); + 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.mockReturnValueOnce(durableAck.promise); + 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: (() => void) | undefined; + mockEnqueueMessageRequestUpdateDurably.mockImplementationOnce( + (_id, _details, options: { onCommitted?: () => void } | 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", + }); + + await expect( + updateMessageRequestDetailsDurably(810, { + durationMs: 1_500, + statusCode: 200, + outputTokens: 10, + providerChain: [{ id: 1, name: "provider-a", groupTag: "openai" }], + model: "gpt-4.1", + }) + ).rejects.toThrow("durable acknowledgement timed out"); + await flushMicrotasks(); + + expect(onCommitted).toBeTypeOf("function"); + expect(mockQueuePublicStatusRollupWrite).not.toHaveBeenCalled(); + + onCommitted?.(); + 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" }); + 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-terminal-public-status-seam.test.ts b/tests/unit/repository/message-terminal-public-status-seam.test.ts new file mode 100644 index 000000000..77845c51c --- /dev/null +++ b/tests/unit/repository/message-terminal-public-status-seam.test.ts @@ -0,0 +1,547 @@ +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 merge publishes one rollup from the committed latest payload", 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(); + await Promise.all([oldFailure, latestSuccess, flush]); + await flushMicrotasks(); + + expect(committedSql).toHaveLength(1); + expect(committedSql[0]?.sql).toMatch(/"?status_code"? IS NULL/); + expect(committedSql[0]?.sql).toContain("RETURNING id"); + expect(row).toMatchObject({ + durationMs: latestSuccessDetails.durationMs, + statusCode: latestSuccessDetails.statusCode, + inputTokens: oldFailureDetails.inputTokens, + outputTokens: latestSuccessDetails.outputTokens, + ttfbMs: latestSuccessDetails.ttfbMs, + providerChain: latestSuccessDetails.providerChain, + providerId: latestSuccessDetails.providerId, + }); + expect(rollupPipelines).toHaveLength(1); + + const rollupIncrementOperations = rollupPipelines[0]!.filter( + (operation) => operation.command === "hincrbyfloat" + ); + expect(rollupIncrementOperations).toHaveLength(5); + const rollupIncrements = Object.fromEntries( + rollupIncrementOperations.map((operation) => [ + String(operation.args[1]), + Number(operation.args[2]), + ]) + ); + expect(rollupIncrements).toEqual({ + "42|gpt-4.1|success": 1, + "42|gpt-4.1|ttfb_sum": latestSuccessDetails.ttfbMs, + "42|gpt-4.1|ttfb_count": 1, + "42|gpt-4.1|tps_sum": 80, + "42|gpt-4.1|tps_count": 1, + }); + + await stopMessageRequestWriteBuffer(); + }); +}); diff --git a/tests/unit/repository/message-write-buffer.test.ts b/tests/unit/repository/message-write-buffer.test.ts index 88c53d1aa..620a02443 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,513 @@ 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"); + await expect( + Promise.race([durablePromise.then(() => "resolved"), Promise.resolve("pending")]) + ).resolves.toBe("pending"); + + 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.toBeUndefined(); + 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([undefined, undefined]); + 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(); + await expect( + Promise.race([durablePromise.then(() => "resolved"), Promise.resolve("pending")]) + ).resolves.toBe("pending"); + + await flushMessageRequestWriteBuffer(); + await expect(durablePromise).resolves.toBeUndefined(); + 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.toBeUndefined(); + 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.toBeUndefined(); + + 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.toBeUndefined(); + 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.toBeUndefined(); + expect(loggerErrorMock).toHaveBeenCalledWith( + "[MessageRequestWriteBuffer] Durable commit callback failed", + expect.objectContaining({ + error: "rollup callback failed", + messageRequestId: 324, + }) + ); + await stopMessageRequestWriteBuffer(); + }); + + 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; + const settlementsBeforeRelease = await Promise.all([ + Promise.race([ + outerStopPromise.then( + () => "fulfilled", + () => "rejected" + ), + Promise.resolve("pending"), + ]), + Promise.race([ + reentrantPromise.then( + () => "fulfilled", + () => "rejected" + ), + Promise.resolve("pending"), + ]), + ]); + + if (shouldReject) { + databaseBarrier.reject(databaseError); + } else { + databaseBarrier.resolve([]); + } + const stopResults = await Promise.allSettled([outerStopPromise, reentrantPromise]); + + expect(settlementsBeforeRelease).toEqual(["pending", "pending"]); + 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 () => { @@ -264,6 +803,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"; From f931ab8a3cd364cb7d7ed0d055f335d4d0c8048c Mon Sep 17 00:00:00 2001 From: ding113 Date: Tue, 14 Jul 2026 19:24:51 +0800 Subject: [PATCH 04/20] perf(rate-limit): batch cost tracking into single pipeline and settle leases atomically Consolidate all rolling-window and fixed-window cost writes into one Redis pipeline per trackCost / trackUserDailyCost call, eliminating sequential round-trips. The unified TRACK_COST_ROLLING_WINDOW script is now write-only (cleanup + append + TTL repair) and no longer scans the full ZSET on every successful request. Add settleLeaseBudgets: a single Lua invocation that validates all twelve lease keys, applies decrements, and writes an idempotency marker in one atomic step. Replace twelve fire-and-forget decrementLeaseBudget calls with one settlement, preventing duplicate deductions on ioredis reconnect. Configure commandTimeout, socketTimeout, and autoResendUnfulfilledCommands on the shared Redis client so a TCP blackhole cannot grow the command queue or replay timed-out writes after reconnection. --- src/lib/rate-limit/lease-service.ts | 254 ++++++++++++++++ src/lib/rate-limit/service.ts | 277 ++++++++++++------ src/lib/redis/client.ts | 10 + src/lib/redis/lua-scripts.ts | 99 ++----- .../lease-settlement-redis.test.ts | 243 +++++++++++++++ tests/integration/rolling-cost-redis.test.ts | 130 ++++++++ .../unit/lib/rate-limit/lease-service.test.ts | 215 ++++++++++++++ .../lib/rate-limit/rolling-window-5h.test.ts | 68 +---- .../rolling-window-cache-warm.test.ts | 13 +- .../unit/lib/rate-limit/service-extra.test.ts | 157 +++++++++- 10 files changed, 1218 insertions(+), 248 deletions(-) create mode 100644 tests/integration/lease-settlement-redis.test.ts create mode 100644 tests/integration/rolling-cost-redis.test.ts diff --git a/src/lib/rate-limit/lease-service.ts b/src/lib/rate-limit/lease-service.ts index 6406afca0..fda8d3fee 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,215 @@ 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 + settlements[#settlements + 1] = {-1, remaining} + 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/tests/integration/lease-settlement-redis.test.ts b/tests/integration/lease-settlement-redis.test.ts new file mode 100644 index 000000000..6d012d74e --- /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.5, + }); + 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/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/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 () => { From c276038e13415b09d506f2a72acf3ef7873502ee Mon Sep 17 00:00:00 2001 From: ding113 Date: Tue, 14 Jul 2026 19:24:52 +0800 Subject: [PATCH 05/20] fix(shutdown): join in-flight tasks and close DB pools after writer quiescence MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AsyncTaskManager.register now accepts a factory invoked after admission rather than a pre-started Promise, enabling shutdown to abort and join all pending generations—including tail tasks registered by abort listeners. shutdownAll returns a shared Promise that loops until the pending set is empty. runApplicationCleanup treats async-task settlement, writer flush, and DB pool close as non-detachable critical barriers: per-step timeouts become soft warnings, and failures propagate to server.js which exits non-zero. The hard-exit watchdog timer is now referenced so the process stays alive long enough to report a truthful exit status. --- server.js | 46 +- src/lib/async-task-manager.ts | 184 +++++-- src/lib/lifecycle/shutdown.ts | 104 ++-- src/lib/price-sync/cloud-price-updater.ts | 4 +- .../async-task-manager-edge-runtime.test.ts | 478 ++++++++++++++++-- tests/unit/lib/shutdown.test.ts | 84 ++- .../price-sync/cloud-price-updater.test.ts | 24 +- tests/unit/server-shutdown.test.ts | 160 +++++- 8 files changed, 952 insertions(+), 132 deletions(-) diff --git a/server.js b/server.js index f7ede45ed..cb6fcc0ae 100644 --- a/server.js +++ b/server.js @@ -742,10 +742,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 +770,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 +804,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 +851,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 }); diff --git a/src/lib/async-task-manager.ts b/src/lib/async-task-manager.ts index 912cf4140..01c604f2e 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,70 @@ 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)); + + for (const taskInfo of activeTasks) { + this.cleanup(taskInfo.taskId, taskInfo); + } + } + })().then(resolveShutdown, (error) => { + this.lifecycleState = "closed"; + rejectShutdown(error); + }); + + return shutdownPromise; + } + /** * 获取当前活跃任务数 */ getActiveTaskCount(): number { - return this.tasks.size; + return this.pendingTasks.size; } /** @@ -288,8 +392,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 +407,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/lifecycle/shutdown.ts b/src/lib/lifecycle/shutdown.ts index 00521d875..dffb6c2ea 100644 --- a/src/lib/lifecycle/shutdown.ts +++ b/src/lib/lifecycle/shutdown.ts @@ -50,11 +50,13 @@ 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,6 +66,7 @@ export async function runApplicationCleanup( const startedAt = Date.now(); logger.info("[Shutdown] application cleanup starting", { signal, totalMs, stepMs }); + let writerQuiescencePending = false; const work = (async () => { // 1. 停止本地周期任务（不需要做 IO，几乎是同步） @@ -113,26 +116,58 @@ export async function runApplicationCleanup( // 5. 取消仍在飞的后台异步任务。 // 必须排在 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" - ); + // 6. 刷写 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. 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); + } - // 7. Langfuse 自带超时（LANGFUSE_SHUTDOWN_TIMEOUT_MS），这里再加一层兜底 + // 8. Langfuse 自带超时（LANGFUSE_SHUTDOWN_TIMEOUT_MS），这里再加一层兜底 await withTimeout( (async () => { const { shutdownLangfuse } = await import("@/lib/langfuse"); @@ -142,7 +177,7 @@ export async function runApplicationCleanup( "shutdownLangfuse" ); - // 8. Redis 连接最后关：上面的步骤可能仍在写日志/缓存 + // 9. Redis 连接最后关：上面的步骤可能仍在写日志/缓存 await withTimeout( (async () => { const { closeRedis } = await import("@/lib/redis"); @@ -152,7 +187,7 @@ export async function runApplicationCleanup( "closeRedis" ); - // 9. API Key Vacuum Filter 订阅清理 —— 同步函数，不需要 timeout + // 10. API Key Vacuum Filter 订阅清理 —— 同步函数，不需要 timeout try { const g = globalThis as unknown as { __CCH_API_KEY_VF_SYNC_CLEANUP__?: (() => void) | null; @@ -164,7 +199,7 @@ export async function runApplicationCleanup( }); } - // 10. 云价格定时同步 + // 11. 云价格定时同步 try { const g = globalThis as unknown as { __CCH_CLOUD_PRICE_SYNC_INTERVAL_ID__?: ReturnType; @@ -180,15 +215,24 @@ export async function runApplicationCleanup( } })(); - 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/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/shutdown.test.ts b/tests/unit/lib/shutdown.test.ts index 3d3a39ae3..67985a113 100644 --- a/tests/unit/lib/shutdown.test.ts +++ b/tests/unit/lib/shutdown.test.ts @@ -14,6 +14,7 @@ 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 }) @@ -23,6 +24,7 @@ describe.sequential("lifecycle/shutdown", () => { afterEach(() => { vi.useRealTimers(); vi.unstubAllEnvs(); + delete (globalThis as unknown as { __ASYNC_TASK_MANAGER__?: unknown }).__ASYNC_TASK_MANAGER__; }); it("markShuttingDown flips isShuttingDown idempotently", async () => { @@ -50,14 +52,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 () => {}); @@ -140,4 +142,76 @@ 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); + }); }); 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/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"'); }); }); From 0c8abf7725088ed6001df681c7bad3cda45f97ca Mon Sep 17 00:00:00 2001 From: ding113 Date: Tue, 14 Jul 2026 19:24:52 +0800 Subject: [PATCH 06/20] fix(langfuse): redact credential headers at the trace boundary Replace the raw headersToRecord helper with redactHeaders so authorization, cookie, x-api-key, and set-cookie values are masked before being written to Langfuse generation metadata, preventing secret leakage in observability exports. --- src/lib/langfuse/trace-proxy-request.ts | 29 +++--------- tests/unit/langfuse/langfuse-trace.test.ts | 52 +++++++++++++++++++--- 2 files changed, 52 insertions(+), 29 deletions(-) diff --git a/src/lib/langfuse/trace-proxy-request.ts b/src/lib/langfuse/trace-proxy-request.ts index abbe359ea..ec4641239 100644 --- a/src/lib/langfuse/trace-proxy-request.ts +++ b/src/lib/langfuse/trace-proxy-request.ts @@ -1,5 +1,6 @@ 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,25 +39,6 @@ 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 = {}; - headers.forEach((value, key) => { - result[key] = value; - }); - return result; -} - const SUCCESS_REASONS = new Set([ "request_success", "retry_success", @@ -242,7 +224,9 @@ export async function traceProxyRequest(ctx: TraceContext): Promise { requestSequence: String(session.getRequestSequence()), }; - // Build generation metadata - all request detail fields, raw headers (no redaction) + const requestHeaders = redactHeaders(session.headers); + const responseHeaders = redactHeaders(ctx.responseHeaders); + const generationMetadata: Record = { // Provider providerId: provider?.id, @@ -277,9 +261,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 diff --git a/tests/unit/langfuse/langfuse-trace.test.ts b/tests/unit/langfuse/langfuse-trace.test.ts index 91bd63899..1d9a9c010 100644 --- a/tests/unit/langfuse/langfuse-trace.test.ts +++ b/tests/unit/langfuse/langfuse-trace.test.ts @@ -234,12 +234,28 @@ 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-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 +265,33 @@ 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); + } }); test("should include provider name and model in tags", async () => { From 81259bd55abcaa5ded3925b04219f49ec09f9797 Mon Sep 17 00:00:00 2001 From: ding113 Date: Tue, 14 Jul 2026 19:25:43 +0800 Subject: [PATCH 07/20] test(redis): add timeout and reconnect regression coverage Verify that buildRedisOptionsForUrl sets commandTimeout, socketTimeout, and autoResendUnfulfilledCommands on both redis:// and rediss:// URLs, and that REDIS_COMMAND_TIMEOUT_MS overrides the defaults. Guards the rate-limit performance commit against silent regressions in Redis client hardening. --- tests/unit/lib/redis/client.test.ts | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/unit/lib/redis/client.test.ts b/tests/unit/lib/redis/client.test.ts index eb32520ec..cac2f337a 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", () => { From e488020b0d79b842462a4d44382bf8477b077178 Mon Sep 17 00:00:00 2001 From: ding113 Date: Tue, 14 Jul 2026 20:13:00 +0800 Subject: [PATCH 08/20] fix(repository): pass committed patch to delayed durable onCommitted callback The public-status rollup test for timed-out durable waiters previously typed onCommitted as () => void and invoked it with no arguments, mismatching the real DurableMessageRequestUpdateOptions contract which receives the committed MessageRequestUpdatePatch. Align the test mock and invocation so the delayed-commit path exercises the actual callback signature. Also fix import ordering in message.ts to satisfy the formatter. --- src/repository/message.ts | 2 +- .../message-public-status-rollup.test.ts | 29 +++++++++++-------- 2 files changed, 18 insertions(+), 13 deletions(-) diff --git a/src/repository/message.ts b/src/repository/message.ts index 22b457ca3..a779e0f1d 100644 --- a/src/repository/message.ts +++ b/src/repository/message.ts @@ -19,9 +19,9 @@ import { EXCLUDE_WARMUP_CONDITION } from "./_shared/message-request-conditions"; import { toMessageRequest } from "./_shared/transformers"; import { type DurableMessageRequestUpdateOptions, - type MessageRequestUpdatePatch, enqueueMessageRequestUpdate, enqueueMessageRequestUpdateDurably, + type MessageRequestUpdatePatch, } from "./message-write-buffer"; type PublicStatusRequestSeed = { diff --git a/tests/unit/repository/message-public-status-rollup.test.ts b/tests/unit/repository/message-public-status-rollup.test.ts index 0e2e3633d..0783664ff 100644 --- a/tests/unit/repository/message-public-status-rollup.test.ts +++ b/tests/unit/repository/message-public-status-rollup.test.ts @@ -1,4 +1,8 @@ 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()); @@ -289,9 +293,9 @@ describe("repository/message public status rollup hook", () => { it("publishes the public-status rollup when a timed-out durable waiter commits later", async () => { mockGetEnvConfig.mockReturnValue({ MESSAGE_REQUEST_WRITE_MODE: "async" }); - let onCommitted: (() => void) | undefined; + let onCommitted: DurableMessageRequestUpdateOptions["onCommitted"]; mockEnqueueMessageRequestUpdateDurably.mockImplementationOnce( - (_id, _details, options: { onCommitted?: () => void } | undefined) => { + (_id, _details, options: DurableMessageRequestUpdateOptions | undefined) => { onCommitted = options?.onCommitted; return Promise.reject(new Error("durable acknowledgement timed out")); } @@ -336,21 +340,22 @@ describe("repository/message public status rollup hook", () => { original_model: "gpt-4.1", }); - await expect( - updateMessageRequestDetailsDurably(810, { - durationMs: 1_500, - statusCode: 200, - outputTokens: 10, - providerChain: [{ id: 1, name: "provider-a", groupTag: "openai" }], - model: "gpt-4.1", - }) - ).rejects.toThrow("durable acknowledgement timed out"); + 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(); - onCommitted?.(); + await onCommitted?.(terminalPatch); await flushMicrotasks(); expect(mockQueuePublicStatusRollupWrite).toHaveBeenCalledTimes(1); From 9052d9fd175b6783ad6ccc32cd7f13425c490f97 Mon Sep 17 00:00:00 2001 From: ding113 Date: Fri, 17 Jul 2026 00:18:21 +0800 Subject: [PATCH 09/20] fix(ws): bound outbound backpressure and drain request bodies safely Replace fire-and-forget safeSend with a per-WebSocket outbound queue that caps pending bytes at 1 MiB, serializes sends behind a single in-flight callback, and pauses the upstream SSE response until the client drains. Late callbacks from a closed socket are invalidated by generation counter so they cannot deliver stale frames. Guard the internal HTTP request body write with a 30 s drain timeout and destroy both the request and response when the client vanishes mid-payload. Error and close paths now send a structured error frame before initiating the WebSocket close handshake so clients always receive a terminal event. --- server.js | 338 +++++++++++++++--- ...server-response-write-backpressure.test.ts | 286 +++++++++++++++ 2 files changed, 568 insertions(+), 56 deletions(-) create mode 100644 tests/unit/server-response-write-backpressure.test.ts diff --git a/server.js b/server.js index cb6fcc0ae..3b5a50eda 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, + }); + } +} + +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); } - return false; + 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) { @@ -136,18 +260,25 @@ async function handleWebSocketConnection(ws, req) { // 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; 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 = () => { @@ -178,6 +309,7 @@ async function handleWebSocketConnection(ws, req) { closed = true; abortCurrentInternalReq(); dropPendingFrames(); + invalidateOutboundSends(ws); cleanupUpstreamWsSession(); }; @@ -201,6 +333,7 @@ async function handleWebSocketConnection(ws, req) { closed = true; abortCurrentInternalReq(); dropPendingFrames(); + invalidateOutboundSends(ws); cleanupUpstreamWsSession(); log("info", "ws_client_close_initiated", { code, reason }); try { @@ -209,6 +342,10 @@ async function handleWebSocketConnection(ws, req) { log("warn", "ws_client_close_failed", { error: String(err) }); } }; + const sendErrorAndClose = (error, close) => { + const finish = () => requestClose(close.code, close.reason); + safeSend(ws, { type: "error", error }, { onSuccess: finish, onFailure: finish }); + }; ws.on("close", finalize); ws.on("error", (err) => { @@ -222,8 +359,10 @@ async function handleWebSocketConnection(ws, req) { if (closed) return; 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; } @@ -272,13 +411,15 @@ async function handleWebSocketConnection(ws, req) { req, body, responsesWsSessionId, - (clientReq) => { + (clientReq, clientRes) => { currentInternalReq = clientReq; + if (clientRes) currentInternalRes = clientRes; }, requestClose ); if (!closed) { currentInternalReq = null; + currentInternalRes = null; } }; @@ -298,8 +439,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 queued request"); - requestClose(1011, "internal_error"); + sendErrorAndClose( + { code: "internal_error", message: "Failed to process queued request" }, + { code: 1011, reason: "internal_error" } + ); }); } } @@ -308,8 +451,10 @@ async function handleWebSocketConnection(ws, req) { ws.on("message", (data, isBinary) => { if (closed) 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 +465,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 +477,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" } + ); }); }); } @@ -424,6 +573,9 @@ async function forwardToInternalHttp( headers: internalHeaders, }, (res) => { + if (typeof registerInternalReq === "function") { + registerInternalReq(req, res); + } const contentType = (res.headers["content-type"] || "").toLowerCase(); const isSse = contentType.includes("text/event-stream"); let responseSettled = false; @@ -433,6 +585,22 @@ async function forwardToInternalHttp( resolve(); return true; }; + const settleAndClose = (reason) => { + initiateClose(1011, reason); + settleResponse(); + }; + 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 @@ -450,47 +618,47 @@ 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: settleResponse, 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: settleResponse, 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, + 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; } @@ -505,8 +673,8 @@ async function forwardToInternalHttp( const failIfUnsettled = (code, message, closeReason) => { if (responseSettled) return; if (!sawTerminal) { - emitErrorEvent(ws, code, message); - initiateClose(1011, closeReason); + sendFatalError(code, message, closeReason); + return; } settleResponse(); }; @@ -532,7 +700,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: settleResponse, + onFailure: settleAndClose, + }); sawTerminal = true; } continue; @@ -542,11 +714,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 ? settleResponse : undefined, + onFailure: settleAndClose, + }); + if (isTerminalEvent) { sawTerminal = true; terminalEventType = event.type; log("info", "ws_terminal_event_sent", { type: event.type, source: "sse" }); @@ -568,12 +749,11 @@ async function forwardToInternalHttp( 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 @@ -581,7 +761,7 @@ async function forwardToInternalHttp( // errors initiate a close handshake. log("info", "ws_turn_completed", { terminalEventType }); } - settleResponse(); + if (!sawTerminal) return; }); res.on("error", (err) => { failIfUnsettled( @@ -619,8 +799,53 @@ async function forwardToInternalHttp( if (typeof registerInternalReq === "function") { registerInternalReq(req); } - req.write(payload); - req.end(); + 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(); + resolve(); + }; + const expireRequestBody = () => { + if (requestBodyFinished) return; + requestBodyFinished = true; + clearRequestBodyListeners(); + if (!req.destroyed) req.destroy(); + initiateClose(1011, "internal_request_drain_timeout"); + resolve(); + }; + 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); + } }); } @@ -880,6 +1105,7 @@ module.exports = { registerOrchestratedShutdown, WS_MAX_PAYLOAD_BYTES, MAX_PENDING_BYTES, + MAX_PENDING_OUTBOUND_BYTES, }; if (require.main === module) { 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..938521c20 --- /dev/null +++ b/tests/unit/server-response-write-backpressure.test.ts @@ -0,0 +1,286 @@ +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) => 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(), 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.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("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"); + }); +}); From 0d8fd7d64def183f54125d90735ca76148415f74 Mon Sep 17 00:00:00 2001 From: ding113 Date: Fri, 17 Jul 2026 00:18:22 +0800 Subject: [PATCH 10/20] fix(proxy): track concurrency ownership to prevent leak on early exit Capture the sessionId at increment time into a dedicated variable so the finally block decrements exactly the session it acquired, even if ProxySession.fromContext or the guard pipeline mutates session state before forwarding begins. Previously the finally block re-read session.sessionId, which could be null or different by the time the handler reached its cleanup path, leaking a concurrency slot. Add unit tests covering early guard rejection, successful forwarding, post-session error translation, and pre-session decode failures. --- src/app/v1/_lib/proxy-handler.ts | 6 +- ...roxy-handler-concurrency-ownership.test.ts | 103 ++++++++++ .../proxy/proxy-handler-public-errors.test.ts | 177 ++++++++++++++++++ .../proxy-handler-public-success.test.ts | 177 ++++++++++++++++++ 4 files changed, 461 insertions(+), 2 deletions(-) create mode 100644 tests/unit/proxy/proxy-handler-concurrency-ownership.test.ts create mode 100644 tests/unit/proxy/proxy-handler-public-errors.test.ts create mode 100644 tests/unit/proxy/proxy-handler-public-success.test.ts diff --git a/src/app/v1/_lib/proxy-handler.ts b/src/app/v1/_lib/proxy-handler.ts index 1fbde5c27..957d96452 100644 --- a/src/app/v1/_lib/proxy-handler.ts +++ b/src/app/v1/_lib/proxy-handler.ts @@ -18,6 +18,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 { @@ -88,6 +89,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. 记录请求开始 @@ -146,8 +148,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/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: [], + }); + }); +}); From 0799d90151ba5488d80bfc67396f3abb984a5cfc Mon Sep 17 00:00:00 2001 From: ding113 Date: Fri, 17 Jul 2026 00:18:22 +0800 Subject: [PATCH 11/20] refactor(proxy): drive Gemini passthrough through demand-driven pump Replace the cloned-response reader pattern with a single demand-driven response pump that feeds the client stream and observes chunks for stats collection. The client now receives a new Response wrapping the pump stream instead of the original upstream Response, eliminating the unbounded clone that held a second copy of every streaming byte in memory. The pump arms a 60 s deadline on each unconsumed lookahead chunk so a connected-but-non-reading client cannot pin the upstream connection indefinitely. Client cancellation propagates to the source even when the cancel observer throws, and reentrant cancel calls preserve the first hard-cancel owner. Response handler terminal paths now use durable persistence with conditional fallback and bounded failure deadlines so a hanging database cannot stall the streaming task indefinitely. --- .../_lib/proxy/demand-driven-response-pump.ts | 107 +++++-- src/app/v1/_lib/proxy/response-handler.ts | 184 ++++++------ .../connected-non-reader-lifetime.test.ts | 258 +++++++++++++++++ ...sponse-handler-exported-finalizers.test.ts | 202 +++++++++++++ ...gemini-stream-passthrough-timeouts.test.ts | 7 +- .../response-handler-gemini-terminal.test.ts | 228 +++++++++++++++ ...ponse-handler-hedge-loser-priority.test.ts | 103 ++++--- ...esponse-handler-nonstream-terminal.test.ts | 262 +++++++++++++++++ .../response-handler-public-dispatch.test.ts | 213 ++++++++++++++ .../response-handler-stream-terminal.test.ts | 265 ++++++++++++++++++ 10 files changed, 1643 insertions(+), 186 deletions(-) create mode 100644 tests/unit/proxy/connected-non-reader-lifetime.test.ts create mode 100644 tests/unit/proxy/response-handler-exported-finalizers.test.ts create mode 100644 tests/unit/proxy/response-handler-gemini-terminal.test.ts create mode 100644 tests/unit/proxy/response-handler-nonstream-terminal.test.ts create mode 100644 tests/unit/proxy/response-handler-public-dispatch.test.ts create mode 100644 tests/unit/proxy/response-handler-stream-terminal.test.ts diff --git a/src/app/v1/_lib/proxy/demand-driven-response-pump.ts b/src/app/v1/_lib/proxy/demand-driven-response-pump.ts index 438a732f9..c6ff5b6bd 100644 --- a/src/app/v1/_lib/proxy/demand-driven-response-pump.ts +++ b/src/app/v1/_lib/proxy/demand-driven-response-pump.ts @@ -23,6 +23,8 @@ export interface DemandDrivenResponsePump { wasClientAborted: () => boolean; } +const PENDING_CHUNK_DEADLINE_MS = 60_000; + function toError(error: unknown): Error { return error instanceof Error ? error : new Error(String(error)); } @@ -39,6 +41,7 @@ export function createDemandDrivenResponsePump( 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; @@ -54,15 +57,41 @@ export function createDemandDrivenResponsePump( } }; - const settle = (streamEndedNormally: boolean, error: Error | null) => { + 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) => { @@ -92,6 +121,43 @@ export function createDemandDrivenResponsePump( let scheduleDrain = () => {}; + const startDrain = (_reason?: unknown) => { + if (settled || state === "finalizing" || state === "closed") return; + if (state === "draining") { + scheduleDrain(); + return; + } + 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); + cancelSource(error); + }, PENDING_CHUNK_DEADLINE_MS); + }; + const ensureRead = (): Promise => { if (settled || pendingChunk || readInFlight) { return readInFlight ?? Promise.resolve(); @@ -116,7 +182,9 @@ export function createDemandDrivenResponsePump( } options.onChunk(result.value); + if (settled) return; pendingChunk = result.value; + armPendingChunkDeadline(); }, (error) => finishWithError(error) ) @@ -143,6 +211,7 @@ export function createDemandDrivenResponsePump( continue; } if (pendingChunk) { + clearPendingChunkDeadline(); pendingChunk = null; continue; } @@ -156,24 +225,6 @@ export function createDemandDrivenResponsePump( }); }; - const startDrain = (_reason?: unknown) => { - if (settled || state === "finalizing" || state === "closed") return; - if (state === "draining") { - scheduleDrain(); - return; - } - clientAborted = true; - state = "draining"; - try { - clientController?.error( - _reason == null ? new Error("Client disconnected") : toError(_reason) - ); - } catch { - // The ReadableStream cancel algorithm may have already detached the controller. - } - scheduleDrain(); - }; - const stream = new ReadableStream( { start(controller) { @@ -186,6 +237,7 @@ export function createDemandDrivenResponsePump( if (settled || state !== "client-active" || !pendingChunk) return; const chunk = pendingChunk; + clearPendingChunkDeadline(); pendingChunk = null; try { clientController?.enqueue(chunk); @@ -197,8 +249,11 @@ export function createDemandDrivenResponsePump( void ensureRead(); }, cancel(reason) { - options.onClientCancel?.(reason); - startDrain(reason); + try { + options.onClientCancel?.(reason); + } finally { + startDrain(reason); + } }, }, { highWaterMark: 0 } @@ -210,15 +265,7 @@ export function createDemandDrivenResponsePump( stream, completion, startDrain, - cancelSource(reason) { - if (settled) return; - const normalized = reason == null ? new Error("Source cancelled") : toError(reason); - const cancelPromise = reader.cancel(normalized); - settle(false, normalized); - void cancelPromise.catch(() => { - // The pump has already recorded the hard-cancel cause. - }); - }, + cancelSource, errorClient(error) { if (settled || state !== "client-active") return; state = "draining"; diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index f4114683b..e569ebe14 100644 --- a/src/app/v1/_lib/proxy/response-handler.ts +++ b/src/app/v1/_lib/proxy/response-handler.ts @@ -549,9 +549,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", { @@ -560,6 +560,7 @@ function discardBeforeResponseBodySnapshot(session: ProxySession): void { error, }); }); + return true; } export type UsageMetrics = { @@ -1445,7 +1446,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(); } @@ -2418,24 +2420,40 @@ 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 passthroughPump: DemandDrivenResponsePump; + passthroughPump = createDemandDrivenResponsePump({ + source: response.body, + onReadStart: () => observePassthroughReadStart(), + onChunk: (value) => observePassthroughChunk(value), + onClientCancel: (reason) => { + passthroughPump.startDrain(reason); + passthroughPump.cancelSource(reason); + }, + }); + const cleanupPassthroughClientAbortListener = bindClientAbortListener( + session.clientAbortSignal, + () => { + const reason = session.clientAbortSignal?.reason; + passthroughPump.startDrain(reason); + passthroughPump.cancelSource(reason); + } + ); const statusCode = response.status; const taskId = `stream-passthrough-${messageContext.id}`; @@ -2454,14 +2472,15 @@ 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; @@ -2501,6 +2520,10 @@ export class ProxyResponseHandler { }, idleTimeoutMs); }; + observePassthroughReadStart = () => { + if (!isFirstChunk) startIdleTimer(); + }; + const clearResponseTimeoutOnce = (firstChunkSize?: number) => { if (responseTimeoutCleared) return; if (!sessionWithCleanup.clearResponseTimeout) return; @@ -2529,9 +2552,20 @@ export class ProxyResponseHandler { return flushAndSnapshot().text; }; + observePassthroughChunk = (value) => { + if (isFirstChunk) { + isFirstChunk = false; + session.recordTtfb(); + clearResponseTimeoutOnce(value.byteLength); + } + streamTextAccumulator.pushBytes(value); + AsyncTaskManager.touch(taskId); + }; + const releaseTransportResources = () => { if (transportReleased) return; transportReleased = true; + cleanupPassthroughClientAbortListener(); cleanupTaskAbortBinding(); clearIdleTimer(); try { @@ -2557,90 +2591,20 @@ export class ProxyResponseHandler { } ); } - try { - const cancelPromise = reader?.cancel(); - cancelPromise?.catch((error) => { - logger.warn("[ResponseHandler] Gemini passthrough: Failed to cancel stats reader", { - taskId, - providerId: provider.id, - providerName: provider.name, - error: error instanceof Error ? error.message : String(error), - }); - }); - } catch (error) { - logger.warn("[ResponseHandler] Gemini passthrough: Failed to cancel stats reader", { - taskId, - providerId: provider.id, - providerName: provider.name, - error: error instanceof Error ? error.message : String(error), - }); - } - try { - reader?.releaseLock(); - } catch (error) { - logger.warn("[ResponseHandler] Gemini passthrough: Failed to release reader lock", { - taskId, - providerId: provider.id, - providerName: provider.name, - error: error instanceof Error ? error.message : String(error), - }); - } - reader = null; releaseSessionAgent(session); }; 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); - } - - // 首块数据到达后才启动 idle timer（避免与首字节超时职责重叠） - if (!isFirstChunk) { - startIdleTimer(); - } - } + 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分钟过期） @@ -2689,6 +2653,7 @@ export class ProxyResponseHandler { // 使用共享的统计处理方法 const duration = Date.now() - session.startTime; + terminalFinalizationStarted = true; const finalized = await finalizeDeferredStreamingFinalizationIfNeeded( session, allContent, @@ -2726,7 +2691,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"); @@ -2756,6 +2722,10 @@ export class ProxyResponseHandler { }); try { + if (terminalFinalizationStarted) { + throw err; + } + terminalFinalizationStarted = true; clearIdleTimer(); const allContent = flushAndJoin(); const duration = Date.now() - session.startTime; @@ -2787,13 +2757,29 @@ export class ProxyResponseHandler { }); } } catch (finalizeError) { + const persistenceDeadlineAtMs = Date.now() + STREAM_FAILURE_PERSISTENCE_MAX_MS; + 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, + awaitPersistence: (promise: Promise) => + raceWithDeadline( + promise, + persistenceDeadlineAtMs, + "stream_failure_persistence_timeout" + ), }); } } finally { @@ -2841,8 +2827,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", { @@ -3136,8 +3125,7 @@ export class ProxyResponseHandler { session.shouldPersistSessionDebugArtifacts() && !streamSnapshot?.truncated ) { - const beforeBody = - (await awaitFinalization(consumeBeforeResponseBodySnapshot(session))) ?? allContent; + const beforeBody = allContent; void SessionManager.storeSessionResponse( session.sessionId, allContent, 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..8f6c7ca11 --- /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: true }); + 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/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 7011abacd..6d2dee051 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 @@ -393,7 +393,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({ @@ -412,10 +412,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(); }, }), @@ -431,7 +432,7 @@ 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 }) 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..dbac49755 --- /dev/null +++ b/tests/unit/proxy/response-handler-gemini-terminal.test.ts @@ -0,0 +1,228 @@ +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 }; +const mocks = vi.hoisted(() => ({ + conditional: vi.fn<(id: number, details: object) => Promise>(), + details: vi.fn<(id: number, details: object) => Promise>(), + durable: vi.fn<(id: number, details: object) => 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) { + await Promise.allSettled(mocks.tasks.splice(0, mocks.tasks.length)); + } +} + +describe("ProxyResponseHandler.dispatch Gemini terminal behavior", () => { + afterEach(() => vi.useRealTimers()); + + beforeEach(() => { + vi.useRealTimers(); + mocks.tasks.length = 0; + vi.clearAllMocks(); + mocks.conditional.mockResolvedValue(undefined); + mocks.details.mockResolvedValue(undefined); + mocks.durable.mockResolvedValue(undefined); + }); + + it("cancels the Gemini source when the returned body is cancelled", async () => { + const cancelSource = vi.fn(); + let failSource = () => {}; + const source = new ReadableStream({ + cancel: cancelSource, + start(controller) { + failSource = () => controller.error(new DOMException("closed", "AbortError")); + }, + }); + const session = await createSession({}); + const returned = await ProxyResponseHandler.dispatch(session, geminiResponse(source)); + + await returned.body?.cancel(new Error("client cancelled body")); + const cancellationCount = cancelSource.mock.calls.length; + failSource(); + await settleTasks(); + + expect(cancellationCount).toBe(1); + }); + + 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({ statusCode: 500 }) + ); + }); + + 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 4affc7bcd..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(() => ({ @@ -76,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 { @@ -140,7 +142,7 @@ function createCodexProvider(overrides: Partial = {}): Provider { }; } -function createLoserSession( +async function createLoserSession( provider: Provider, overrides: { sessionId?: string | null; @@ -148,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", () => { @@ -210,7 +205,7 @@ describe("finalizeHedgeLoserBilling Codex priority snapshot", () => { 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, @@ -221,7 +216,7 @@ describe("finalizeHedgeLoserBilling Codex priority snapshot", () => { const billed = await finalizeHedgeLoserBilling({ messageRequestId: 123, messageRequestCreatedAtMs: new Date("2026-06-08T00:00:00.000Z").getTime(), - loserSession: loserSession as any, + loserSession, provider, attemptNumber: 1, upstreamStatusCode: 200, @@ -252,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, }); @@ -266,7 +261,7 @@ describe("finalizeHedgeLoserBilling Codex priority snapshot", () => { const billed = await finalizeHedgeLoserBilling({ messageRequestId: 123, messageRequestCreatedAtMs: new Date("2026-06-08T00:00:00.000Z").getTime(), - loserSession: loserSession as any, + loserSession, provider: loserProvider, attemptNumber: 1, upstreamStatusCode: 200, @@ -306,16 +301,14 @@ describe("finalizeHedgeLoserBilling Codex priority snapshot", () => { 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 = createLoserSession(loserProvider, { sessionId: null }) as ReturnType< - typeof createLoserSession - > & { messageContext: null; sessionId: null }; + 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: loserSession as any, + loserSession, provider: loserProvider, attemptNumber: 2, upstreamStatusCode: 200, 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..bededce2c --- /dev/null +++ b/tests/unit/proxy/response-handler-nonstream-terminal.test.ts @@ -0,0 +1,262 @@ +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 }; + +const mocks = vi.hoisted(() => ({ + conditional: vi.fn<(id: number, details: object) => Promise>(), + durable: vi.fn<(id: number, details: object) => 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.mockResolvedValue(undefined); + mocks.durable.mockResolvedValue(undefined); + 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({ statusCode: 503 }) + ); + 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..db6e58ae9 --- /dev/null +++ b/tests/unit/proxy/response-handler-stream-terminal.test.ts @@ -0,0 +1,265 @@ +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) => 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(undefined); + }); + + 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(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(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(cancelSource).toHaveBeenCalledOnce(); + expect(releaseAgent).toHaveBeenCalledOnce(); + }); +}); From e6bfbdb1dbd37957fc24420a87130726d725a4ac Mon Sep 17 00:00:00 2001 From: ding113 Date: Fri, 17 Jul 2026 00:18:22 +0800 Subject: [PATCH 12/20] fix(proxy): persist terminal error details durably before status tracking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Switch error-handler from separate duration and details writes to a single updateMessageRequestDetailsDurably call that includes durationMs, ensuring the Langfuse trace, persistence commit, and status-tracker endRequest fire in deterministic order. The trace is emitted first, the durable write is awaited, and only then is the status tracker released — so a database failure surfaces as a rejected handler rather than a silently orphaned trace. Add comprehensive unit tests for client-safe message sanitisation, override application, terminal status mapping, durable persistence ordering, and the end-to-end terminal outcome contract. --- src/app/v1/_lib/proxy/error-handler.ts | 6 +- .../error-handler-client-message.test.ts | 188 +++++++++++++ .../error-handler-durable-persistence.test.ts | 215 +++++++++++++++ .../error-handler-langfuse-trace.test.ts | 220 ++++----------- .../proxy/error-handler-overrides.test.ts | 176 ++++++++++++ .../error-handler-terminal-status.test.ts | 212 +++++++++++++++ .../proxy/terminal-outcome-contract.test.ts | 250 ++++++++++++++++++ 7 files changed, 1090 insertions(+), 177 deletions(-) create mode 100644 tests/unit/proxy/error-handler-client-message.test.ts create mode 100644 tests/unit/proxy/error-handler-durable-persistence.test.ts create mode 100644 tests/unit/proxy/error-handler-overrides.test.ts create mode 100644 tests/unit/proxy/error-handler-terminal-status.test.ts create mode 100644 tests/unit/proxy/terminal-outcome-contract.test.ts diff --git a/src/app/v1/_lib/proxy/error-handler.ts b/src/app/v1/_lib/proxy/error-handler.ts index 13b4e4d6e..be253da97 100644 --- a/src/app/v1/_lib/proxy/error-handler.ts +++ b/src/app/v1/_lib/proxy/error-handler.ts @@ -9,7 +9,7 @@ import { emitProxyLangfuseTrace } from "@/lib/langfuse/emit-proxy-trace"; import { logger } from "@/lib/logger"; import { ProxyStatusTracker } from "@/lib/proxy-status-tracker"; 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"; @@ -623,7 +623,6 @@ export class ProxyErrorHandler { } const duration = Date.now() - session.startTime; - await updateMessageRequestDuration(session.messageContext.id, duration); // 如果是限流错误，将元数据附加到错误消息中 let finalErrorMessage = errorMessage; @@ -632,7 +631,8 @@ export class ProxyErrorHandler { } // 保存错误信息和决策链 - await updateMessageRequestDetails(session.messageContext.id, { + await updateMessageRequestDetailsDurably(session.messageContext.id, { + durationMs: duration, errorMessage: finalErrorMessage, providerChain: session.getProviderChain(), statusCode: statusCode, 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..17afa8c13 --- /dev/null +++ b/tests/unit/proxy/error-handler-client-message.test.ts @@ -0,0 +1,188 @@ +import { Context } from "hono"; +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 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"); + }); +}); 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..6df3511ee --- /dev/null +++ b/tests/unit/proxy/error-handler-durable-persistence.test.ts @@ -0,0 +1,215 @@ +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"; +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(); + }); +}); 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..ed5350590 --- /dev/null +++ b/tests/unit/proxy/error-handler-terminal-status.test.ts @@ -0,0 +1,212 @@ +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/terminal-outcome-contract.test.ts b/tests/unit/proxy/terminal-outcome-contract.test.ts new file mode 100644 index 000000000..e51143087 --- /dev/null +++ b/tests/unit/proxy/terminal-outcome-contract.test.ts @@ -0,0 +1,250 @@ +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(); + + await expect( + Promise.race([handlePromise.then(() => "resolved"), Promise.resolve("pending")]) + ).resolves.toBe("pending"); + 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(); + }); +}); From 9bf599a3b2230d5dd2a92be3d15256807bd0ec0f Mon Sep 17 00:00:00 2001 From: ding113 Date: Fri, 17 Jul 2026 00:18:22 +0800 Subject: [PATCH 13/20] refactor(proxy): replace combined abort signal with listener-bound transport controller MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace combineAbortSignals polyfill with a dedicated AbortController whose abort is driven by lightweight bindClientAbortListener registrations on the response and client signals. The client-signal listener is cleaned up as soon as the forwarder obtains the upstream response, so a client disconnect after headers no longer aborts the in-flight transport — only the response controller retains that authority through stream consumption. Add an integration test exercising real loopback transports through the hedge lifecycle: winner settlement fences loser timers, each launched transport releases its agent exactly once, database overload is classified before fanout, and hedge winner/loser billing fires exactly once per request. --- src/app/v1/_lib/proxy/forwarder.ts | 37 +- .../integration/proxy-hedge-lifecycle.test.ts | 627 ++++++++++++++++++ 2 files changed, 650 insertions(+), 14 deletions(-) create mode 100644 tests/integration/proxy-hedge-lifecycle.test.ts diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 40f0b5ef8..17b42d60e 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -63,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 { @@ -3007,23 +3006,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 } : {}), }; @@ -3109,7 +3116,7 @@ export class ProxyForwarder { body: requestBodyJson, sessionId: getResponsesWsSessionId(session.headers), endpointId: responsesWsEndpointId, - abortSignal: combinedSignal, + abortSignal: transportController.signal, }); if ("response" in wsResult) { @@ -3672,6 +3679,8 @@ export class ProxyForwarder { } } + cleanupClientTransportSignal(); + // 检查 HTTP 错误状态（4xx/5xx 均视为失败，触发重试） // 注意：用户要求所有 4xx 都重试，包括 401、403、429 等 if (!response.ok) { @@ -3725,7 +3734,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 // 标志，重复调用安全。 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(); + } + }); +}); From ac6e6f86683e26de855c152b97d7158006537a82 Mon Sep 17 00:00:00 2001 From: ding113 Date: Fri, 17 Jul 2026 00:18:22 +0800 Subject: [PATCH 14/20] test(integration): cover async write-buffer recovery and CAS semantics Add a PostgreSQL integration test that exercises the async message write buffer under row-level locks: mixed durable/ordinary batches settle only after commit, timed-out primaries yield to fallback ownership, retired pending patches are excluded from reinserted generations, a saturated 5000-entry queue evicts non-terminal patches while retaining terminal priority, and winner-cost retries respect the bounded cadence with authoritative loser-cost accumulation. --- tests/configs/integration.config.ts | 2 + ...age-write-buffer-recovery-postgres.test.ts | 283 ++++++++++++++++++ 2 files changed, 285 insertions(+) create mode 100644 tests/integration/message-write-buffer-recovery-postgres.test.ts diff --git a/tests/configs/integration.config.ts b/tests/configs/integration.config.ts index 513a74658..ca2152353 100644 --- a/tests/configs/integration.config.ts +++ b/tests/configs/integration.config.ts @@ -12,6 +12,8 @@ export default createTestRunnerConfig({ "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/message-write-buffer-recovery-postgres.test.ts b/tests/integration/message-write-buffer-recovery-postgres.test.ts new file mode 100644 index 000000000..80bffd4e9 --- /dev/null +++ b/tests/integration/message-write-buffer-recovery-postgres.test.ts @@ -0,0 +1,283 @@ +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(); + + await expect( + Promise.race([durable.then(() => "settled"), Promise.resolve("pending")]) + ).resolves.toBe("pending"); + 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" }); + }); +}); From 99c53ef3e8c71bb475cbeda889555ac77bca8f07 Mon Sep 17 00:00:00 2001 From: ding113 Date: Fri, 17 Jul 2026 00:18:22 +0800 Subject: [PATCH 15/20] test(dashboard): pass explicit UTC timezone to resolveTimePresetDates Align test assertions with the production call site, which now requires a timezone argument to resolve date presets deterministically. --- tests/unit/dashboard/user-insights-page.test.tsx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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(); From 76b11845e3c49c5d99e903064257e79abd127a16 Mon Sep 17 00:00:00 2001 From: ding113 Date: Fri, 17 Jul 2026 00:18:22 +0800 Subject: [PATCH 16/20] test(repository): cover message terminal writes, CAS, cost accounting, and readback Add a shared drizzle query tracing harness and comprehensive unit tests for the message repository: terminal detail writes with durable acknowledgement and CAS-guarded unfinalized claims, winner and hedge-loser cost accounting with idempotent retries, session stats aggregation with provider/model/cache-TTL grouping, paged session request queries with adjacent-sequence lookup, usage-log filtering with ledger fallback, and public readback projections for single-request, session, and audit lookups. --- ...e-aggregate-multiple-session-stats.test.ts | 190 +++++++++++++ .../message-aggregate-session-stats.test.ts | 181 ++++++++++++ .../message-public-readback.test.ts | 206 ++++++++++++++ .../repository/message-query-test-support.ts | 101 +++++++ .../message-session-readback.test.ts | 261 ++++++++++++++++++ .../message-session-request-query.test.ts | 171 ++++++++++++ .../message-terminal-cas-durable.test.ts | 234 ++++++++++++++++ .../message-terminal-cost-accounting.test.ts | 133 +++++++++ .../message-terminal-write-apis.test.ts | 182 ++++++++++++ .../message-usage-logs-query.test.ts | 203 ++++++++++++++ 10 files changed, 1862 insertions(+) create mode 100644 tests/unit/repository/message-aggregate-multiple-session-stats.test.ts create mode 100644 tests/unit/repository/message-aggregate-session-stats.test.ts create mode 100644 tests/unit/repository/message-public-readback.test.ts create mode 100644 tests/unit/repository/message-query-test-support.ts create mode 100644 tests/unit/repository/message-session-readback.test.ts create mode 100644 tests/unit/repository/message-session-request-query.test.ts create mode 100644 tests/unit/repository/message-terminal-cas-durable.test.ts create mode 100644 tests/unit/repository/message-terminal-cost-accounting.test.ts create mode 100644 tests/unit/repository/message-terminal-write-apis.test.ts create mode 100644 tests/unit/repository/message-usage-logs-query.test.ts 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-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..6257c52ee --- /dev/null +++ b/tests/unit/repository/message-terminal-cas-durable.test.ts @@ -0,0 +1,234 @@ +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() { + const where = vi.fn(async (_condition: unknown) => []); + 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).toBeUndefined(); + 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", + }); + + await expect( + Promise.race([completion.then(() => "committed"), Promise.resolve("pending")]) + ).resolves.toBe("pending"); + 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"); + + await updateMessageRequestDetailsDurably(804, { durationMs: 900, statusCode: 201 }); + + expect(set).toHaveBeenCalledWith( + expect.objectContaining({ durationMs: 900, statusCode: 201, updatedAt: expect.any(Date) }) + ); + expect(writerExecute).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-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]); + }); +}); From 6a3e92924355ea37b5cdf13cbdf070112a1455fd Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 16 Jul 2026 16:20:00 +0000 Subject: [PATCH 17/20] chore: format code (perf-full-path-benchmark-optimization-2fb8cc3) --- tests/unit/lib/redis/client.test.ts | 28 +- .../connected-non-reader-lifetime.test.ts | 40 +- .../error-handler-terminal-status.test.ts | 67 +-- ...esponse-handler-client-abort-drain.test.ts | 126 ++--- ...essage-terminal-public-status-seam.test.ts | 466 +++++++++--------- .../repository/message-write-buffer.test.ts | 133 ++--- 6 files changed, 437 insertions(+), 423 deletions(-) diff --git a/tests/unit/lib/redis/client.test.ts b/tests/unit/lib/redis/client.test.ts index cac2f337a..c6dfa6c49 100644 --- a/tests/unit/lib/redis/client.test.ts +++ b/tests/unit/lib/redis/client.test.ts @@ -61,20 +61,20 @@ describe("buildRedisOptionsForUrl", () => { 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); - }); + 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/proxy/connected-non-reader-lifetime.test.ts b/tests/unit/proxy/connected-non-reader-lifetime.test.ts index 8f6c7ca11..85e8876fe 100644 --- a/tests/unit/proxy/connected-non-reader-lifetime.test.ts +++ b/tests/unit/proxy/connected-non-reader-lifetime.test.ts @@ -212,27 +212,27 @@ describe("connected non-reader response lifetime", () => { 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"]), - }; - }); + 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(); - }); + 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(); diff --git a/tests/unit/proxy/error-handler-terminal-status.test.ts b/tests/unit/proxy/error-handler-terminal-status.test.ts index ed5350590..98b2d5a60 100644 --- a/tests/unit/proxy/error-handler-terminal-status.test.ts +++ b/tests/unit/proxy/error-handler-terminal-status.test.ts @@ -159,36 +159,43 @@ describe("ProxyErrorHandler.handle terminal status", () => { ); }); - 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.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(); 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 53391b2cd..0f0beb6ad 100644 --- a/tests/unit/proxy/response-handler-client-abort-drain.test.ts +++ b/tests/unit/proxy/response-handler-client-abort-drain.test.ts @@ -2899,51 +2899,54 @@ describe("ProxyResponseHandler stream client abort finalization", () => { 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 - ); + ] 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(); + 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(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"), + }), + ] + : [], + }) + ); + } + ); it("rejects non-stream processing when both terminal persistence attempts fail", async () => { vi.mocked(updateMessageRequestDetailsDurably).mockRejectedValueOnce( @@ -3019,25 +3022,28 @@ describe("ProxyResponseHandler stream client abort finalization", () => { 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" }, - }); + ] 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(); + 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, - }), - ]); - }); + expect(recordFailure).not.toHaveBeenCalled(); + expect(session.getProviderChain()).toEqual([ + expect.objectContaining({ + id: 1, + reason: "resource_not_found", + statusCode: 404, + }), + ]); + } + ); it("persists Gemini non-stream duration before completing terminal stats", async () => { const session = createSession(new AbortController().signal, { diff --git a/tests/unit/repository/message-terminal-public-status-seam.test.ts b/tests/unit/repository/message-terminal-public-status-seam.test.ts index 77845c51c..f539bd675 100644 --- a/tests/unit/repository/message-terminal-public-status-seam.test.ts +++ b/tests/unit/repository/message-terminal-public-status-seam.test.ts @@ -52,246 +52,246 @@ describe("message terminal public-status public seam", () => { 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 }]; - }); + 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 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 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, - }, - ]), + })); + 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", - }, - ], + update: vi.fn(), }, - ], - }); - 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]); + 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", + }, + ], }, - }; - }), - }; - 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); + ], + }); + 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(); } - 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 merge publishes one rollup from the committed latest payload", async () => { vi.resetModules(); diff --git a/tests/unit/repository/message-write-buffer.test.ts b/tests/unit/repository/message-write-buffer.test.ts index 620a02443..707f12c38 100644 --- a/tests/unit/repository/message-write-buffer.test.ts +++ b/tests/unit/repository/message-write-buffer.test.ts @@ -562,75 +562,76 @@ describe("message_request 异步批量写入", () => { 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; + ])( + "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"); } - 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; - const settlementsBeforeRelease = await Promise.all([ - Promise.race([ - outerStopPromise.then( - () => "fulfilled", - () => "rejected" - ), - Promise.resolve("pending"), - ]), - Promise.race([ - reentrantPromise.then( - () => "fulfilled", - () => "rejected" - ), - Promise.resolve("pending"), - ]), - ]); - - if (shouldReject) { - databaseBarrier.reject(databaseError); - } else { - databaseBarrier.resolve([]); - } - const stopResults = await Promise.allSettled([outerStopPromise, reentrantPromise]); - - expect(settlementsBeforeRelease).toEqual(["pending", "pending"]); - 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 }, + const samePromise = outerStopPromise === reentrantPromise; + const settlementsBeforeRelease = await Promise.all([ + Promise.race([ + outerStopPromise.then( + () => "fulfilled", + () => "rejected" + ), + Promise.resolve("pending"), + ]), + Promise.race([ + reentrantPromise.then( + () => "fulfilled", + () => "rejected" + ), + Promise.resolve("pending"), + ]), ]); + + if (shouldReject) { + databaseBarrier.reject(databaseError); + } else { + databaseBarrier.resolve([]); + } + const stopResults = await Promise.allSettled([outerStopPromise, reentrantPromise]); + + expect(settlementsBeforeRelease).toEqual(["pending", "pending"]); + 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); } - expect(samePromise).toBe(true); - }); + ); it("stop 无法刷写剩余终态时所有调用都应持续拒绝同一错误", async () => { process.env.MESSAGE_REQUEST_WRITE_MODE = "async"; From ed78d3799eda40171754d6a7d5a4b58349a9113b Mon Sep 17 00:00:00 2001 From: ding113 Date: Fri, 17 Jul 2026 14:23:31 +0800 Subject: [PATCH 18/20] fix(security): redact secrets from crash reports and error responses Diagnostic reports now exclude the entire process environment via --report-exclude-env in both Dockerfiles. Crash handlers wrap database errors through findSafeDatabaseError before writing report files or stderr fallbacks, preventing DrizzleQueryError SQL text and bound parameters from reaching any public boundary. The proxy error handler returns a generic 503/500 for database failures without loading system settings or awaiting durable persistence. Langfuse traces and session metadata strip x-cch-* reserved internal headers. clearSessionProvider uses a Redis CAS compare so late cleanup from a prior provider cannot delete a binding already moved to a new one. --- Dockerfile | 4 +- deploy/Dockerfile | 4 +- src/app/v1/_lib/proxy-handler.ts | 14 ++- src/app/v1/_lib/proxy/error-handler.ts | 56 ++++++++-- src/app/v1/_lib/proxy/errors.ts | 12 ++- src/app/v1/_lib/proxy/forwarder.ts | 74 +++++++++---- src/app/v1/_lib/proxy/provider-selector.ts | 12 +-- src/drizzle/admitted-client.ts | 102 +++++++++++++++++- src/instrumentation.ts | 84 +++++++++++---- src/lib/langfuse/emit-proxy-trace.ts | 16 ++- src/lib/langfuse/trace-proxy-request.ts | 38 ++++++- src/lib/session-manager.ts | 49 +++++++-- .../db-pool-isolation-postgres.test.ts | 35 +++--- tests/unit/deploy-dockerfile-contract.test.ts | 19 ++++ tests/unit/drizzle/db-admission.test.ts | 61 +++++++++++ .../instrumentation-crash-handler.test.ts | 33 +++++- tests/unit/langfuse/emit-proxy-trace.test.ts | 43 ++++++++ tests/unit/langfuse/langfuse-trace.test.ts | 5 + .../unit/lib/session-manager-helpers.test.ts | 11 ++ .../session-manager-terminate-session.test.ts | 24 +++++ .../build-request-details-redaction.test.ts | 15 +++ .../error-handler-client-message.test.ts | 49 +++++++++ .../error-handler-durable-persistence.test.ts | 18 ++++ ...provider-selector-cross-type-model.test.ts | 6 +- ...er-selector-model-mismatch-binding.test.ts | 12 ++- .../proxy-forwarder-hedge-first-byte.test.ts | 8 +- 26 files changed, 701 insertions(+), 103 deletions(-) create mode 100644 tests/unit/deploy-dockerfile-contract.test.ts create mode 100644 tests/unit/langfuse/emit-proxy-trace.test.ts 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/src/app/v1/_lib/proxy-handler.ts b/src/app/v1/_lib/proxy-handler.ts index 957d96452..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"; @@ -30,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); @@ -136,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); } diff --git a/src/app/v1/_lib/proxy/error-handler.ts b/src/app/v1/_lib/proxy/error-handler.ts index be253da97..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,6 +9,7 @@ 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 { updateMessageRequestDetailsDurably } from "@/repository/message"; import type { SystemSettings } from "@/types/system-config"; @@ -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, @@ -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 1731b63df..647b4bfbd 100644 --- a/src/app/v1/_lib/proxy/errors.ts +++ b/src/app/v1/_lib/proxy/errors.ts @@ -13,6 +13,7 @@ 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. */ @@ -1126,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", @@ -1207,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); @@ -1234,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 = @@ -1241,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 17b42d60e..cb5dd36f8 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -4,7 +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 } from "@/drizzle/admitted-client"; +import { findDbPoolAdmissionError, findSafeDatabaseError } from "@/drizzle/admitted-client"; import { applyAnthropicProviderOverridesWithAudit } from "@/lib/anthropic/provider-overrides"; import { getCircuitState, @@ -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,9 +1788,9 @@ export class ProxyForwarder { throw lastError; } - if (errorCategory === ErrorCategory.LOCAL_OVERLOAD) { + if (databaseError) { const admission = findDbPoolAdmissionError(lastError); - logger.warn("ProxyForwarder: Local database admission rejected", { + logger.warn("ProxyForwarder: Local database operation failed", { providerId: currentProvider.id, providerName: currentProvider.name, endpointId: activeEndpoint.endpointId, @@ -1799,13 +1804,15 @@ export class ProxyForwarder { reason: "system_error", circuitState: getCircuitState(currentProvider.id), attemptNumber: attemptCount, - errorMessage, + errorMessage: databaseError.message, errorDetails: { system: { - errorType: "DbPoolAdmissionError", - errorName: "DbPoolAdmissionError", - errorMessage: admission?.message ?? errorMessage, - errorCode: admission?.code, + errorType: + databaseError.kind === "admission" ? "DbPoolAdmissionError" : "DatabaseError", + errorName: + databaseError.kind === "admission" ? "DbPoolAdmissionError" : "DatabaseError", + errorMessage: databaseError.message, + errorCode: databaseError.code, }, request: buildRequestDetails(session), }, @@ -2368,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 } @@ -3679,8 +3689,6 @@ export class ProxyForwarder { } } - cleanupClientTransportSignal(); - // 检查 HTTP 错误状态（4xx/5xx 均视为失败，触发重试） // 注意：用户要求所有 4xx 都重试，包括 401、403、429 等 if (!response.ok) { @@ -3710,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 & { @@ -3849,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 }); }; @@ -4269,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 = {}; @@ -4308,6 +4325,7 @@ export class ProxyForwarder { 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, @@ -4322,13 +4340,13 @@ export class ProxyForwarder { ...attempt.endpointAudit, reason: "system_error", attemptNumber: attempt.sequence, - errorMessage, + errorMessage: safeAdmissionMessage, circuitState: getCircuitState(attempt.provider.id), errorDetails: { system: { errorType: "DbPoolAdmissionError", errorName: "DbPoolAdmissionError", - errorMessage: admission?.message ?? errorMessage, + errorMessage: safeAdmissionMessage, errorCode: admission?.code, }, request: buildRequestDetails(session), @@ -4971,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/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/drizzle/admitted-client.ts b/src/drizzle/admitted-client.ts index b55eaa98d..e871c15e8 100644 --- a/src/drizzle/admitted-client.ts +++ b/src/drizzle/admitted-client.ts @@ -7,6 +7,14 @@ export interface DbPoolAdmissionErrorDetails { 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; @@ -57,6 +65,68 @@ 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; @@ -165,8 +235,22 @@ function wrapPendingQuery(pending: unknown, release: () => void): unknown { 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" || property === "cancel") { + if (property === "execute" || property === "forEach") { void track().catch(() => undefined); } return proxy; @@ -251,6 +335,22 @@ export function createAdmittedSqlClient( }; 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; 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/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 ec4641239..cf5d3e2d4 100644 --- a/src/lib/langfuse/trace-proxy-request.ts +++ b/src/lib/langfuse/trace-proxy-request.ts @@ -1,3 +1,4 @@ +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"; @@ -39,6 +40,34 @@ function getStatusCategory(statusCode: number): string { return `${Math.floor(statusCode / 100)}xx`; } +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) => { + if (!key.toLowerCase().startsWith("x-cch-")) { + externalHeaders.append(key, value); + } + }); + return redactHeaders(externalHeaders); +} + const SUCCESS_REASONS = new Set([ "request_success", "retry_success", @@ -224,15 +253,15 @@ export async function traceProxyRequest(ctx: TraceContext): Promise { requestSequence: String(session.getRequestSequence()), }; - const requestHeaders = redactHeaders(session.headers); - const responseHeaders = redactHeaders(ctx.responseHeaders); + 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(), @@ -327,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/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/tests/integration/db-pool-isolation-postgres.test.ts b/tests/integration/db-pool-isolation-postgres.test.ts index 821527cf1..197046577 100644 --- a/tests/integration/db-pool-isolation-postgres.test.ts +++ b/tests/integration/db-pool-isolation-postgres.test.ts @@ -5,16 +5,10 @@ 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); -if (!process.env.DSN && process.env.DATABASE_URL) { - process.env.DSN = process.env.DATABASE_URL; -} - -const previousPoolMax = process.env.DB_POOL_MAX; -process.env.DB_POOL_MAX = "6"; -vi.resetModules(); - 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[]; @@ -43,15 +37,30 @@ run("PostgreSQL pool isolation integration", () => { } 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 () => { - await dbModule.closeDbPools(); - if (previousPoolMax === undefined) { - delete process.env.DB_POOL_MAX; - } else { - process.env.DB_POOL_MAX = previousPoolMax; + 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; + } } }); 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 index f3fae9b24..39e9bade4 100644 --- a/tests/unit/drizzle/db-admission.test.ts +++ b/tests/unit/drizzle/db-admission.test.ts @@ -161,4 +161,65 @@ describe("drizzle/db outstanding admission", () => { 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/instrumentation-crash-handler.test.ts b/tests/unit/instrumentation-crash-handler.test.ts index af06d461f..0f736d520 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,13 +122,43 @@ 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: database wrappers are redacted at every crash sink", () => { + const { uncaughtException } = captureHandlers(); + 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); }); it("uncaughtException: a non-benign transport code (ECONNREFUSED) exits with code 1", () => { 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 1d9a9c010..92dcea2f8 100644 --- a/tests/unit/langfuse/langfuse-trace.test.ts +++ b/tests/unit/langfuse/langfuse-trace.test.ts @@ -248,6 +248,8 @@ describe("traceProxyRequest", () => { 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", }), }), @@ -292,6 +294,9 @@ describe("traceProxyRequest", () => { 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/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/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/error-handler-client-message.test.ts b/tests/unit/proxy/error-handler-client-message.test.ts index 17afa8c13..3c5ba79b9 100644 --- a/tests/unit/proxy/error-handler-client-message.test.ts +++ b/tests/unit/proxy/error-handler-client-message.test.ts @@ -1,4 +1,5 @@ import { Context } from "hono"; +import { DrizzleQueryError } from "drizzle-orm"; import { beforeEach, describe, expect, test, vi } from "vitest"; import { ProxyErrorHandler, @@ -6,6 +7,7 @@ import { } 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(() => ({ @@ -185,4 +187,51 @@ describe("ProxyErrorHandler.handle client message", () => { 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 index 6df3511ee..803a294d8 100644 --- a/tests/unit/proxy/error-handler-durable-persistence.test.ts +++ b/tests/unit/proxy/error-handler-durable-persistence.test.ts @@ -1,8 +1,10 @@ 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"; @@ -212,4 +214,20 @@ describe("ProxyErrorHandler.handle durable persistence", () => { 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/provider-selector-cross-type-model.test.ts b/tests/unit/proxy/provider-selector-cross-type-model.test.ts index 1bab9a1d4..f171f6ffe 100644 --- a/tests/unit/proxy/provider-selector-cross-type-model.test.ts +++ b/tests/unit/proxy/provider-selector-cross-type-model.test.ts @@ -302,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 ); }); @@ -334,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-hedge-first-byte.test.ts b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts index c4e83c99f..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(); @@ -1646,7 +1646,7 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { 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"); + expect(mocks.clearSessionProvider).toHaveBeenCalledWith("sess-hedge", 1); } finally { vi.useRealTimers(); } @@ -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({ @@ -1751,7 +1751,7 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { expect(mocks.pickRandomProviderWithExclusion).not.toHaveBeenCalled(); expect(mocks.recordEndpointFailure).not.toHaveBeenCalled(); expect(mocks.recordFailure).not.toHaveBeenCalled(); - expect(mocks.clearSessionProvider).toHaveBeenCalledWith("sess-hedge"); + expect(mocks.clearSessionProvider).toHaveBeenCalledWith("sess-hedge", 1); expect(session.getProviderChain()).toEqual([ expect.objectContaining({ id: provider.id, From 48ddcd3d49cc97ba1f391271ab9a53918ea42339 Mon Sep 17 00:00:00 2001 From: ding113 Date: Fri, 17 Jul 2026 14:23:32 +0800 Subject: [PATCH 19/20] fix(proxy): gate post-terminal side effects on durable SQL commit The first durable write claimant owns the terminal patch; later contenders observe its acknowledgement but cannot overwrite the committed outcome. Request duration is merged into the same CAS patch so overflow or process exit cannot leave a permanently active record. Post-terminal side effects (provider circuit mutation, session binding clear, Codex cache binding) fire from an onCommitted callback after the SQL write is acknowledged rather than speculatively before it. server.js WebSocket turns track per-turn request, response, and settle lifecycle: sendErrorAndClose aborts the active internal request before queuing a fatal frame, and terminal events wait for their WS send acknowledgement before releasing turn ownership. --- server.js | 237 +++++-- .../_lib/proxy/demand-driven-response-pump.ts | 6 +- src/app/v1/_lib/proxy/response-handler.ts | 658 ++++++++++++------ src/repository/message-write-buffer.ts | 55 +- src/repository/message.ts | 103 ++- ...age-write-buffer-recovery-postgres.test.ts | 14 +- .../connected-non-reader-lifetime.test.ts | 2 +- ...esponse-handler-client-abort-drain.test.ts | 393 ++++++----- ...handler-endpoint-circuit-isolation.test.ts | 12 +- ...gemini-stream-passthrough-timeouts.test.ts | 49 +- .../response-handler-gemini-terminal.test.ts | 45 +- .../response-handler-lease-decrement.test.ts | 9 +- ...esponse-handler-nonstream-terminal.test.ts | 30 +- .../response-handler-stream-terminal.test.ts | 25 +- .../proxy/terminal-outcome-contract.test.ts | 14 +- .../message-public-status-rollup.test.ts | 24 +- .../message-terminal-cas-durable.test.ts | 62 +- ...essage-terminal-public-status-seam.test.ts | 33 +- .../repository/message-write-buffer.test.ts | 165 ++++- ...server-response-write-backpressure.test.ts | 211 +++++- 20 files changed, 1535 insertions(+), 612 deletions(-) diff --git a/server.js b/server.js index 3b5a50eda..8c4469a34 100644 --- a/server.js +++ b/server.js @@ -256,11 +256,14 @@ 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 = () => { const reqToDestroy = currentInternalReq; @@ -307,6 +310,8 @@ async function handleWebSocketConnection(ws, req) { const finalize = () => { if (closed) return; closed = true; + currentTurnSettle?.(); + currentTurnSettle = null; abortCurrentInternalReq(); dropPendingFrames(); invalidateOutboundSends(ws); @@ -331,6 +336,8 @@ async function handleWebSocketConnection(ws, req) { return; } closed = true; + currentTurnSettle?.(); + currentTurnSettle = null; abortCurrentInternalReq(); dropPendingFrames(); invalidateOutboundSends(ws); @@ -343,6 +350,10 @@ async function handleWebSocketConnection(ws, req) { } }; 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 }); }; @@ -356,7 +367,35 @@ 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") { sendErrorAndClose( @@ -406,20 +445,21 @@ async function handleWebSocketConnection(ws, req) { hasPreviousResponseId: typeof body.previous_response_id === "string", }); - await forwardToInternalHttp( - ws, - req, - body, - responsesWsSessionId, - (clientReq, clientRes) => { - currentInternalReq = clientReq; - if (clientRes) currentInternalRes = clientRes; - }, - requestClose - ); - if (!closed) { - currentInternalReq = null; - currentInternalRes = 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 }); } }; @@ -434,7 +474,7 @@ 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), @@ -449,7 +489,7 @@ async function handleWebSocketConnection(ws, req) { }; ws.on("message", (data, isBinary) => { - if (closed) return; + if (closed || closing) return; if (isBinary) { sendErrorAndClose( { code: "invalid_frame_type", message: "Only text WebSocket frames are supported" }, @@ -564,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", @@ -573,21 +622,39 @@ async function forwardToInternalHttp( headers: internalHeaders, }, (res) => { - if (typeof registerInternalReq === "function") { - registerInternalReq(req, res); - } 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); - settleResponse(); + forceSettleResponse(); }; const sendFatalError = (code, message, closeReason) => { const sent = safeSend( @@ -609,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 { @@ -628,7 +696,7 @@ async function forwardToInternalHttp( ? parsed.error : { code: `http_${res.statusCode}`, message: text.slice(0, 512) }, }, - { response: res, onSuccess: settleResponse, onFailure: settleAndClose } + { response: res, onSuccess: acknowledgeTerminalSend, onFailure: settleAndClose } ); log("info", "ws_terminal_event_sent", { type: "error", @@ -639,7 +707,7 @@ async function forwardToInternalHttp( safeSend( ws, { type: "response.completed", response: parsed }, - { response: res, onSuccess: settleResponse, onFailure: settleAndClose } + { response: res, onSuccess: acknowledgeTerminalSend, onFailure: settleAndClose } ); log("info", "ws_terminal_event_sent", { type: "response.completed", source: "json" }); } @@ -654,6 +722,11 @@ async function forwardToInternalHttp( }); res.on("close", () => { if (responseSettled) return; + if (responseBodyEnded || res.complete) { + responseBodyEnded = true; + settleResponse(); + return; + } sendFatalError( "internal_response_closed", "Internal response closed before a complete JSON body was received", @@ -672,11 +745,16 @@ async function forwardToInternalHttp( const EVENT_DELIMITER = /\r?\n\r?\n/; const failIfUnsettled = (code, message, closeReason) => { if (responseSettled) return; - if (!sawTerminal) { - sendFatalError(code, message, 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 = () => { @@ -702,7 +780,7 @@ async function forwardToInternalHttp( // clean terminal event. safeSend(ws, { type: "response.completed", response: null }, { response: res, - onSuccess: settleResponse, + onSuccess: acknowledgeTerminalSend, onFailure: settleAndClose, }); sawTerminal = true; @@ -724,7 +802,7 @@ async function forwardToInternalHttp( event && typeof event.type === "string" && TERMINAL_EVENT_TYPES.has(event.type); safeSend(ws, event, { response: res, - onSuccess: isTerminalEvent ? settleResponse : undefined, + onSuccess: isTerminalEvent ? acknowledgeTerminalSend : undefined, onFailure: settleAndClose, }); if (isTerminalEvent) { @@ -743,6 +821,7 @@ async function forwardToInternalHttp( }); res.on("end", () => { if (responseSettled) return; + responseBodyEnded = true; // Flush any remaining buffered event if (buffer.trim().length > 0) { buffer += "\n\n"; @@ -760,6 +839,7 @@ async function forwardToInternalHttp( // response.create. Do not close here; only fatal transport/protocol // errors initiate a close handshake. log("info", "ws_turn_completed", { terminalEventType }); + settleResponse(); } if (!sawTerminal) return; }); @@ -771,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", @@ -780,25 +865,42 @@ 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( - ws, - "internal_request_error", - String(err && err.message ? err.message : err) - ); - initiateClose(1011, "internal_request_error"); - } - resolve(); - }); - if (typeof registerInternalReq === "function") { - registerInternalReq(req); + 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; @@ -827,15 +929,56 @@ async function forwardToInternalHttp( if (requestBodyFinished) return; requestBodyFinished = true; clearRequestBodyListeners(); + if (!turnFinished) { + turnFinished = true; + const finish = () => { + initiateClose(1011, "internal_request_body_closed"); + resolve(); + }; + const sent = safeSend( + ws, + { + type: "error", + error: { + code: "internal_request_body_closed", + message: "Internal request body closed before it was fully written", + }, + }, + { onSuccess: finish, onFailure: finish } + ); + if (!sent) finish(); + return; + } resolve(); }; const expireRequestBody = () => { if (requestBodyFinished) return; requestBodyFinished = true; clearRequestBodyListeners(); + turnFinished = true; if (!req.destroyed) req.destroy(); - initiateClose(1011, "internal_request_drain_timeout"); - resolve(); + 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(); diff --git a/src/app/v1/_lib/proxy/demand-driven-response-pump.ts b/src/app/v1/_lib/proxy/demand-driven-response-pump.ts index c6ff5b6bd..137db3dd9 100644 --- a/src/app/v1/_lib/proxy/demand-driven-response-pump.ts +++ b/src/app/v1/_lib/proxy/demand-driven-response-pump.ts @@ -121,13 +121,13 @@ export function createDemandDrivenResponsePump( let scheduleDrain = () => {}; - const startDrain = (_reason?: unknown) => { + const startDrain = (_reason?: unknown, markClientAborted = true) => { if (settled || state === "finalizing" || state === "closed") return; if (state === "draining") { scheduleDrain(); return; } - clientAborted = true; + if (markClientAborted) clientAborted = true; state = "draining"; try { clientController?.error( @@ -153,7 +153,7 @@ export function createDemandDrivenResponsePump( `Client response body was not consumed within ${PENDING_CHUNK_DEADLINE_MS}ms`, "AbortError" ); - startDrain(error); + startDrain(error, false); cancelSource(error); }, PENDING_CHUNK_DEADLINE_MS); }; diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index e569ebe14..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,10 +38,8 @@ import { import { addMessageRequestHedgeLoserCost, updateMessageRequestCostWithBreakdown, - updateMessageRequestDetails, updateMessageRequestDetailsDurably, updateMessageRequestDetailsIfUnfinalized, - updateMessageRequestDuration, updateMessageRequestWinnerCost, } from "@/repository/message"; import type { HedgeLoserBilling, StoredCostBreakdown } from "@/types/cost-breakdown"; @@ -132,34 +131,55 @@ async function persistNonStreamTerminalDetails(options: { messageRequestId: number; durationMs: number; details: MessageRequestTerminalDetails; -}): Promise { + onCommitted?: () => void | Promise; +}): Promise { const completeTerminalDetails = { ...options.details, durationMs: options.durationMs, }; try { - await updateMessageRequestDetailsDurably(options.messageRequestId, completeTerminalDetails); - return; + 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: primaryError, + error: + databaseError?.message ?? + (primaryError instanceof Error ? primaryError.message : String(primaryError)), + errorCode: databaseError?.code, + errorPool: databaseError?.pool, }); } try { - await updateMessageRequestDetailsIfUnfinalized( - options.messageRequestId, - completeTerminalDetails - ); + 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, - fallbackError, + error: + databaseError?.message ?? + (fallbackError instanceof Error ? fallbackError.message : String(fallbackError)), + errorCode: databaseError?.code, + errorPool: databaseError?.pool, }); throw markNonStreamTerminalPersistenceError(fallbackError); } @@ -196,40 +216,128 @@ function raceWithDeadline( 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: () => Promise; -}): void { + commit: (signal: AbortSignal) => Promise; +}): Promise { const effectTaskId = `${options.taskId}-post-terminal-effects`; - AsyncTaskManager.register( - effectTaskId, - async () => { + 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()); + commitPromise = Promise.resolve(options.commit(signal)); } catch (error) { commitPromise = Promise.reject(error); } - await raceWithTimeout( - commitPromise, - STREAM_FINALIZATION_MAX_MS, - "post_terminal_side_effect_timeout" - ).catch((error) => { - logger.warn("[ResponseHandler] Post-terminal side effects did not complete", { + 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), }); - }); - }, - { - taskType: "post-terminal-side-effects", - staleTimeoutMs: STREAM_FINALIZATION_MAX_MS, + } 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); + } + }) ); } @@ -1038,12 +1146,12 @@ function finalizeDeferredStreamingFinalizationIfNeeded( ): 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; @@ -1058,12 +1166,7 @@ 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; } @@ -1108,9 +1211,16 @@ 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; @@ -1126,7 +1236,7 @@ function finalizeDeferredStreamingFinalizationIfNeeded( } const shouldClearSessionBindingOnFailure = - (!streamEndedNormally && !clientAbortCompleteSuccess) || + ((clientAborted || !streamEndedNormally) && !clientAbortCompleteSuccess) || detected.isError || (upstreamStatusCode >= 400 && errorMessage !== null); @@ -1174,7 +1284,7 @@ function finalizeDeferredStreamingFinalizationIfNeeded( // 同时，为了让故障转移/熔断能正确工作： // - 客户端主动中断：不计入熔断器（这通常不是供应商问题） // - 非客户端中断：计入 provider/endpoint 熔断失败（与 timeout 路径保持一致） - if (!streamEndedNormally && !clientAbortCompleteSuccess) { + if ((clientAborted || !streamEndedNormally) && !clientAbortCompleteSuccess) { session.addProviderToChain(providerForChain, { endpointId: meta.endpointId, endpointUrl: meta.endpointUrl, @@ -1631,9 +1741,20 @@ export class ProxyResponseHandler { // 使用共享的统计处理方法 const duration = Date.now() - session.startTime; - if (messageContext) { - await updateMessageRequestDuration(messageContext.id, duration); - } + 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, @@ -1641,18 +1762,9 @@ export class ProxyResponseHandler { duration, errorMessageForFinalize, undefined, - false // Gemini 非流式透传 + false, // Gemini 非流式透传 + scheduleProviderFailure ); - - if (commitProviderFailure) { - schedulePostTerminalSideEffects({ - taskId, - providerId: provider.id, - sessionId: session.sessionId, - commit: commitProviderFailure, - }); - } - emitProxyLangfuseTrace(session, { responseHeaders: response.headers, responseText, @@ -1690,29 +1802,11 @@ export class ProxyResponseHandler { }); } - if (messageContext) { - const duration = Date.now() - session.startTime; - await updateMessageRequestDuration(messageContext.id, duration); - await updateMessageRequestDetailsDurably(messageContext.id, { - 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, - }); - const tracker = ProxyStatusTracker.getInstance(); - tracker.endRequest(messageContext.user.id, messageContext.id); - } - const postTerminalSideEffects: Array<() => Promise> = []; if (session.sessionId) { const sessionId = session.sessionId; postTerminalSideEffects.push(async () => { - await SessionManager.clearSessionProvider(sessionId); + await SessionManager.clearSessionProvider(sessionId, provider.id); }); } if ( @@ -1735,15 +1829,42 @@ export class ProxyResponseHandler { } }); } - if (postTerminalSideEffects.length > 0) { - schedulePostTerminalSideEffects({ + let postTerminalSideEffectsScheduled = false; + const scheduleCommittedSideEffects = () => { + if (postTerminalSideEffects.length === 0 || postTerminalSideEffectsScheduled) return; + postTerminalSideEffectsScheduled = true; + return schedulePostTerminalSideEffects({ taskId, providerId: provider.id, sessionId: session.sessionId, - commit: async () => { - await Promise.all(postTerminalSideEffects.map((effect) => effect())); - }, + 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(); @@ -1851,36 +1972,11 @@ export class ProxyResponseHandler { options.statusCode ?? (session.clientAbortSignal?.aborted ? 499 : statusCode); const errorDetails = options.error === undefined ? undefined : buildProcessingErrorDetails(options.error); - 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, - }); - } finally { - tracker.endRequest(messageContext.user.id, messageContext.id); - } - } - const postTerminalSideEffects = [...(options.postTerminalSideEffects ?? [])]; if (session.sessionId) { const sessionId = session.sessionId; postTerminalSideEffects.push(async () => { - await SessionManager.clearSessionProvider(sessionId); + await SessionManager.clearSessionProvider(sessionId, provider.id); const sessionUsagePayload: SessionUsageUpdate = { status: @@ -1900,16 +1996,41 @@ export class ProxyResponseHandler { } }); } - - if (postTerminalSideEffects.length > 0) { - schedulePostTerminalSideEffects({ + let postTerminalSideEffectsScheduled = false; + const scheduleCommittedSideEffects = () => { + if (postTerminalSideEffects.length === 0 || postTerminalSideEffectsScheduled) return; + postTerminalSideEffectsScheduled = true; + return schedulePostTerminalSideEffects({ taskId, providerId: provider.id, sessionId: session.sessionId, - commit: async () => { - await Promise.all(postTerminalSideEffects.map((effect) => effect())); - }, + 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); + } } }; @@ -2181,6 +2302,18 @@ export class ProxyResponseHandler { }); } + 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 = { @@ -2213,23 +2346,13 @@ export class ProxyResponseHandler { messageRequestId: messageContext.id, durationMs: duration, details: terminalDetails, + onCommitted: scheduleCommittedSideEffects, }); } finally { tracker.endRequest(messageContext.user.id, messageContext.id); } } - if (postTerminalSideEffects.length > 0) { - schedulePostTerminalSideEffects({ - taskId, - providerId: provider.id, - sessionId: session.sessionId, - commit: async () => { - await Promise.all(postTerminalSideEffects.map((effect) => effect())); - }, - }); - } - logger.debug("ResponseHandler: Non-stream response processed", { taskId, providerId: provider.id, @@ -2436,22 +2559,41 @@ export class ProxyResponseHandler { 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) => { - passthroughPump.startDrain(reason); - passthroughPump.cancelSource(reason); + startPassthroughDrain(reason); }, }); const cleanupPassthroughClientAbortListener = bindClientAbortListener( session.clientAbortSignal, () => { const reason = session.clientAbortSignal?.reason; - passthroughPump.startDrain(reason); - passthroughPump.cancelSource(reason); + startPassthroughDrain(reason); } ); const statusCode = response.status; @@ -2483,6 +2625,21 @@ export class ProxyResponseHandler { 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 = @@ -2523,6 +2680,19 @@ export class ProxyResponseHandler { 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; @@ -2553,6 +2723,7 @@ export class ProxyResponseHandler { }; observePassthroughChunk = (value) => { + clearIdleTimer(); if (isFirstChunk) { isFirstChunk = false; session.recordTtfb(); @@ -2565,6 +2736,7 @@ export class ProxyResponseHandler { const releaseTransportResources = () => { if (transportReleased) return; transportReleased = true; + clearPassthroughDrainTimeout(); cleanupPassthroughClientAbortListener(); cleanupTaskAbortBinding(); clearIdleTimer(); @@ -2662,6 +2834,7 @@ export class ProxyResponseHandler { clientAborted, abortReason ); + latestCommitSideEffects = finalized.commitSideEffects; const finalizedUsage = await finalizeRequestStats( session, allContent, @@ -2669,7 +2842,8 @@ 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, @@ -2681,14 +2855,6 @@ export class ProxyResponseHandler { isStreaming: true, errorMessage: finalized.errorMessage ?? undefined, }); - if (finalized.commitSideEffects) { - schedulePostTerminalSideEffects({ - taskId, - providerId: provider.id, - sessionId: session.sessionId, - commit: finalized.commitSideEffects, - }); - } } catch (error) { const err = error instanceof Error ? error : new Error(String(error)); const clientAborted = @@ -2738,26 +2904,19 @@ 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) ); - if (finalized.commitSideEffects) { - schedulePostTerminalSideEffects({ - taskId, - providerId: provider.id, - sessionId: session.sessionId, - commit: finalized.commitSideEffects, - }); - } } catch (finalizeError) { - const persistenceDeadlineAtMs = Date.now() + STREAM_FAILURE_PERSISTENCE_MAX_MS; const fallbackStatusCode = statusCode >= 400 ? statusCode @@ -2774,12 +2933,14 @@ export class ProxyResponseHandler { taskId, phase: "stream", detailsWriter: updateMessageRequestDetailsIfUnfinalized, + onCommitted: () => scheduleCommitSideEffects(latestCommitSideEffects), awaitPersistence: (promise: Promise) => - raceWithDeadline( + awaitTerminalPersistenceWithOwnership({ promise, - persistenceDeadlineAtMs, - "stream_failure_persistence_timeout" - ), + taskId, + operation: "gemini-stream-fallback", + timeoutMs: STREAM_FAILURE_PERSISTENCE_MAX_MS, + }), }); } } finally { @@ -3066,31 +3227,47 @@ export class ProxyResponseHandler { }; 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) { - const persistenceDeadlineAtMs = Date.now() + STREAM_FAILURE_PERSISTENCE_MAX_MS; streamFailurePersistencePromise = persistRequestFailure({ ...options, detailsWriter: updateMessageRequestDetailsIfUnfinalized, + onCommitted: options.onCommitted ?? scheduleStreamCommitSideEffects, awaitPersistence: (promise: Promise) => - raceWithDeadline( + awaitTerminalPersistenceWithOwnership({ promise, - persistenceDeadlineAtMs, - "stream_failure_persistence_timeout" - ), - }).catch((error) => { - logger.error("ResponseHandler: Stream failure fallback threw", { - taskId, - messageId: messageContext.id, - error, + 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; + return streamFailurePersistencePromise ?? Promise.resolve(); }; let streamFinalizationPromise: Promise | null = null; @@ -3113,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; @@ -3166,7 +3346,6 @@ export class ProxyResponseHandler { } const duration = Date.now() - session.startTime; - await awaitFinalization(updateMessageRequestDuration(messageContext.id, duration)); const tracker = ProxyStatusTracker.getInstance(); tracker.endRequest(messageContext.user.id, messageContext.id); @@ -3394,35 +3573,7 @@ export class ProxyResponseHandler { ? anthropicModelDetection.actualResponseModel : extractActualResponseModelForProvider(provider.providerType, true, allContent); - // 保存扩展信息（status code, tokens, provider chain） - await awaitFinalization( - updateMessageRequestDetailsDurably(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 = true; - - const postTerminalSideEffects: Array<() => Promise> = []; - const commitDeferredSideEffects = finalized.commitSideEffects; - if (commitDeferredSideEffects) { - postTerminalSideEffects.push(commitDeferredSideEffects); - } + const postTerminalSideEffects = [...latestStreamCommitSideEffects]; if (codexCacheBinding) { const { sessionId, promptCacheKey, providerId, keyId } = codexCacheBinding; postTerminalSideEffects.push(async () => { @@ -3438,16 +3589,37 @@ export class ProxyResponseHandler { } }); } - if (postTerminalSideEffects.length > 0) { - schedulePostTerminalSideEffects({ - taskId, - providerId: providerIdForPersistence ?? provider.id, - sessionId: session.sessionId, - commit: async () => { - await Promise.all(postTerminalSideEffects.map((effect) => effect())); + latestStreamCommitSideEffects = postTerminalSideEffects; + + // 保存扩展信息（status code, tokens, provider chain） + 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, @@ -4881,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) { @@ -4953,8 +5126,9 @@ export async function finalizeRequestStats( }); } - await updateMessageRequestDetailsDurably(messageContext.id, { + const terminalDetails = { statusCode: statusCode, + durationMs: duration, ...(errorMessage ? { errorMessage } : {}), ttfbMs: session.ttfbMs ?? duration, providerChain: session.getProviderChain(), @@ -4968,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; } @@ -5056,8 +5237,9 @@ export async function finalizeRequestStats( } // 7. 更新请求详情 - await updateMessageRequestDetailsDurably(messageContext.id, { + const terminalDetails = { statusCode: statusCode, + durationMs: duration, inputTokens: normalizedUsage.input_tokens, outputTokens: normalizedUsage.output_tokens, ttfbMs: session.ttfbMs ?? duration, @@ -5078,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()) { @@ -5234,6 +5421,11 @@ function buildProcessingErrorDetails(error: unknown): { errorStack?: string; errorCause?: string; } { + const databaseError = findSafeDatabaseError(error); + if (databaseError) { + return { errorMessage: databaseError.message }; + } + const maxErrorStackLength = 8192; const maxErrorCauseLength = 4096; const errorMessage = formatProcessingError(error); @@ -5272,44 +5464,47 @@ async function persistRequestFailure(options: { taskId: string; phase: "stream" | "non-stream"; awaitPersistence?: (promise: Promise) => Promise; - detailsWriter?: typeof updateMessageRequestDetails; -}): 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 ?? updateMessageRequestDetails; + 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, errorStack, errorCause } = buildProcessingErrorDetails(error); const duration = Date.now() - session.startTime; + let committed = false; try { - // 更新请求持续时间 - await awaitPersistence(updateMessageRequestDuration(messageContext.id, duration)); - - // 更新错误详情和 provider chain - await awaitPersistence( - detailsWriter(messageContext.id, { - statusCode, - errorMessage, - errorStack, - errorCause, - ttfbMs: phase === "non-stream" ? (session.ttfbMs ?? duration) : session.ttfbMs, - providerChain: session.getProviderChain(), - model: session.getCurrentModel() ?? undefined, - providerId: session.provider?.id, // 更新最终供应商ID（重试切换后） - context1mApplied: session.getContext1mApplied(), - swapCacheTtlApplied: session.provider?.swapCacheTtlBilling ?? false, - specialSettings: session.getSpecialSettings() ?? undefined, - }) - ); + // duration 与 terminal status 必须属于同一个 CAS patch，避免 ordinary + // metadata 在 overflow 或进程退出时丢失而留下永久 active 记录。 + const terminalDetails = { + statusCode, + durationMs: duration, + errorMessage, + errorStack, + errorCause, + ttfbMs: phase === "non-stream" ? (session.ttfbMs ?? duration) : session.ttfbMs, + providerChain: session.getProviderChain(), + model: session.getCurrentModel() ?? undefined, + providerId: session.provider?.id, // 更新最终供应商ID（重试切换后） + 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()) { @@ -5332,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 { // 确保无论数据库操作成功与否，都清理追踪状态 @@ -5364,6 +5563,7 @@ async function persistRequestFailure(options: { sseEventCount: phase === "stream" ? 0 : undefined, errorMessage, }); + return committed; } /** diff --git a/src/repository/message-write-buffer.ts b/src/repository/message-write-buffer.ts index 8ff18ead0..e58f871e8 100644 --- a/src/repository/message-write-buffer.ts +++ b/src/repository/message-write-buffer.ts @@ -2,6 +2,7 @@ import "server-only"; import type { SQL } from "drizzle-orm"; import { sql } from "drizzle-orm"; +import { findSafeDatabaseError } from "@/drizzle/admitted-client"; import { getMessageWriterDb } from "@/drizzle/db"; import { getEnvConfig } from "@/lib/config/env.schema"; import { logger } from "@/lib/logger"; @@ -404,6 +405,7 @@ class MessageRequestWriteBuffer { private overflowLastDroppedId: number | undefined; private flushAgainAfterCurrent = false; private flushInFlight: Promise | null = null; + private readonly commitCallbacksInFlight = new Set>(); private stopping = false; constructor(config: WriterConfig) { @@ -423,28 +425,17 @@ class MessageRequestWriteBuffer { id: number, patch: MessageRequestUpdatePatch, options: DurableMessageRequestUpdateOptions = {} - ): Promise { + ): Promise { if (this.stopping) { return Promise.reject(new Error("message_request writer is stopping")); } const activeAcknowledgement = this.durableAcknowledgements.get(id); if (activeAcknowledgement && !activeAcknowledgement.settled) { - if (activeAcknowledgement.state === "pending") { - const existing = this.pending.get(id); - if (existing?.durableAcknowledgement === activeAcknowledgement) { - if (options.onCommitted) { - activeAcknowledgement.onCommittedCallbacks.add(options.onCommitted); - } - this.setPending(id, mergePatch(existing.patch, patch), existing.durableAcknowledgement); - this.scheduleFlushIfNeeded(); - return activeAcknowledgement.promise; - } - } - - return Promise.reject( - new Error(`durable message_request update already in flight for id ${id}`) - ); + // 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) { @@ -461,11 +452,11 @@ class MessageRequestWriteBuffer { acknowledgement, new Error("durable message_request queue is full") ); - return acknowledgement.promise; + return acknowledgement.promise.then(() => true); } this.scheduleFlushIfNeeded(); - return acknowledgement.promise; + return acknowledgement.promise.then(() => true); } private createDurableAcknowledgement( @@ -520,12 +511,18 @@ class MessageRequestWriteBuffer { try { const result = callback(patch); if (result && typeof result.then === "function") { - void result.catch((error: unknown) => { - logger.error("[MessageRequestWriteBuffer] Durable commit callback failed", { - error: error instanceof Error ? error.message : String(error), - messageRequestId: acknowledgement.id, + 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", { @@ -789,8 +786,13 @@ class MessageRequestWriteBuffer { } 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, }); @@ -833,6 +835,9 @@ class MessageRequestWriteBuffer { 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(); @@ -870,7 +875,7 @@ export function enqueueMessageRequestUpdateDurably( id: number, patch: MessageRequestUpdatePatch, options?: DurableMessageRequestUpdateOptions -): Promise { +): Promise { if (getEnvConfig().MESSAGE_REQUEST_WRITE_MODE !== "async") { return Promise.reject( new Error("durable message_request buffer API requires async write mode") diff --git a/src/repository/message.ts b/src/repository/message.ts index a779e0f1d..4d72a6944 100644 --- a/src/repository/message.ts +++ b/src/repository/message.ts @@ -143,7 +143,7 @@ async function readPublicStatusRequestSeedFallback( function queuePublicStatusRollupForFinalDetails( id: number, details: PublicStatusFinalDetails -): void { +): Promise | undefined { if (!isPublicStatusFinalDetails(details) || !markPublicStatusRequestInFlight(id)) { return; } @@ -152,7 +152,7 @@ function queuePublicStatusRollupForFinalDetails( return; } - void (async () => { + return (async () => { try { const seed = peekPublicStatusRequestSeed(id) ?? (await readPublicStatusRequestSeedFallback(id)); @@ -210,12 +210,12 @@ function queuePublicStatusRollupForFinalDetails( function publishCommittedMessageRequestDetails( id: number, details: PublicStatusFinalDetails -): void { +): Promise | undefined { if (details.durationMs !== undefined) { updatePublicStatusRequestSeed(id, { durationMs: details.durationMs }); } if (details.providerChain !== undefined && details.statusCode !== undefined) { - queuePublicStatusRollupForFinalDetails(id, details); + return queuePublicStatusRollupForFinalDetails(id, details); } } @@ -505,12 +505,17 @@ export type MessageRequestDetailsUpdate = { export async function updateMessageRequestDetails( id: number, details: MessageRequestDetailsUpdate, - options: { onlyIfUnfinalized?: boolean } = {} -): Promise { + 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; + } enqueueMessageRequestUpdate(id, details); - publishCommittedMessageRequestDetails(id, details); - return; + return true; } const updateData: Record = { @@ -579,25 +584,56 @@ export async function updateMessageRequestDetails( } if (options.onlyIfUnfinalized) { - const updated = await getMessageWriterDb() + 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; + return false; } } else { await db.update(messageRequest).set(updateData).where(eq(messageRequest.id, id)); } - publishCommittedMessageRequestDetails(id, details); + const rollupPromise = publishCommittedMessageRequestDetails(id, details); + if (options.awaitCommitObservers === false) { + void rollupPromise; + } else { + await rollupPromise; + } + return true; } export async function updateMessageRequestDetailsIfUnfinalized( id: number, - details: MessageRequestDetailsUpdate -): Promise { - await updateMessageRequestDetails(id, details, { onlyIfUnfinalized: true }); + 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; } /** @@ -608,25 +644,50 @@ export async function updateMessageRequestDetailsDurably( id: number, details: MessageRequestDetailsUpdate, options?: DurableMessageRequestUpdateOptions -): Promise { +): Promise { if (getEnvConfig().MESSAGE_REQUEST_WRITE_MODE !== "async") { - await updateMessageRequestDetails(id, details); - return; + 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; - publishCommittedMessageRequestDetails(id, committedPatch); - return options?.onCommitted?.(committedPatch); + const rollupPromise = publishCommittedMessageRequestDetails(id, committedPatch); + const callbackResult = options?.onCommitted?.(committedPatch); + if (rollupPromise && callbackResult) { + return Promise.all([rollupPromise, callbackResult]).then(() => undefined); + } + return callbackResult ?? rollupPromise; }; - await enqueueMessageRequestUpdateDurably(id, details, { + const committed = await enqueueMessageRequestUpdateDurably(id, details, { ...options, onCommitted: publishCommit, }); - publishCommit(details); + return committed; } /** diff --git a/tests/integration/message-write-buffer-recovery-postgres.test.ts b/tests/integration/message-write-buffer-recovery-postgres.test.ts index 80bffd4e9..2686fd379 100644 --- a/tests/integration/message-write-buffer-recovery-postgres.test.ts +++ b/tests/integration/message-write-buffer-recovery-postgres.test.ts @@ -129,9 +129,17 @@ run.sequential("message write buffer PostgreSQL recovery", () => { }); const flush = writeBuffer.flushMessageRequestWriteBuffer(); - await expect( - Promise.race([durable.then(() => "settled"), Promise.resolve("pending")]) - ).resolves.toBe("pending"); + 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 }; }); diff --git a/tests/unit/proxy/connected-non-reader-lifetime.test.ts b/tests/unit/proxy/connected-non-reader-lifetime.test.ts index 85e8876fe..f41e2a7c6 100644 --- a/tests/unit/proxy/connected-non-reader-lifetime.test.ts +++ b/tests/unit/proxy/connected-non-reader-lifetime.test.ts @@ -122,7 +122,7 @@ describe("connected non-reader response lifetime", () => { expect(cancel).toHaveBeenCalledOnce(); expect(cancel).toHaveBeenCalledWith(deadlineError); - expect(completion).toMatchObject({ streamEndedNormally: false, clientAborted: true }); + expect(completion).toMatchObject({ streamEndedNormally: false, clientAborted: false }); expect(completion.error).toEqual(deadlineError); expect(completions).toEqual([completion]); expect(pump.getState()).toBe("closed"); 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 0f0beb6ad..ee817d348 100644 --- a/tests/unit/proxy/response-handler-client-abort-drain.test.ts +++ b/tests/unit/proxy/response-handler-client-abort-drain.test.ts @@ -969,6 +969,18 @@ async function expectTaskToResolveWithoutWaiting(promise: Promise): Promis 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) { @@ -995,6 +1007,28 @@ describe("ProxyResponseHandler stream client abort finalization", () => { 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 () => { @@ -1017,9 +1051,10 @@ describe("ProxyResponseHandler stream client abort finalization", () => { it("keeps shutdown pending until generic non-stream recovery persistence settles", async () => { const recoveryStarted = createDeferred(); const releaseRecovery = createDeferred(); - vi.mocked(updateMessageRequestDuration).mockImplementationOnce(async () => { + vi.mocked(updateMessageRequestDetailsIfUnfinalized).mockImplementationOnce(async () => { recoveryStarted.resolve(); await releaseRecovery.promise; + return true; }); let shutdownPromise: Promise | undefined; @@ -1038,14 +1073,16 @@ describe("ProxyResponseHandler stream client abort finalization", () => { await expectPromiseToRemainPending(shutdownPromise); expect(updateMessageRequestDetails).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); releaseRecovery.resolve(); await shutdownPromise; await expect(processingTask).rejects.toThrow("error classification failed"); - expect(updateMessageRequestDetails).toHaveBeenCalledWith( + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledWith( 123, expect.objectContaining({ + durationMs: expect.any(Number), statusCode: 500, errorMessage: "Error: error classification failed", }) @@ -1060,9 +1097,10 @@ describe("ProxyResponseHandler stream client abort finalization", () => { it("keeps shutdown pending until generic stream recovery persistence settles", async () => { const recoveryStarted = createDeferred(); const releaseRecovery = createDeferred(); - vi.mocked(updateMessageRequestDuration).mockImplementationOnce(async () => { + vi.mocked(updateMessageRequestDetailsIfUnfinalized).mockImplementationOnce(async () => { recoveryStarted.resolve(); await releaseRecovery.promise; + return true; }); let downstreamRead: Promise | undefined; let shutdownPromise: Promise | undefined; @@ -1082,7 +1120,7 @@ describe("ProxyResponseHandler stream client abort finalization", () => { shutdownPromise = shutdownAllAsyncTasks(); await expectPromiseToRemainPending(shutdownPromise); - expect(updateMessageRequestDetailsIfUnfinalized).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); releaseRecovery.resolve(); await shutdownPromise; @@ -1092,9 +1130,11 @@ describe("ProxyResponseHandler stream client abort finalization", () => { 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(); @@ -1182,7 +1222,12 @@ describe("ProxyResponseHandler stream client abort finalization", () => { await reader?.cancel("test cleanup"); await drainAsyncTasks(); - expect(updateMessageRequestDuration).toHaveBeenCalledWith(123, expect.any(Number)); + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ durationMs: expect.any(Number) }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); }); it("does not count a pending chunk for an active slow consumer as Provider idle", async () => { @@ -1245,7 +1290,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { statusCode: 200, inputTokens: 463, outputTokens: 11, - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); expect(recordFailure).not.toHaveBeenCalled(); } finally { @@ -1371,7 +1417,7 @@ describe("ProxyResponseHandler stream client abort finalization", () => { await drainAsyncTasks(); expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); - expect(updateMessageRequestDuration).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); expect(emitProxyLangfuseTrace).toHaveBeenCalledTimes(1); }); @@ -1420,7 +1466,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect.objectContaining({ statusCode: 200, providerId: 1, - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); expect(emitProxyLangfuseTrace).toHaveBeenCalledTimes(1); expect( @@ -1483,7 +1530,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { statusCode: 502, errorMessage: "STREAM_UPSTREAM_ABORTED", providerId: 1, - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); expect( vi.mocked(AsyncTaskManager.register).mock.calls.some((call) => { @@ -1504,8 +1552,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { const removeSpy = vi.spyOn(clientController.signal, "removeEventListener"); const responseController = new AbortController(); let resolvePersistence!: () => void; - const blockedPersistence = new Promise((resolve) => { - resolvePersistence = resolve; + const blockedPersistence = new Promise((resolve) => { + resolvePersistence = () => resolve(true); }); let markPersistenceStarted!: () => void; const persistenceStarted = new Promise((resolve) => { @@ -1533,7 +1581,7 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); - vi.mocked(updateMessageRequestDuration).mockImplementationOnce(() => { + vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce(() => { markPersistenceStarted(); return blockedPersistence; }); @@ -1545,7 +1593,7 @@ describe("ProxyResponseHandler stream client abort finalization", () => { await downstream.text(); await persistenceStarted; - expect(updateMessageRequestDuration).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); expect(updateMessageRequestDetailsIfUnfinalized).not.toHaveBeenCalled(); expect(session.clearResponseTimeout).toHaveBeenCalledTimes(1); expect(releaseAgent).toHaveBeenCalledTimes(1); @@ -1568,7 +1616,11 @@ describe("ProxyResponseHandler stream client abort finalization", () => { timersRestored = true; const tasks = asyncTasks.splice(0, asyncTasks.length); await expectAllFulfilled(tasks); - expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ durationMs: expect.any(Number), statusCode: 200 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); } finally { resolvePersistence(); if (!timersRestored) { @@ -1597,8 +1649,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); - vi.mocked(updateMessageRequestDuration).mockImplementationOnce( - () => new Promise(() => {}) + vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce( + () => new Promise(() => {}) ); const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); @@ -1611,9 +1663,11 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledWith( 123, expect.objectContaining({ + durationMs: expect.any(Number), statusCode: 500, errorMessage: "Error: stream_finalization_timeout", - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); } finally { vi.clearAllTimers(); @@ -1621,13 +1675,14 @@ describe("ProxyResponseHandler stream client abort finalization", () => { } }); - it("does not continue a timed-out failure fallback in the background", 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 resolveFallbackDuration!: () => void; + let resolveFallback!: () => void; + let committedCallback: (() => void | Promise) | undefined; try { const clientController = new AbortController(); const session = createSession(clientController.signal); @@ -1643,36 +1698,41 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); - vi.mocked(updateMessageRequestDuration) - .mockImplementationOnce(() => new Promise(() => {})) - .mockImplementationOnce( - () => - new Promise((resolve) => { - resolveFallbackDuration = resolve; - }) - ); + 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; + } + ); const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); await downstream.text(); await vi.advanceTimersByTimeAsync(120_000); await flushMicrotasks(); - expect(updateMessageRequestDuration).toHaveBeenCalledTimes(2); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); await vi.advanceTimersByTimeAsync(5_000); await flushMicrotasks(); - const tasks = asyncTasks.splice(0, asyncTasks.length); - await expectAllFulfilled(tasks); - expect(updateMessageRequestDetailsDurably).not.toHaveBeenCalled(); - expect(updateMessageRequestDetailsIfUnfinalized).not.toHaveBeenCalled(); + expect(asyncTasks.length).toBeGreaterThanOrEqual(2); - resolveFallbackDuration(); + resolveFallback(); await flushMicrotasks(); + expect(committedCallback).toBeTypeOf("function"); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); - expect(updateMessageRequestDetailsDurably).not.toHaveBeenCalled(); - expect(updateMessageRequestDetailsIfUnfinalized).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); } finally { - resolveFallbackDuration?.(); + resolveFallback?.(); await flushMicrotasks(); vi.clearAllTimers(); vi.useRealTimers(); @@ -1681,86 +1741,32 @@ describe("ProxyResponseHandler stream client abort finalization", () => { } }); - it("shares one deadline across sequential stream finalization steps", async () => { - vi.useFakeTimers(); - const flushMicrotasks = (remaining = 32): Promise => - remaining === 0 - ? Promise.resolve() - : Promise.resolve().then(() => flushMicrotasks(remaining - 1)); - let resolveDuration!: () => void; - let resolveTerminalDetails!: () => void; - let markDurationStarted!: () => void; - const durationStarted = new Promise((resolve) => { - markDurationStarted = resolve; - }); - let markTerminalDetailsStarted!: () => void; - const terminalDetailsStarted = new Promise((resolve) => { - markTerminalDetailsStarted = resolve; + 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, }); - 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(updateMessageRequestDuration).mockImplementationOnce(() => { - markDurationStarted(); - return new Promise((resolve) => { - resolveDuration = resolve; - }); - }); - vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce(() => { - markTerminalDetailsStarted(); - return new Promise((resolve) => { - resolveTerminalDetails = resolve; - }); - }); - const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); - await downstream.text(); - await durationStarted; - await vi.advanceTimersByTimeAsync(119_000); - resolveDuration(); - await terminalDetailsStarted; - await flushMicrotasks(); - - expect(updateMessageRequestDuration).toHaveBeenCalledTimes(1); - expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); - expect(updateMessageRequestDetailsIfUnfinalized).not.toHaveBeenCalled(); - - await vi.advanceTimersByTimeAsync(1_000); - await flushMicrotasks(); + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); + await drainAsyncTasks(); - expect(updateMessageRequestDuration).toHaveBeenCalledTimes(2); - expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); - expect(vi.mocked(updateMessageRequestDetailsIfUnfinalized).mock.calls[0]).toEqual([ - 123, - expect.objectContaining({ - statusCode: 500, - errorMessage: "Error: stream_finalization_timeout", - }), - ]); - } finally { - resolveDuration?.(); - resolveTerminalDetails?.(); - await flushMicrotasks(); - vi.clearAllTimers(); - vi.useRealTimers(); - const tasks = asyncTasks.splice(0, asyncTasks.length); - await Promise.race([ - expectAllFulfilled(tasks), - new Promise((resolve) => setTimeout(resolve, 100)), - ]); - } + 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 () => { @@ -1775,10 +1781,10 @@ describe("ProxyResponseHandler stream client abort finalization", () => { const primaryDetailsStarted = new Promise((resolve) => { markPrimaryDetailsStarted = resolve; }); - let resolveFallbackDuration!: () => void; - let markFallbackDurationStarted!: () => void; - const fallbackDurationStarted = new Promise((resolve) => { - markFallbackDurationStarted = resolve; + let resolveFallback!: () => void; + let markFallbackStarted!: () => void; + const fallbackStarted = new Promise((resolve) => { + markFallbackStarted = resolve; }); try { const clientController = new AbortController(); @@ -1795,32 +1801,29 @@ describe("ProxyResponseHandler stream client abort finalization", () => { endpointUrl: "https://api.test.invalid/v1", upstreamStatusCode: 200, }); - vi.mocked(updateMessageRequestDuration) - .mockImplementationOnce(async () => undefined) - .mockImplementationOnce(() => { - markFallbackDurationStarted(); - return new Promise((resolve) => { - resolveFallbackDuration = resolve; - }); - }); vi.mocked(updateMessageRequestDetailsDurably) .mockImplementationOnce((_id, details) => { markPrimaryDetailsStarted(); - return new Promise((resolve) => { + return new Promise((resolve) => { resolvePrimaryDetails = () => { terminalStatusCode = details.statusCode ?? null; - resolve(); + 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); + }); } ); @@ -1828,13 +1831,13 @@ describe("ProxyResponseHandler stream client abort finalization", () => { await downstream.text(); await primaryDetailsStarted; await vi.advanceTimersByTimeAsync(120_000); - await fallbackDurationStarted; + await fallbackStarted; resolvePrimaryDetails(); await flushMicrotasks(); expect(terminalStatusCode).toBe(200); - resolveFallbackDuration(); + resolveFallback(); await flushMicrotasks(); const tasks = asyncTasks.splice(0, asyncTasks.length); await expectAllFulfilled(tasks); @@ -1843,7 +1846,7 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect(terminalStatusCode).toBe(200); } finally { resolvePrimaryDetails?.(); - resolveFallbackDuration?.(); + resolveFallback?.(); await flushMicrotasks(); vi.clearAllTimers(); vi.useRealTimers(); @@ -1894,14 +1897,15 @@ describe("ProxyResponseHandler stream client abort finalization", () => { await drainAsyncTasks(); expect(AsyncTaskManager.cancel).not.toHaveBeenCalled(); - expect(updateMessageRequestDuration).toHaveBeenCalledWith(123, expect.any(Number)); + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( 123, expect.objectContaining({ statusCode: 200, inputTokens: 463, outputTokens: 11, - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); }); @@ -1935,7 +1939,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { statusCode: 200, inputTokens: 463, outputTokens: 11, - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); expect(SessionManager.storeSessionResponse).not.toHaveBeenCalled(); @@ -1980,7 +1985,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { statusCode: 200, inputTokens: 463, outputTokens: 11, - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); const traceCall = vi.mocked(emitProxyLangfuseTrace).mock.calls.at(-1); @@ -2024,7 +2030,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { statusCode: 200, inputTokens: 463, outputTokens: 11, - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); expect(SessionManager.storeSessionResponse).not.toHaveBeenCalled(); }); @@ -2062,7 +2069,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { statusCode: 200, }), ], - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); }); @@ -2092,7 +2100,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect.objectContaining({ statusCode: 499, errorMessage: "CLIENT_ABORTED", - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); }); @@ -2126,7 +2135,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { 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 = ( @@ -2193,7 +2203,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect.objectContaining({ statusCode: 200, inputTokens: 463, - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); const expectedCost = 0.001554; @@ -2277,7 +2288,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect.objectContaining({ statusCode: 499, errorMessage: "CLIENT_ABORTED", - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); } finally { vi.useRealTimers(); @@ -2325,7 +2337,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect.objectContaining({ statusCode: 499, errorMessage: "CLIENT_ABORTED", - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); } finally { vi.useRealTimers(); @@ -2376,7 +2389,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect.objectContaining({ statusCode: 499, errorMessage: "CLIENT_ABORTED", - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); } finally { vi.useRealTimers(); @@ -2425,7 +2439,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect.objectContaining({ statusCode: 502, errorMessage: "STREAM_IDLE_TIMEOUT", - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); expect(recordFailure).toHaveBeenCalledTimes(1); } finally { @@ -2471,7 +2486,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect.objectContaining({ statusCode: 499, errorMessage: "CLIENT_ABORTED", - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); } finally { vi.useRealTimers(); @@ -2516,7 +2532,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect.objectContaining({ statusCode: 499, errorMessage: "CLIENT_ABORTED", - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); } finally { if (!upstreamController.signal.aborted) { @@ -2556,7 +2573,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect.objectContaining({ statusCode: 499, errorMessage: "CLIENT_ABORTED", - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); }); @@ -2587,7 +2605,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect.objectContaining({ statusCode: 502, errorMessage: "STREAM_UPSTREAM_ABORTED", - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); }); @@ -2619,14 +2638,11 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); expect(updateMessageRequestDetails).not.toHaveBeenCalled(); expect(updateMessageRequestDetailsIfUnfinalized).toHaveBeenCalledTimes(1); - expect(updateMessageRequestDuration).toHaveBeenCalledTimes(2); - expect(vi.mocked(updateMessageRequestDuration).mock.calls[1]).toEqual([ - 123, - expect.any(Number), - ]); + 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, @@ -2639,6 +2655,7 @@ describe("ProxyResponseHandler stream client abort finalization", () => { }), ], }), + expect.objectContaining({ onCommitted: expect.any(Function) }), ]); }); @@ -2723,7 +2740,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect.objectContaining({ statusCode: 502, errorMessage: "STREAM_IDLE_TIMEOUT", - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); } finally { vi.useRealTimers(); @@ -2762,7 +2780,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect.objectContaining({ statusCode: 502, errorMessage: "STREAM_RESPONSE_TIMEOUT", - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); upstream.close(); }); @@ -2801,7 +2820,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect.objectContaining({ statusCode: 502, errorMessage: "STREAM_RESPONSE_TIMEOUT", - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); expect(recordFailure).toHaveBeenCalledTimes(1); expect(recordFailure).toHaveBeenCalledWith( @@ -2814,7 +2834,16 @@ describe("ProxyResponseHandler stream client abort finalization", () => { it("waits for durable non-stream failure details before mutating the Provider circuit", async () => { const durableAck = createDeferred(); vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce( - async () => await durableAck.promise + 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"}}', { @@ -2874,7 +2903,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { 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( @@ -2892,7 +2922,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { errorMessage: "FAKE_200_JSON_ERROR_MESSAGE_NON_EMPTY", }), ], - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); }); @@ -2943,7 +2974,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { }), ] : [], - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); } ); @@ -2972,7 +3004,11 @@ describe("ProxyResponseHandler stream client abort finalization", () => { it("waits for durable Gemini non-stream failure details before mutating the Provider circuit", async () => { const durableAck = createDeferred(); vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce( - async () => await durableAck.promise + async (_id, details, options) => { + await durableAck.promise; + publishCommitObserver(details, options); + return true; + } ); const session = createSession(new AbortController().signal, { providerType: "gemini", @@ -3045,7 +3081,7 @@ describe("ProxyResponseHandler stream client abort finalization", () => { } ); - it("persists Gemini non-stream duration before completing terminal stats", async () => { + it("persists Gemini non-stream duration atomically with terminal stats", async () => { const session = createSession(new AbortController().signal, { providerType: "gemini", originalFormat: "gemini", @@ -3060,18 +3096,22 @@ describe("ProxyResponseHandler stream client abort finalization", () => { await ProxyResponseHandler.dispatch(session, response); await drainAsyncTasks(); - expect(updateMessageRequestDuration).toHaveBeenCalledTimes(1); - expect(updateMessageRequestDuration).toHaveBeenCalledWith(123, expect.any(Number)); + expect(updateMessageRequestDuration).not.toHaveBeenCalled(); expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( 123, - expect.objectContaining({ statusCode: 200 }) + 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 () => await durableAck.promise + async (_id, details, options) => { + await durableAck.promise; + publishCommitObserver(details, options); + return true; + } ); const responseController = new AbortController(); const session = createSession(new AbortController().signal); @@ -3109,7 +3149,11 @@ describe("ProxyResponseHandler stream client abort finalization", () => { it("waits for durable non-stream details before updating the Codex cache binding", async () => { const durableAck = createDeferred(); vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce( - async () => await durableAck.promise + 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); @@ -3145,7 +3189,11 @@ describe("ProxyResponseHandler stream client abort finalization", () => { const cacheBinding = createDeferred(); const cacheBindingStarted = createDeferred(); vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce( - async () => await durableAck.promise + 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 () => { @@ -3222,7 +3270,8 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( 123, - expect.objectContaining({ statusCode: 500 }) + expect.objectContaining({ statusCode: 500 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); expect(SessionManager.updateSessionWithCodexCacheKey).not.toHaveBeenCalled(); }); @@ -3249,13 +3298,15 @@ describe("ProxyResponseHandler stream client abort finalization", () => { await ProxyResponseHandler.dispatch(session, response); await drainAsyncTasks(); - expect(updateMessageRequestDuration).toHaveBeenCalledWith(123, expect.any(Number)); + 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 ea5030e99..c3bdf3a5d 100644 --- a/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts +++ b/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts @@ -68,6 +68,7 @@ vi.mock("@/repository/message", () => ({ updateMessageRequestCostWithBreakdown: vi.fn(), updateMessageRequestDetails: vi.fn(), updateMessageRequestDetailsDurably: vi.fn(), + updateMessageRequestDetailsIfUnfinalized: vi.fn(), updateMessageRequestDuration: vi.fn(), })); @@ -356,7 +357,7 @@ function setupCommonMocks() { updatedAt: new Date(), }); vi.mocked(updateMessageRequestDetails).mockResolvedValue(undefined); - vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(true); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(SessionManager.clearSessionProvider).mockResolvedValue(undefined); @@ -408,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( @@ -437,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(); }); @@ -453,7 +454,7 @@ describe("Endpoint circuit breaker isolation", () => { 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( @@ -514,7 +515,8 @@ describe("Endpoint circuit breaker isolation", () => { ); expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( 1, - expect.objectContaining({ statusCode: 200, providerId: 1 }) + expect.objectContaining({ statusCode: 200, providerId: 1 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); }); 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 6d2dee051..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 @@ -35,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) => { @@ -683,7 +694,7 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { } }); - test("客户端中断流式透传后应清理 session provider 绑定，避免下次继续复用旧供应商", async () => { + test("客户端中断流式透传后应继续 drain 并提交尾部 usage", async () => { asyncTasks.length = 0; const { baseUrl, close } = await startSseServer((_req, res) => { res.writeHead(200, { @@ -695,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(); @@ -743,9 +756,16 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { clientAbortController.abort(new Error("client_cancelled")); await expectAllFulfilled(asyncTasks); - expect(vi.mocked(SessionManager.clearSessionProvider)).toHaveBeenCalledWith( - "gemini-abort-session" + 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(); @@ -808,19 +828,23 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { statusCode: 200, inputTokens: 463, outputTokens: 11, - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); expect(SessionManager.storeSessionResponse).not.toHaveBeenCalled(); }); - test("Gemini 终态副作用挂起时应先释放传输资源并在共享 deadline 后收敛", async () => { + 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(() => {}); + return new Promise((resolve) => { + releaseSideEffect = resolve; + }); }); try { @@ -873,8 +897,15 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { 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); - await expectAllFulfilled(asyncTasks.splice(0, asyncTasks.length)); + expect(tasksSettled).toBe(false); + + releaseSideEffect(); + await pendingTasks; expect(releaseAgent).toHaveBeenCalledTimes(1); expect(updateMessageRequestDetailsDurably).toHaveBeenCalledTimes(1); diff --git a/tests/unit/proxy/response-handler-gemini-terminal.test.ts b/tests/unit/proxy/response-handler-gemini-terminal.test.ts index dbac49755..c6d0706e8 100644 --- a/tests/unit/proxy/response-handler-gemini-terminal.test.ts +++ b/tests/unit/proxy/response-handler-gemini-terminal.test.ts @@ -6,10 +6,13 @@ 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) => Promise>(), + 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) => Promise>(), + durable: + vi.fn<(id: number, details: object, options?: TerminalWriterOptions) => Promise>(), tasks: Array.from>([]), })); @@ -118,7 +121,13 @@ const geminiResponse = (body: BodyInit) => async function settleTasks(): Promise { while (mocks.tasks.length > 0) { - await Promise.allSettled(mocks.tasks.splice(0, mocks.tasks.length)); + 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"); + } } } @@ -129,29 +138,40 @@ describe("ProxyResponseHandler.dispatch Gemini terminal behavior", () => { vi.useRealTimers(); mocks.tasks.length = 0; vi.clearAllMocks(); - mocks.conditional.mockResolvedValue(undefined); + mocks.conditional.mockResolvedValue(true); mocks.details.mockResolvedValue(undefined); - mocks.durable.mockResolvedValue(undefined); + mocks.durable.mockResolvedValue(true); }); - it("cancels the Gemini source when the returned body is cancelled", async () => { + it("drains the Gemini source after the returned body is cancelled", async () => { const cancelSource = vi.fn(); - let failSource = () => {}; + let sourceController: ReadableStreamDefaultController | null = null; + const encoder = new TextEncoder(); const source = new ReadableStream({ cancel: cancelSource, start(controller) { - failSource = () => controller.error(new DOMException("closed", "AbortError")); + 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")); - const cancellationCount = cancelSource.mock.calls.length; - failSource(); + sourceController?.enqueue(encoder.encode('{"usageMetadata":{"promptTokenCount":1}}\n')); + sourceController?.close(); await settleTasks(); - expect(cancellationCount).toBe(1); + 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 () => { @@ -192,7 +212,8 @@ describe("ProxyResponseHandler.dispatch Gemini terminal behavior", () => { expect(mocks.conditional).toHaveBeenCalledWith( MESSAGE.id, - expect.objectContaining({ statusCode: 500 }) + expect.objectContaining({ durationMs: expect.any(Number), statusCode: 500 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); }); diff --git a/tests/unit/proxy/response-handler-lease-decrement.test.ts b/tests/unit/proxy/response-handler-lease-decrement.test.ts index 5cf34b13d..dd76fdc05 100644 --- a/tests/unit/proxy/response-handler-lease-decrement.test.ts +++ b/tests/unit/proxy/response-handler-lease-decrement.test.ts @@ -73,6 +73,7 @@ vi.mock("@/repository/message", () => ({ updateMessageRequestCostWithBreakdown: vi.fn(), updateMessageRequestDetails: vi.fn(), updateMessageRequestDetailsDurably: vi.fn(), + updateMessageRequestDetailsIfUnfinalized: vi.fn(), updateMessageRequestDuration: vi.fn(), })); @@ -548,7 +549,7 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { vi.mocked(findLatestPriceByModel).mockResolvedValue( makePriceRecord(originalModel, testPriceData) ); - vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(undefined); + vi.mocked(updateMessageRequestDetailsDurably).mockResolvedValue(true); vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponsePhaseSnapshot).mockResolvedValue(undefined); @@ -630,7 +631,8 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { statusCode: 200, inputTokens: usage.input_tokens, outputTokens: usage.output_tokens, - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); }); @@ -719,7 +721,8 @@ describe("Lease Budget Decrement after trackCostToRedis", () => { statusCode: 200, inputTokens: usage.input_tokens, outputTokens: usage.output_tokens, - }) + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) ); }); diff --git a/tests/unit/proxy/response-handler-nonstream-terminal.test.ts b/tests/unit/proxy/response-handler-nonstream-terminal.test.ts index bededce2c..f72581739 100644 --- a/tests/unit/proxy/response-handler-nonstream-terminal.test.ts +++ b/tests/unit/proxy/response-handler-nonstream-terminal.test.ts @@ -7,10 +7,13 @@ 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) => Promise>(), - durable: vi.fn<(id: number, details: object) => Promise>(), + 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(), @@ -206,8 +209,24 @@ describe("ProxyResponseHandler.dispatch nonstream terminal behavior", () => { beforeEach(() => { mocks.tasks.length = 0; vi.clearAllMocks(); - mocks.conditional.mockResolvedValue(undefined); - mocks.durable.mockResolvedValue(undefined); + 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); }); @@ -229,7 +248,8 @@ describe("ProxyResponseHandler.dispatch nonstream terminal behavior", () => { expect(settlements.every(({ status }) => status === "fulfilled")).toBe(true); expect(mocks.conditional).toHaveBeenCalledWith( 41, - expect.objectContaining({ statusCode: 503 }) + 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( diff --git a/tests/unit/proxy/response-handler-stream-terminal.test.ts b/tests/unit/proxy/response-handler-stream-terminal.test.ts index db6e58ae9..efe5ddd8e 100644 --- a/tests/unit/proxy/response-handler-stream-terminal.test.ts +++ b/tests/unit/proxy/response-handler-stream-terminal.test.ts @@ -9,7 +9,10 @@ import type { User } from "@/types/user"; type TaskOptions = { readonly abortController?: AbortController }; const mocks = vi.hoisted(() => ({ - durable: vi.fn<(id: number, details: object) => Promise>(), + durable: + vi.fn< + (id: number, details: object, options?: { onCommitted?: () => void }) => Promise + >(), tasks: Array.from>([]), trackerEnd: vi.fn(), })); @@ -206,7 +209,7 @@ describe("ProxyResponseHandler.dispatch stream terminal behavior", () => { beforeEach(() => { mocks.tasks.length = 0; vi.clearAllMocks(); - mocks.durable.mockResolvedValue(undefined); + mocks.durable.mockResolvedValue(true); }); it("persists a naturally completed stream and releases its transport", async () => { @@ -219,7 +222,11 @@ describe("ProxyResponseHandler.dispatch stream terminal behavior", () => { await returned.text(); await settleTasks(); - expect(mocks.durable).toHaveBeenCalledWith(51, expect.objectContaining({ statusCode: 200 })); + 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(); }); @@ -242,7 +249,11 @@ describe("ProxyResponseHandler.dispatch stream terminal behavior", () => { abortSource(); await settleTasks(); - expect(mocks.durable).toHaveBeenCalledWith(51, expect.objectContaining({ statusCode: 499 })); + expect(mocks.durable).toHaveBeenCalledWith( + 51, + expect.objectContaining({ statusCode: 499 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); expect(releaseAgent).toHaveBeenCalledOnce(); }); @@ -258,7 +269,11 @@ describe("ProxyResponseHandler.dispatch stream terminal behavior", () => { await expect(bodyRead).rejects.toThrow("response deadline exceeded"); await settleTasks(); - expect(mocks.durable).toHaveBeenCalledWith(51, expect.objectContaining({ statusCode: 502 })); + 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/terminal-outcome-contract.test.ts b/tests/unit/proxy/terminal-outcome-contract.test.ts index e51143087..8a9d1f859 100644 --- a/tests/unit/proxy/terminal-outcome-contract.test.ts +++ b/tests/unit/proxy/terminal-outcome-contract.test.ts @@ -227,9 +227,17 @@ describe("terminal outcome contract", () => { const flushPromise = flushMessageRequestWriteBuffer(); await flushMicrotasks(); - await expect( - Promise.race([handlePromise.then(() => "resolved"), Promise.resolve("pending")]) - ).resolves.toBe("pending"); + 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"); diff --git a/tests/unit/repository/message-public-status-rollup.test.ts b/tests/unit/repository/message-public-status-rollup.test.ts index 0783664ff..864a1af78 100644 --- a/tests/unit/repository/message-public-status-rollup.test.ts +++ b/tests/unit/repository/message-public-status-rollup.test.ts @@ -178,7 +178,7 @@ describe("repository/message public status rollup hook", () => { }); it("writes timeout fallback details only while the request is unfinalized", async () => { - mockWriterDbUpdateReturning.mockResolvedValueOnce([{ id: 606 }]); + mockDbUpdateReturning.mockResolvedValueOnce([{ id: 606 }]); const { updateMessageRequestDetailsIfUnfinalized } = await import("@/repository/message"); await updateMessageRequestDetailsIfUnfinalized(606, { @@ -186,11 +186,10 @@ describe("repository/message public status rollup hook", () => { errorMessage: "Error: stream_finalization_timeout", }); - 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(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"); }); @@ -217,7 +216,7 @@ describe("repository/message public status rollup hook", () => { }); it("does not queue terminal rollup when the timeout fallback loses the terminal CAS", async () => { - mockWriterDbUpdateReturning.mockResolvedValueOnce([]); + mockDbUpdateReturning.mockResolvedValueOnce([]); const { updateMessageRequestDetailsIfUnfinalized } = await import("@/repository/message"); await updateMessageRequestDetailsIfUnfinalized(607, { @@ -227,14 +226,20 @@ describe("repository/message public status rollup hook", () => { }); await flushMicrotasks(); - expect(mockGetMessageWriterDb).toHaveBeenCalledTimes(1); + 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.mockReturnValueOnce(durableAck.promise); + mockEnqueueMessageRequestUpdateDurably.mockImplementationOnce( + async (_id, details, options: DurableMessageRequestUpdateOptions | undefined) => { + await durableAck.promise; + await options?.onCommitted?.(details); + return true; + } + ); mockDbInsertReturning.mockResolvedValueOnce([ { id: 808, @@ -368,6 +373,7 @@ describe("repository/message public status rollup hook", () => { 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"), diff --git a/tests/unit/repository/message-terminal-cas-durable.test.ts b/tests/unit/repository/message-terminal-cas-durable.test.ts index 6257c52ee..f4f26ce8e 100644 --- a/tests/unit/repository/message-terminal-cas-durable.test.ts +++ b/tests/unit/repository/message-terminal-cas-durable.test.ts @@ -130,8 +130,11 @@ function installAsyncDurableBoundaries(execute: (query: SqlQuery) => Promise []); +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(); @@ -182,7 +185,7 @@ describe("message terminal CAS and durable acknowledgement", () => { providerChain: [{ id: 4, name: "fallback", groupTag: "openai" }], }); - expect(result).toBeUndefined(); + expect(result).toBe(false); expect(boundary.writerReturning).toHaveBeenCalledTimes(1); expect(boundary.getRedisClient).not.toHaveBeenCalled(); expect(boundary.defaultUpdate).not.toHaveBeenCalled(); @@ -205,9 +208,17 @@ describe("message terminal CAS and durable acknowledgement", () => { model: "gpt-4.1", }); - await expect( - Promise.race([completion.then(() => "committed"), Promise.resolve("pending")]) - ).resolves.toBe("pending"); + 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(); @@ -223,12 +234,49 @@ describe("message terminal CAS and durable acknowledgement", () => { vi.resetModules(); const { set, writerExecute } = installSyncDurableBoundaries(); const { updateMessageRequestDetailsDurably } = await import("@/repository/message"); + const onCommitted = vi.fn(); - await updateMessageRequestDetailsDurably(804, { durationMs: 900, statusCode: 201 }); + 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-public-status-seam.test.ts b/tests/unit/repository/message-terminal-public-status-seam.test.ts index f539bd675..e684da8ef 100644 --- a/tests/unit/repository/message-terminal-public-status-seam.test.ts +++ b/tests/unit/repository/message-terminal-public-status-seam.test.ts @@ -293,7 +293,7 @@ describe("message terminal public-status public seam", () => { } ); - it("same-ID pending durable merge publishes one rollup from the committed latest payload", async () => { + it("same-ID pending durable contention publishes one rollup from the first owner", async () => { vi.resetModules(); vi.useFakeTimers(); @@ -507,40 +507,41 @@ describe("message terminal public-status public seam", () => { expect(row.statusCode).toBeNull(); releaseCommit.resolve(); - await Promise.all([oldFailure, latestSuccess, flush]); + 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: latestSuccessDetails.durationMs, - statusCode: latestSuccessDetails.statusCode, + durationMs: oldFailureDetails.durationMs, + statusCode: oldFailureDetails.statusCode, inputTokens: oldFailureDetails.inputTokens, - outputTokens: latestSuccessDetails.outputTokens, - ttfbMs: latestSuccessDetails.ttfbMs, - providerChain: latestSuccessDetails.providerChain, - providerId: latestSuccessDetails.providerId, + 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(5); + 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|success": 1, - "42|gpt-4.1|ttfb_sum": latestSuccessDetails.ttfbMs, - "42|gpt-4.1|ttfb_count": 1, - "42|gpt-4.1|tps_sum": 80, - "42|gpt-4.1|tps_count": 1, - }); + expect(rollupIncrements).toEqual({ "42|gpt-4.1|failure": 1 }); await stopMessageRequestWriteBuffer(); }); diff --git a/tests/unit/repository/message-write-buffer.test.ts b/tests/unit/repository/message-write-buffer.test.ts index 707f12c38..cdac4fc6c 100644 --- a/tests/unit/repository/message-write-buffer.test.ts +++ b/tests/unit/repository/message-write-buffer.test.ts @@ -191,9 +191,17 @@ describe("message_request 异步批量写入", () => { expect(executeMock).toHaveBeenCalledTimes(1); const built = toSqlText(executeMock.mock.calls[0]?.[0]); expect(built.sql).toContain("RETURNING id"); - await expect( - Promise.race([durablePromise.then(() => "resolved"), Promise.resolve("pending")]) - ).resolves.toBe("pending"); + 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; @@ -228,7 +236,7 @@ describe("message_request 异步批量写入", () => { deferred.resolve([{ id: 51001 }, { id: 52002 }]); await flushPromise; - await expect(durablePromise).resolves.toBeUndefined(); + await expect(durablePromise).resolves.toBe(true); await stopMessageRequestWriteBuffer(); }); @@ -255,7 +263,47 @@ describe("message_request 异步批量写入", () => { deferred.resolve([{ id: 11 }, { id: 12 }]); await flushPromise; - await expect(Promise.all([first, second])).resolves.toEqual([undefined, undefined]); + 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(); }); @@ -273,12 +321,20 @@ describe("message_request 异步批量写入", () => { const durablePromise = enqueueMessageRequestUpdateDurably(21, { statusCode: 200 }); await flushMessageRequestWriteBuffer(); - await expect( - Promise.race([durablePromise.then(() => "resolved"), Promise.resolve("pending")]) - ).resolves.toBe("pending"); + 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.toBeUndefined(); + await expect(durablePromise).resolves.toBe(true); expect(executeMock).toHaveBeenCalledTimes(2); await stopMessageRequestWriteBuffer(); }); @@ -325,7 +381,7 @@ describe("message_request 异步批量写入", () => { const retry = enqueueMessageRequestUpdateDurably(31, { statusCode: 200 }); await flushMessageRequestWriteBuffer(); - await expect(retry).resolves.toBeUndefined(); + await expect(retry).resolves.toBe(true); await stopMessageRequestWriteBuffer(); }); @@ -354,7 +410,7 @@ describe("message_request 异步批量写入", () => { const retry = enqueueMessageRequestUpdateDurably(311, { statusCode: 502 }); await flushMessageRequestWriteBuffer(); - await expect(retry).resolves.toBeUndefined(); + await expect(retry).resolves.toBe(true); const built = toSqlText(executeMock.mock.calls[0]?.[0]); expect(built.params).not.toContain("stale-primary-generation"); @@ -382,7 +438,7 @@ describe("message_request 异步批量写入", () => { firstExecute.resolve([]); await flushPromise; - await expect(retry).resolves.toBeUndefined(); + await expect(retry).resolves.toBe(true); expect(executeMock).toHaveBeenCalledTimes(2); await stopMessageRequestWriteBuffer(); }); @@ -482,7 +538,7 @@ describe("message_request 异步批量写入", () => { ); await expect(flushMessageRequestWriteBuffer()).resolves.toBeUndefined(); - await expect(durable).resolves.toBeUndefined(); + await expect(durable).resolves.toBe(true); expect(loggerErrorMock).toHaveBeenCalledWith( "[MessageRequestWriteBuffer] Durable commit callback failed", expect.objectContaining({ @@ -493,6 +549,35 @@ describe("message_request 异步批量写入", () => { 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"; @@ -592,22 +677,26 @@ describe("message_request 异步批量写入", () => { throw new Error("executor did not synchronously re-enter stop"); } const samePromise = outerStopPromise === reentrantPromise; - const settlementsBeforeRelease = await Promise.all([ - Promise.race([ - outerStopPromise.then( - () => "fulfilled", - () => "rejected" - ), - Promise.resolve("pending"), - ]), - Promise.race([ - reentrantPromise.then( - () => "fulfilled", - () => "rejected" - ), - Promise.resolve("pending"), - ]), - ]); + 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); @@ -616,7 +705,7 @@ describe("message_request 异步批量写入", () => { } const stopResults = await Promise.allSettled([outerStopPromise, reentrantPromise]); - expect(settlementsBeforeRelease).toEqual(["pending", "pending"]); + expect(settlementsBeforeRelease).toEqual([false, false]); if (shouldReject) { const shutdownError = "message_request writer shutdown persistence failed"; expect(stopResults).toEqual([ @@ -699,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; diff --git a/tests/unit/server-response-write-backpressure.test.ts b/tests/unit/server-response-write-backpressure.test.ts index 938521c20..006c36aca 100644 --- a/tests/unit/server-response-write-backpressure.test.ts +++ b/tests/unit/server-response-write-backpressure.test.ts @@ -22,7 +22,11 @@ type ServerModule = { request: { headers: Record; url: string }, body: Record, sessionId: string, - registerRequest?: (request: http.ClientRequest) => void, + registerRequest?: ( + request: http.ClientRequest, + response?: http.IncomingMessage | null, + settleTurn?: () => boolean + ) => boolean | void, close?: (code: number, reason: string) => void ) => Promise; }; @@ -61,7 +65,11 @@ function createIncomingResponse(): http.IncomingMessage { function requestInput() { return { - ws: { readyState: 1, send: vi.fn(), close: vi.fn() }, + 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" }, }; @@ -111,7 +119,10 @@ async function startSseBridge(send: WebSocketLike["send"]) { return { events, request, response, ws }; } -afterEach(vi.restoreAllMocks); +afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); +}); describe("server response write backpressure", () => { it("waits for request drain before ending a backpressured payload", async () => { @@ -126,6 +137,76 @@ describe("server response write backpressure", () => { 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); @@ -283,4 +364,128 @@ describe("server response write backpressure", () => { 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"); + }); }); From 048caedb62f68b60d1fb5531d1f841307641f6ca Mon Sep 17 00:00:00 2001 From: ding113 Date: Fri, 17 Jul 2026 14:23:33 +0800 Subject: [PATCH 20/20] fix(shutdown): await scheduler quiescence before closing DB and Redis Background schedulers (probe, public-status, log-cleanup) now expose their in-flight promise so shutdown awaits true quiescence instead of racing a fixed timeout. Bull queues are joined before DB pool and Redis closure; failures are collected into an AggregateError so critical cleanup still runs. Ledger backfill registers as a cancellable AsyncTaskManager job with abort-signal checkpoints between batches. The lease settlement Lua script consumes the cached slice when a request exceeds the remaining budget, preventing repeated overshoot within the refresh window. --- src/lib/async-task-manager.ts | 4 - src/lib/ledger-backfill/service.ts | 8 +- src/lib/lifecycle/shutdown.ts | 79 +++++++++++++--- .../provider-endpoints/probe-log-cleanup.ts | 32 +++++-- src/lib/provider-endpoints/probe-scheduler.ts | 21 ++++- src/lib/public-status/scheduler.ts | 22 ++++- src/lib/rate-limit/lease-service.ts | 12 ++- .../lease-settlement-redis.test.ts | 2 +- tests/unit/lib/shutdown.test.ts | 93 +++++++++++++++++++ tests/unit/usage-ledger/backfill.test.ts | 41 +++++++- 10 files changed, 280 insertions(+), 34 deletions(-) diff --git a/src/lib/async-task-manager.ts b/src/lib/async-task-manager.ts index 01c604f2e..4c5ae8c47 100644 --- a/src/lib/async-task-manager.ts +++ b/src/lib/async-task-manager.ts @@ -367,10 +367,6 @@ class AsyncTaskManagerClass { } await Promise.allSettled(activeTasks.map((taskInfo) => taskInfo.promise)); - - for (const taskInfo of activeTasks) { - this.cleanup(taskInfo.taskId, taskInfo); - } } })().then(resolveShutdown, (error) => { this.lifecycleState = "closed"; 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 dffb6c2ea..42aa4f800 100644 --- a/src/lib/lifecycle/shutdown.ts +++ b/src/lib/lifecycle/shutdown.ts @@ -46,6 +46,31 @@ 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; @@ -67,10 +92,11 @@ 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(); @@ -80,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(); @@ -102,18 +128,41 @@ 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 才能把这些尾部更新真正落库。 const asyncTasksWarningTimer = setTimeout(() => { @@ -131,7 +180,7 @@ export async function runApplicationCleanup( clearTimeout(asyncTasksWarningTimer); } - // 6. 刷写 message_request 异步写缓冲。这里不能用可脱离的单步 timeout： + // 7. 刷写 message_request 异步写缓冲。这里不能用可脱离的单步 timeout： // closeDbPools 必须等 writer 真正 settled，否则会关闭仍在执行终态 SQL 的连接。 writerQuiescencePending = true; const writerWarningTimer = setTimeout(() => { @@ -150,7 +199,7 @@ export async function runApplicationCleanup( writerQuiescencePending = false; } - // 7. writer flush 完成后再关闭数据库 pool。pool close 也是 critical barrier， + // 8. writer flush 完成后再关闭数据库 pool。pool close 也是 critical barrier， // 单步 deadline 只能告警，不能让底层 client.end() 脱离 shutdown 生命周期。 const dbWarningTimer = setTimeout(() => { logger.warn("[Shutdown] closeDbPools still pending", { ms: stepMs }); @@ -167,7 +216,7 @@ export async function runApplicationCleanup( clearTimeout(dbWarningTimer); } - // 8. Langfuse 自带超时（LANGFUSE_SHUTDOWN_TIMEOUT_MS），这里再加一层兜底 + // 9. Langfuse 自带超时（LANGFUSE_SHUTDOWN_TIMEOUT_MS），这里再加一层兜底 await withTimeout( (async () => { const { shutdownLangfuse } = await import("@/lib/langfuse"); @@ -177,7 +226,7 @@ export async function runApplicationCleanup( "shutdownLangfuse" ); - // 9. Redis 连接最后关：上面的步骤可能仍在写日志/缓存 + // 10. Redis 连接最后关：上面的步骤可能仍在写日志/缓存 await withTimeout( (async () => { const { closeRedis } = await import("@/lib/redis"); @@ -187,7 +236,7 @@ export async function runApplicationCleanup( "closeRedis" ); - // 10. 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; @@ -199,7 +248,7 @@ export async function runApplicationCleanup( }); } - // 11. 云价格定时同步 + // 12. 云价格定时同步 try { const g = globalThis as unknown as { __CCH_CLOUD_PRICE_SYNC_INTERVAL_ID__?: ReturnType; @@ -213,6 +262,10 @@ 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 totalWarningTimer = setTimeout(() => { 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/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 fda8d3fee..99ba9a29f 100644 --- a/src/lib/rate-limit/lease-service.ts +++ b/src/lib/rate-limit/lease-service.ts @@ -472,7 +472,17 @@ export class LeaseService { if not remaining or ttl <= 0 then settlements[#settlements + 1] = {0, -1} elseif remaining < cost then - settlements[#settlements + 1] = {-1, remaining} + -- 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 diff --git a/tests/integration/lease-settlement-redis.test.ts b/tests/integration/lease-settlement-redis.test.ts index 6d012d74e..ecf4a299a 100644 --- a/tests/integration/lease-settlement-redis.test.ts +++ b/tests/integration/lease-settlement-redis.test.ts @@ -232,7 +232,7 @@ run("lease settlement Lua integration", () => { }); await expect(redis.ttl(noTtl.key)).resolves.toBe(-1); expect(JSON.parse((await redis.get(insufficient.key))!) as BudgetLease).toMatchObject({ - remainingBudget: 0.5, + remainingBudget: 0, }); await expect(redis.type(wrongType.key)).resolves.toBe("list"); diff --git a/tests/unit/lib/shutdown.test.ts b/tests/unit/lib/shutdown.test.ts index 67985a113..264a1b682 100644 --- a/tests/unit/lib/shutdown.test.ts +++ b/tests/unit/lib/shutdown.test.ts @@ -19,12 +19,16 @@ describe.sequential("lifecycle/shutdown", () => { .__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 () => { @@ -110,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) => { @@ -214,4 +261,50 @@ describe.sequential("lifecycle/shutdown", () => { 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/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); + }); });