diff --git a/messages/en/settings/config.json b/messages/en/settings/config.json index 77fa27be1..82e6da41b 100644 --- a/messages/en/settings/config.json +++ b/messages/en/settings/config.json @@ -164,7 +164,7 @@ "affinityIgnoreClientSessionId": "Ignore Client Session ID", "affinityIgnoreClientSessionIdDesc": "When enabled, fingerprintable requests are forced to use longest-prefix affinity for provider stickiness (skipping client Session ID binding); non-fingerprintable requests still use session reuse. Default on.", "streamGateMode": "Stream Content Gate", - "streamGateModeDesc": "Buffers upstream output until the first valid content frame arrives, and automatically fails over to another provider on error frames or empty streams; shadow mode only records divergence statistics without affecting forwarding. Default enabled.", + "streamGateModeDesc": "Controls the first-valid-content gate for ordinary requests: enabled mode fails over on error frames or empty streams, shadow mode only records divergence, and off mode forwards directly. Replay owners always retain the pre-content safety gate regardless of this setting. Default enabled.", "streamGateModeOptions": { "off": "Off", "shadow": "Shadow mode", diff --git a/messages/ja/settings/config.json b/messages/ja/settings/config.json index ddb081fb3..f518682e5 100644 --- a/messages/ja/settings/config.json +++ b/messages/ja/settings/config.json @@ -164,7 +164,7 @@ "affinityIgnoreClientSessionId": "クライアント Session ID を無視", "affinityIgnoreClientSessionIdDesc": "有効にすると、フィンガープリント可能なリクエストは最長プレフィックス親和性によるプロバイダー固定を強制します（クライアント Session ID バインディングをスキップ）。フィンガープリント不可能なリクエストは従来どおりセッション再利用を使用します。デフォルトで有効。", "streamGateMode": "ストリーム内容ゲート", - "streamGateModeDesc": "最初の有効なコンテンツフレームが到着するまでバッファリングし、エラーフレームや空ストリームの場合はプロバイダーを自動的に切り替えて再試行します。シャドウモードは転送に影響せず、判定乖離の統計のみを記録します。デフォルトで有効。", + "streamGateModeDesc": "通常リクエストの最初の有効なコンテンツフレームゲートを制御します。有効時はエラーフレームや空ストリームでプロバイダーを切り替え、シャドウモードは判定乖離のみを記録し、オフ時はそのまま転送します。Replay owner は設定に関係なく、最初のコンテンツ前の安全判定を常に維持します。デフォルトで有効。", "streamGateModeOptions": { "off": "オフ", "shadow": "シャドウモード", diff --git a/messages/ru/settings/config.json b/messages/ru/settings/config.json index 5781aab60..8c603f529 100644 --- a/messages/ru/settings/config.json +++ b/messages/ru/settings/config.json @@ -164,7 +164,7 @@ "affinityIgnoreClientSessionId": "Игнорировать клиентский Session ID", "affinityIgnoreClientSessionIdDesc": "Если включено, запросы с отпечатком принудительно используют аффинность по самому длинному префиксу для закрепления за провайдером (пропуская привязку по клиентскому Session ID); запросы без отпечатка по-прежнему используют переиспользование сессии. По умолчанию включено.", "streamGateMode": "Шлюз потокового контента", - "streamGateModeDesc": "Буферизует вывод до появления первого валидного кадра контента и автоматически переключает провайдера при кадрах ошибок или пустых потоках; теневой режим лишь записывает статистику расхождений, не влияя на пересылку. По умолчанию включено.", + "streamGateModeDesc": "Управляет шлюзом первого валидного кадра для обычных запросов: включенный режим переключает провайдера при ошибках или пустом потоке, теневой режим только фиксирует расхождения, а выключенный режим пересылает ответ напрямую. Replay owner всегда сохраняет проверку безопасности до первого контента независимо от настройки. По умолчанию включено.", "streamGateModeOptions": { "off": "Выключено", "shadow": "Теневой режим", diff --git a/messages/zh-CN/settings/config.json b/messages/zh-CN/settings/config.json index ae3d6ff7c..98f14cefe 100644 --- a/messages/zh-CN/settings/config.json +++ b/messages/zh-CN/settings/config.json @@ -88,7 +88,7 @@ "allowNonConversationEndpointProviderFallback": "允许非对话端点跨供应商 fallback", "allowNonConversationEndpointProviderFallbackDesc": "控制 /v1/messages/count_tokens 与 /v1/responses/compact 在当前供应商失败时，是否沿用现有决策链切换到兼容供应商重试。默认开启，并继续保持 raw passthrough 与非计费语义。", "streamGateMode": "流式内容门控", - "streamGateModeDesc": "在首个有效内容帧到达前先行缓冲，遇到错误帧或空流时自动切换供应商重试；影子模式仅旁路统计判定分歧，不影响转发。默认启用。", + "streamGateModeDesc": "控制普通请求的首个有效内容帧门控：启用时遇到错误帧或空流会自动切换供应商重试，影子模式仅旁路统计判定分歧，关闭时直接转发。Replay owner 无论选择哪种模式，始终保留首内容前的安全判定。默认启用。", "streamGateModeOptions": { "off": "关闭", "shadow": "影子模式", diff --git a/messages/zh-TW/settings/config.json b/messages/zh-TW/settings/config.json index 991b84948..12d93122c 100644 --- a/messages/zh-TW/settings/config.json +++ b/messages/zh-TW/settings/config.json @@ -164,7 +164,7 @@ "affinityIgnoreClientSessionId": "忽略用戶端 Session ID", "affinityIgnoreClientSessionIdDesc": "開啟後，可指紋化的請求強制使用最長前綴親和做供應商黏性（跳過用戶端 Session ID 綁定）；不可指紋化的請求仍走會話複用。預設開啟。", "streamGateMode": "串流內容閘控", - "streamGateModeDesc": "在首個有效內容影格到達前先行緩衝，遇到錯誤影格或空串流時自動切換供應商重試；影子模式僅旁路統計判定分歧，不影響轉發。預設啟用。", + "streamGateModeDesc": "控制一般請求的首個有效內容影格閘控：啟用時遇到錯誤影格或空串流會自動切換供應商重試，影子模式僅旁路統計判定分歧，關閉時直接轉發。Replay owner 無論選擇哪種模式，始終保留首內容前的安全判定。預設啟用。", "streamGateModeOptions": { "off": "關閉", "shadow": "影子模式", diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 7c4bc351a..8e78cfdea 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -107,6 +107,8 @@ import { validateOpenAIImageRequest, } from "./openai-image-compat"; import { ProxyProviderResolver } from "./provider-selector"; +import { abortReplayOwnership, releaseReplayOwnership } from "./replay/replay-spool"; +import { isJsonResponseContentType, isMalformedJsonResponseBody } from "./response-content-type"; import { finalizeHedgeLoserBilling, hasStreamCompletionMarker } from "./response-handler"; import type { ProxySession } from "./session"; import { @@ -1343,6 +1345,15 @@ function applyClaudeMetadataUserIdInjectionWithAudit( export class ProxyForwarder { static async send(session: ProxySession): Promise { + try { + return await ProxyForwarder.sendInternal(session); + } catch (error) { + await abortReplayOwnership(session, "forward_failed"); + throw error; + } + } + + private static async sendInternal(session: ProxySession): Promise { if (!session.provider || !session.authState?.success) { throw new Error("代理上下文缺少供应商或鉴权信息"); } @@ -1653,7 +1664,7 @@ export class ProxyForwarder { }; try { - const response = await ProxyForwarder.doForward( + let response = await ProxyForwarder.doForward( session, currentProvider, activeEndpoint.baseUrl, @@ -1668,9 +1679,7 @@ export class ProxyForwarder { const isHtml = normalizedContentType.includes("text/html") || normalizedContentType.includes("application/xhtml+xml"); - const isJson = - normalizedContentType.includes("application/json") || - normalizedContentType.includes("+json"); + const isJson = isJsonResponseContentType(contentType); // ========== 流式响应：延迟成功判定（避免“假 200”）========== // 背景：上游可能返回 HTTP 200，但 SSE 内容为错误 JSON（如 {"error": "..."}）。 @@ -1682,7 +1691,7 @@ export class ProxyForwarder { // 解决：Forwarder 只负责尽快把 Response 返回给下游开始透传， // 把最终成功/失败结算延迟到 ResponseHandler：等 SSE 正常结束后再基于最终 body 补充检查并更新内部状态。 if (isSSE) { - // ========== F1 流式内容门控（enforce 模式）========== + // ========== F1 流式内容门控（enforce 或 Replay owner）========== // 在向客户端提交响应前等待首个有效内容帧： // - 中性前缀（ping/metadata/usage-only）缓冲后随提交一并冲刷； // - error/malformed/空流在此抛错 -> 外层 catch 归类 -> 换供应商（客户端零字节）； @@ -1691,8 +1700,10 @@ export class ProxyForwarder { let streamingResponse = response; let gateChainAudit: ProviderChainItem["streamGate"]; const gateMode = resolveStreamGateMode(); + const shouldRunPrecommitGate = + gateMode === "enforce" || session.replayState?.role === "owner"; if ( - gateMode === "enforce" && + shouldRunPrecommitGate && response.body && session.getEndpointPolicy().kind !== "raw_passthrough" ) { @@ -1854,7 +1865,11 @@ export class ProxyForwarder { // 因此这里在进入成功分支前做一次强信号检测：仅当 body 看起来是完整 HTML 文档时才视为错误。 let inspectedText: string | undefined; let inspectedTruncated = false; - // 注意：这里不会对“大体积 JSON”做假 200 检测（例如 Content-Length > 32KiB）。 + // Replay owner 的 buffered JSON 必须在提交前完整验证，确保 malformed body 能进入 + // 现有 provider fallback，且不会被 ResponseHandler 持久化为 completed Replay。 + const shouldStrictValidateReplayJson = isJson && session.replayState?.role === "owner"; + let replayJsonValidationExceededLimit = false; + // 普通非 Replay 请求仍不会对“大体积 JSON”做假 200 检测（例如 Content-Length > 32KiB）。 // 原因： // - 非流式路径需要 clone 并额外读取响应体，会带来额外的内存/延迟开销； // - 大体积 JSON 更可能是正常响应（而不是网关/WAF 的短错误 JSON）。 @@ -1864,7 +1879,22 @@ export class ProxyForwarder { hasValidContentLength && contentLengthBytes <= NON_STREAM_BODY_INSPECTION_MAX_BYTES; const shouldInspectBody = isHtml || !hasValidContentLength || shouldInspectJson; - if (shouldInspectBody) { + if (shouldStrictValidateReplayJson) { + const validationLimit = getEnvConfig().REPLAY_MAX_PAYLOAD_BYTES; + if (contentLengthBytes !== null && contentLengthBytes > validationLimit) { + replayJsonValidationExceededLimit = true; + releaseReplayOwnership(session); + } else { + const validation = await ProxyForwarder.bufferReplayJsonResponse( + response, + validationLimit + ); + response = validation.response; + inspectedText = validation.text; + replayJsonValidationExceededLimit = validation.exceededLimit; + if (validation.exceededLimit) releaseReplayOwnership(session); + } + } else if (shouldInspectBody) { // 注意：Response.clone() 会 tee 底层 ReadableStream，可能带来一定的瞬时内存开销； // 这里通过“最多读取 32 KiB”并在截断时 cancel 克隆分支来控制开销。 const clonedResponse = response.clone(); @@ -1913,9 +1943,25 @@ export class ProxyForwarder { } } + if ( + shouldStrictValidateReplayJson && + inspectedText !== undefined && + isMalformedJsonResponseBody(contentType, inspectedText) + ) { + const rawBodyMaxChars = 4096; + throw new ProxyError("MALFORMED_BUFFERED_JSON", 502, { + body: "Upstream returned malformed JSON", + providerId: currentProvider.id, + providerName: currentProvider.name, + rawBody: inspectedText.slice(0, rawBodyMaxChars), + rawBodyTruncated: inspectedText.length > rawBodyMaxChars, + isSyntheticFake200: true, + }); + } + // 对于缺失或非法 Content-Length 的情况，需要 clone 并检查响应体 // 注意：这会增加一定的性能开销，但对于非流式响应是可接受的 - if (!contentLength || !hasValidContentLength) { + if ((!contentLength || !hasValidContentLength) && !replayJsonValidationExceededLimit) { const responseText = inspectedText ?? ""; if (!responseText || responseText.trim() === "") { @@ -4721,10 +4767,11 @@ export class ProxyForwarder { attempt.reader = response.body.getReader(); try { - // F1 门控（enforce）：胜者判定从「首个非空字节」升级为「首个有效内容帧」。 + // F1 门控（enforce 或 Replay owner）：胜者判定从「首个非空字节」升级为 + // 「首个有效内容帧」。 // 级联阈值计时器保持不动——内容慢的 attempt 不提交，自动触发下一候选竞速。 const hedgeGateFamily = - resolveStreamGateMode() === "enforce" && + (resolveStreamGateMode() === "enforce" || session.replayState?.role === "owner") && session.getEndpointPolicy().kind !== "raw_passthrough" ? mapProviderTypeToFamily(attempt.provider.providerType) : null; @@ -7924,6 +7971,51 @@ export class ProxyForwarder { } } + private static async bufferReplayJsonResponse( + response: Response, + maxBytes: number + ): Promise<{ response: Response; text: string | undefined; exceededLimit: boolean }> { + const reader = response.body?.getReader(); + if (!reader) { + return { response, text: "", exceededLimit: false }; + } + + const chunks: Uint8Array[] = []; + let totalBytes = 0; + while (true) { + const { done, value } = await reader.read(); + if (done) { + reader.releaseLock(); + const combined = chunks.length > 0 ? concatChunks(chunks) : null; + const body = combined ? new Uint8Array(combined) : new Uint8Array(0); + return { + response: new Response(body, { + status: response.status, + statusText: response.statusText, + headers: response.headers, + }), + text: new TextDecoder().decode(body), + exceededLimit: false, + }; + } + if (!value || value.byteLength === 0) continue; + + chunks.push(value); + totalBytes += value.byteLength; + if (totalBytes > maxBytes) { + return { + response: new Response(ProxyForwarder.buildBufferedPrefixStream(chunks, reader), { + status: response.status, + statusText: response.statusText, + headers: response.headers, + }), + text: undefined, + exceededLimit: true, + }; + } + } + } + private static buildBufferedPrefixStream( prefixChunks: Uint8Array[], reader: ReadableStreamDefaultReader diff --git a/src/app/v1/_lib/proxy/replay/replay-guard.ts b/src/app/v1/_lib/proxy/replay/replay-guard.ts index 065366192..e31a994d4 100644 --- a/src/app/v1/_lib/proxy/replay/replay-guard.ts +++ b/src/app/v1/_lib/proxy/replay/replay-guard.ts @@ -5,6 +5,7 @@ import { getEnvConfig } from "@/lib/config/env.schema"; import { logger } from "@/lib/logger"; import { getProxyRuntimeSettings } from "@/lib/system-settings/proxy-runtime"; import type { ProxySession } from "../session"; +import { restoreReplayResponseHeaders } from "./replay-headers"; import { deriveReplayIdentity, REPLAY_BYPASS_HEADER, type ReplayIdentity } from "./replay-identity"; import { getReplayStore, type ReplayMeta, type ReplayStore } from "./replay-store"; @@ -108,7 +109,7 @@ export class ProxyReplayGuard { // 热层块已过期：落 PG } else if (meta.status === "owning") { const heartbeatFresh = Date.now() - meta.heartbeatAt < ATTACH_STALL_MS; - if (env.REPLAY_LIVE_DEDUP_ENABLED && heartbeatFresh) { + if (meta.delivery !== "buffered" && env.REPLAY_LIVE_DEDUP_ENABLED && heartbeatFresh) { await ProxyReplayGuard.writeAuditRow(session, identity, meta.statusCode, "attached_live"); return ProxyReplayGuard.buildLiveAttachResponse(identity, meta, store); } @@ -226,9 +227,14 @@ export class ProxyReplayGuard { stored: Record, mode: "completed" | "live" ): Headers { - const headers = new Headers(); - headers.set("content-type", stored["content-type"] ?? "text/event-stream"); - headers.set("cache-control", "no-cache"); + const headers = restoreReplayResponseHeaders(stored); + const contentType = headers.get("content-type")?.toLowerCase() ?? ""; + if (mode === "live" && !contentType) { + headers.set("content-type", "text/event-stream"); + } + if (mode === "live" || contentType.includes("text/event-stream")) { + headers.set("cache-control", "no-cache"); + } headers.set("x-cch-replay", mode); return headers; } diff --git a/src/app/v1/_lib/proxy/replay/replay-headers.ts b/src/app/v1/_lib/proxy/replay/replay-headers.ts new file mode 100644 index 000000000..28cafcc53 --- /dev/null +++ b/src/app/v1/_lib/proxy/replay/replay-headers.ts @@ -0,0 +1,42 @@ +const REPLAY_EXCLUDED_RESPONSE_HEADERS = new Set([ + "connection", + "content-encoding", + "content-length", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "set-cookie", + "set-cookie2", + "te", + "trailer", + "transfer-encoding", + "upgrade", +]); + +export function captureReplayResponseHeaders( + source: Headers, + fallbackContentType?: string +): Record { + const captured: Record = {}; + source.forEach((value, name) => { + const normalizedName = name.toLowerCase(); + if (!REPLAY_EXCLUDED_RESPONSE_HEADERS.has(normalizedName)) { + captured[normalizedName] = value; + } + }); + if (fallbackContentType && !captured["content-type"]) { + captured["content-type"] = fallbackContentType; + } + return captured; +} + +export function restoreReplayResponseHeaders(stored: Record): Headers { + const headers = new Headers(); + for (const [name, value] of Object.entries(stored)) { + const normalizedName = name.toLowerCase(); + if (!REPLAY_EXCLUDED_RESPONSE_HEADERS.has(normalizedName)) { + headers.set(normalizedName, value); + } + } + return headers; +} diff --git a/src/app/v1/_lib/proxy/replay/replay-spool.ts b/src/app/v1/_lib/proxy/replay/replay-spool.ts index ab7f1b9f4..73303eb1c 100644 --- a/src/app/v1/_lib/proxy/replay/replay-spool.ts +++ b/src/app/v1/_lib/proxy/replay/replay-spool.ts @@ -1,8 +1,9 @@ import { getEnvConfig } from "@/lib/config/env.schema"; import { logger } from "@/lib/logger"; import type { ProxySession } from "../session"; +import { captureReplayResponseHeaders } from "./replay-headers"; import { isReplayEnabled, type ReplayIdentity } from "./replay-identity"; -import { getReplayStore, type ReplayMeta } from "./replay-store"; +import { getReplayStore, type ReplayDelivery, type ReplayMeta } from "./replay-store"; /** * F2 owner 侧 spool：把客户端可见字节（pump 处理后流）以 write-behind 方式 @@ -21,6 +22,8 @@ import { getReplayStore, type ReplayMeta } from "./replay-store"; const FLUSH_INTERVAL_MS = 100; const FLUSH_BYTES_THRESHOLD = 64 * 1024; +const OWNER_HEARTBEAT_INTERVAL_MS = 15_000; +const PRE_SPOOL_ABORT_WAIT_MS = 100; let activeSpoolCount = 0; @@ -39,6 +42,8 @@ export class ReplaySpool { private disabled = false; private terminal = false; private flushTimer: ReturnType | null = null; + private ownerHeartbeatTimer: ReturnType | null = null; + private ownerHeartbeatInFlight = false; private writeChain: Promise = Promise.resolve(); private metaWritten = false; @@ -46,9 +51,11 @@ export class ReplaySpool { private readonly identity: ReplayIdentity, private readonly ownerToken: string, private readonly statusCode: number, - private readonly contentType: string + private readonly headers: Record, + private readonly delivery: ReplayDelivery = "stream" ) { activeSpoolCount++; + this.startOwnerHeartbeat(); } /** 已达终态（complete/abort）或已失效（disable/halt）：调用方兜底判断用。 */ @@ -113,21 +120,25 @@ export class ReplaySpool { this.writeChain = this.writeChain.then(async () => { try { if (this.disabled) return; - const appended = await this.store.appendChunks(this.identity.replayId, batch); + const expectedChunkCount = this.chunkCount + batch.length; + const appended = await this.store.writeOwned( + this.identity.replayId, + this.ownerToken, + this.buildMeta("owning", { chunkCount: expectedChunkCount }), + batch + ); if (this.disabled) return; if (appended === null) { - // Redis 不可用：本次 replay 放弃（已写块靠 TTL 清理） + // Redis 不可用：本次 replay 放弃（热层写是原子的，不会留下半批数据） this.disable("redis_unavailable"); return; } - this.chunkCount = appended; - await this.writeMeta("owning"); - if (this.disabled) return; - const leaseHeld = await this.store.renewOwnerLease(this.identity.replayId, this.ownerToken); - if (!leaseHeld) { - // 所有权已失（租约被他人接管）：停止 spool，但绝不删条目 + if (appended === false) { this.halt("owner_lease_lost"); + return; } + this.chunkCount = appended; + this.metaWritten = true; } catch (error) { logger.debug("[ReplaySpool] flush failed, disabling spool", { error: error instanceof Error ? error.message : String(error), @@ -137,16 +148,14 @@ export class ReplaySpool { }); } - private async writeMeta( - status: ReplayMeta["status"], - extra?: Partial - ): Promise { - const meta: ReplayMeta = { + private buildMeta(status: ReplayMeta["status"], extra?: Partial): ReplayMeta { + return { status, verifier: this.identity.verifier, scopeTag: this.identity.scopeTag, statusCode: this.statusCode, - headers: { "content-type": this.contentType }, + headers: this.headers, + delivery: this.delivery, format: this.identity.format, model: this.identity.model, chunkCount: this.chunkCount, @@ -154,8 +163,25 @@ export class ReplaySpool { heartbeatAt: Date.now(), ...extra, }; - await this.store.setMeta(this.identity.replayId, meta); - this.metaWritten = true; + } + + private startOwnerHeartbeat(): void { + this.ownerHeartbeatTimer = setInterval(() => { + if (this.disabled || this.released || this.ownerHeartbeatInFlight) return; + this.ownerHeartbeatInFlight = true; + void this.store + .renewOwnerLease(this.identity.replayId, this.ownerToken) + .then((leaseHeld) => { + if (!leaseHeld && !this.released) this.halt("owner_lease_lost"); + }) + .catch(() => { + if (!this.released) this.halt("owner_lease_lost"); + }) + .finally(() => { + this.ownerHeartbeatInFlight = false; + }); + }, OWNER_HEARTBEAT_INTERVAL_MS); + this.ownerHeartbeatTimer.unref?.(); } /** 立即建立 owning meta（handleStream 创建 spool 时调用，供 attach 读者尽早看到状态）。 */ @@ -163,7 +189,21 @@ export class ReplaySpool { this.writeChain = this.writeChain.then(async () => { try { if (this.disabled || this.metaWritten) return; - await this.writeMeta("owning"); + const chunkCount = await this.store.writeOwned( + this.identity.replayId, + this.ownerToken, + this.buildMeta("owning") + ); + if (chunkCount === null) { + this.disable("redis_unavailable"); + return; + } + if (chunkCount === false) { + this.halt("owner_lease_lost"); + return; + } + this.chunkCount = chunkCount; + this.metaWritten = true; } catch (error) { logger.debug("[ReplaySpool] bootstrap failed, disabling spool", { error: error instanceof Error ? error.message : String(error), @@ -180,7 +220,7 @@ export class ReplaySpool { async completeAfterBilling(messageRequestId: number | null): Promise { if (this.disabled || this.terminal) return; this.terminal = true; - this.clearTimer(); + this.clearFlushTimer(); const tail = this.decoder.decode(); if (tail.length > 0) { this.pending.push(tail); @@ -194,14 +234,22 @@ export class ReplaySpool { let pgPersisted = false; try { if (this.disabled) return; - if (batch.length > 0) { - const appended = await this.store.appendChunks(this.identity.replayId, batch); - if (appended === null) { - // 尾批丢失则热层条目不完整，绝不能置 completed - throw new Error("final replay flush failed"); - } - this.chunkCount = appended; + const expectedChunkCount = this.chunkCount + batch.length; + const appended = await this.store.writeOwned( + this.identity.replayId, + this.ownerToken, + this.buildMeta("owning", { chunkCount: expectedChunkCount }), + batch + ); + if (appended === false) { + throw new Error("replay owner lease lost before completion"); + } + if (appended === null) { + // 尾批或 owning meta 丢失时热层条目不完整，绝不能置 completed + throw new Error("final replay flush failed"); } + this.chunkCount = appended; + this.metaWritten = true; // 先写 PG（持久 payload），再翻 Redis meta 为 completed（热层可服务） await this.store.persistCompleted({ replayId: this.identity.replayId, @@ -212,13 +260,20 @@ export class ReplaySpool { format: this.identity.format, model: this.identity.model, statusCode: this.statusCode, - headers: { "content-type": this.contentType }, + headers: this.headers, payload: this.parts.join(""), byteSize: this.totalBytes, sourceMessageRequestId: messageRequestId, }); pgPersisted = true; - await this.writeMeta("completed", { messageRequestId }); + const completed = await this.store.completeOwned( + this.identity.replayId, + this.ownerToken, + this.buildMeta("completed", { messageRequestId }) + ); + if (!completed) { + throw new Error("replay owner lease lost before completed meta"); + } logger.info("[ReplaySpool] replay entry completed", { replayId: this.identity.replayId.slice(0, 12), chunkCount: this.chunkCount, @@ -231,10 +286,14 @@ export class ReplaySpool { error: error instanceof Error ? error.message : String(error), pgPersisted, }); - await this.writeMeta("aborted", { abortReason: "complete_failed" }).catch(() => undefined); - await this.store.deleteChunks(this.identity.replayId).catch(() => undefined); + await this.store + .abortOwned( + this.identity.replayId, + this.ownerToken, + this.buildMeta("aborted", { abortReason: "complete_failed" }) + ) + .catch(() => false); } finally { - await this.store.releaseOwner(this.identity.replayId, this.ownerToken); this.release(); } }); @@ -252,13 +311,14 @@ export class ReplaySpool { try { // 已失效（disable 已清理 / halt 已让渡所有权）：不得再写 meta 覆盖新 owner if (this.disabled) return; - // aborted meta 保留（短 TTL）供 attach 读者感知终态；块立即删除 - await this.writeMeta("aborted", { abortReason: reason }); - await this.store.deleteChunks(this.identity.replayId); + await this.store.abortOwned( + this.identity.replayId, + this.ownerToken, + this.buildMeta("aborted", { abortReason: reason }) + ); } catch { // 热层清理失败靠 TTL 兜底 } finally { - await this.store.releaseOwner(this.identity.replayId, this.ownerToken); this.release(); } }); @@ -285,10 +345,19 @@ export class ReplaySpool { // 清理顺着 writeChain 串行：与 in-flight append 竞态时绝不出现「删除后又写回」 this.writeChain = this.writeChain.then(async () => { if (deleteEntry) { - await this.store.deleteEntry(this.identity.replayId).catch(() => undefined); + await this.store + .abortOwned( + this.identity.replayId, + this.ownerToken, + this.buildMeta("aborted", { abortReason: reason }) + ) + .catch(() => false); + } else { + // compare-delete 只删自己的 token：所有权已失时为安全 no-op + await this.store + .releaseOwner(this.identity.replayId, this.ownerToken) + .catch(() => undefined); } - // compare-delete 只删自己的 token：所有权已失时为安全 no-op - await this.store.releaseOwner(this.identity.replayId, this.ownerToken).catch(() => undefined); }); logger.debug("[ReplaySpool] spool disabled", { replayId: this.identity.replayId.slice(0, 12), @@ -302,17 +371,70 @@ export class ReplaySpool { private release(): void { if (this.released) return; this.released = true; + this.clearOwnerHeartbeat(); activeSpoolCount = Math.max(0, activeSpoolCount - 1); } + private clearFlushTimer(): void { + if (!this.flushTimer) return; + clearTimeout(this.flushTimer); + this.flushTimer = null; + } + + private clearOwnerHeartbeat(): void { + if (!this.ownerHeartbeatTimer) return; + clearInterval(this.ownerHeartbeatTimer); + this.ownerHeartbeatTimer = null; + } + private clearTimer(): void { - if (this.flushTimer) { - clearTimeout(this.flushTimer); - this.flushTimer = null; - } + this.clearFlushTimer(); + this.clearOwnerHeartbeat(); } } +/** Forwarder 在 spool 创建前终止时，以 owner token 原子封死 Replay 条目。 */ +export async function abortReplayOwnership(session: ProxySession, reason: string): Promise { + const replayState = session.replayState; + if (replayState?.role !== "owner") return; + + // 先清本地角色形成幂等门；Redis 失败时保留 fail-open，租约最终由 TTL 回收。 + session.replayState = null; + const { identity, ownerToken } = replayState; + const meta: ReplayMeta = { + status: "aborted", + verifier: identity.verifier, + scopeTag: identity.scopeTag, + statusCode: 502, + headers: {}, + format: identity.format, + model: identity.model, + chunkCount: 0, + byteSize: 0, + heartbeatAt: Date.now(), + abortReason: reason, + }; + const store = getReplayStore(); + const cleanup = (async () => { + const aborted = await store.abortOwned(identity.replayId, ownerToken, meta).catch(() => false); + if (aborted) return; + + logger.warn("[ReplaySpool] failed to abort pre-spool replay ownership", { + replayId: identity.replayId.slice(0, 12), + reason, + }); + await store.releaseOwner(identity.replayId, ownerToken).catch(() => undefined); + })(); + + let timer: ReturnType | null = null; + const deadline = new Promise((resolve) => { + timer = setTimeout(resolve, PRE_SPOOL_ABORT_WAIT_MS); + timer.unref?.(); + }); + await Promise.race([cleanup, deadline]); + if (timer) clearTimeout(timer); +} + /** * 不会建 spool 的路径统一释放 owner 租约并清角色—— * 否则相同请求的重试会被残留租约挡满 45s。 @@ -329,12 +451,13 @@ export function releaseReplayOwnership(session: ProxySession): void { /** * handleStream 建 pump 时创建 owner spool。 * 前置：guard 阶段已成功 claim owner（session.replayState.role === "owner"）。 - * 并发 spool 超上限 / 非 2xx / 非 SSE / 开关关闭 / 异常时返回 null（本请求不做 - * replay），并立即释放 owner 租约。 + * 并发 spool 超上限 / 非 2xx / 响应类型与 delivery 不匹配 / 开关关闭 / 异常时返回 + * null（本请求不做 replay），并立即释放 owner 租约。 */ export function createReplaySpoolIfOwner( session: ProxySession, - response: Response + response: Response, + delivery: ReplayDelivery = "stream" ): ReplaySpool | null { const replayState = session.replayState; if (replayState?.role !== "owner") return null; @@ -352,14 +475,24 @@ export function createReplaySpoolIfOwner( return declineOwnership(); } if (response.status < 200 || response.status >= 300) return declineOwnership(); - const contentType = response.headers.get("content-type") ?? "text/event-stream"; - if (!contentType.toLowerCase().includes("text/event-stream")) return declineOwnership(); + const contentType = response.headers.get("content-type"); + const isSse = contentType + ? contentType.toLowerCase().includes("text/event-stream") + : delivery === "stream"; + if ((delivery === "stream" && !isSse) || (delivery === "buffered" && isSse)) { + return declineOwnership(); + } + const headers = captureReplayResponseHeaders( + response.headers, + delivery === "stream" ? "text/event-stream" : undefined + ); const spool = new ReplaySpool( replayState.identity, replayState.ownerToken, response.status, - contentType + headers, + delivery ); spool.bootstrap(); return spool; diff --git a/src/app/v1/_lib/proxy/replay/replay-store.ts b/src/app/v1/_lib/proxy/replay/replay-store.ts index f3ea4bd11..64dc6590d 100644 --- a/src/app/v1/_lib/proxy/replay/replay-store.ts +++ b/src/app/v1/_lib/proxy/replay/replay-store.ts @@ -20,6 +20,7 @@ import { RedisListStore } from "@/lib/redis/redis-list-store"; */ export type ReplayStatus = "owning" | "completed" | "aborted"; +export type ReplayDelivery = "stream" | "buffered"; export interface ReplayMeta { status: ReplayStatus; @@ -28,6 +29,8 @@ export interface ReplayMeta { statusCode: number; /** 仅保留承载语义的响应头（content-type 等） */ headers: Record; + /** buffered owner 不公开 live attach；缺失字段按旧版 stream 条目处理。 */ + delivery?: ReplayDelivery; format: string; model: string | null; chunkCount: number; @@ -54,6 +57,38 @@ if redis.call('GET', KEYS[1]) == ARGV[1] then end return 0`; +const LUA_WRITE_OWNED = ` +if redis.call('GET', KEYS[1]) ~= ARGV[1] then + return -1 +end +local len = redis.call('LLEN', KEYS[3]) +if #ARGV > 4 then + len = redis.call('RPUSH', KEYS[3], unpack(ARGV, 5)) + if tonumber(ARGV[2]) > 0 then + redis.call('EXPIRE', KEYS[3], ARGV[2]) + end +end +redis.call('SETEX', KEYS[2], ARGV[2], ARGV[4]) +redis.call('EXPIRE', KEYS[1], ARGV[3]) +return len`; + +const LUA_ABORT_OWNED = ` +if redis.call('GET', KEYS[1]) ~= ARGV[1] then + return 0 +end +redis.call('SETEX', KEYS[2], ARGV[2], ARGV[3]) +redis.call('DEL', KEYS[3]) +redis.call('DEL', KEYS[1]) +return 1`; + +const LUA_COMPLETE_OWNED = ` +if redis.call('GET', KEYS[1]) ~= ARGV[1] then + return 0 +end +redis.call('SETEX', KEYS[2], ARGV[2], ARGV[3]) +redis.call('DEL', KEYS[1]) +return 1`; + type RedisRawClient = Pick & { eval(...args: [script: string, numkeys: number, ...rest: (string | number)[]]): Promise; }; @@ -104,6 +139,43 @@ export class ReplayStore { return this.chunks.rpushBatch(replayId, values, resolveReplayTtlSeconds()); } + /** + * owner 热层写入：token 校验、chunk 追加、owning meta 更新和租约续期在同一 Lua + * 内完成，避免旧 owner 在租约交接窗口污染新 owner 的 chunks/meta。 + * null 表示 Redis 不可用，false 表示 token 已失效，number 为当前 chunk 总数。 + */ + async writeOwned( + replayId: string, + ownerToken: string, + meta: ReplayMeta, + values: string[] = [] + ): Promise { + const redis = this.getRawRedis(); + if (!redis) return null; + try { + const result = await redis.eval( + LUA_WRITE_OWNED, + 3, + `cch:replay:owner:${replayId}`, + `cch:replay:meta:${replayId}`, + `cch:replay:chunks:${replayId}`, + ownerToken, + resolveReplayTtlSeconds(), + OWNER_LEASE_TTL_SECONDS, + JSON.stringify(meta), + ...values + ); + const length = typeof result === "number" ? result : Number(result); + return length === -1 ? false : length; + } catch (error) { + logger.debug("[ReplayStore] fenced owner write failed", { + replayId: replayId.slice(0, 12), + error: error instanceof Error ? error.message : String(error), + }); + return null; + } + } + /** 从 offset（0-based）读到当前末尾；Redis 不可用返回 null。 */ async readChunks(replayId: string, fromIndex: number): Promise { return this.chunks.lrangeFrom(replayId, fromIndex); @@ -172,6 +244,55 @@ export class ReplayStore { } } + /** 仅当前 token 仍持有租约时，原子终止条目并清理热层响应块。 */ + async abortOwned(replayId: string, ownerToken: string, meta: ReplayMeta): Promise { + const redis = this.getRawRedis(); + if (!redis) return false; + try { + const result = await redis.eval( + LUA_ABORT_OWNED, + 3, + `cch:replay:owner:${replayId}`, + `cch:replay:meta:${replayId}`, + `cch:replay:chunks:${replayId}`, + ownerToken, + resolveReplayTtlSeconds(), + JSON.stringify(meta) + ); + return result === 1; + } catch (error) { + logger.debug("[ReplayStore] fenced abort failed", { + replayId: replayId.slice(0, 12), + error: error instanceof Error ? error.message : String(error), + }); + return false; + } + } + + /** 仅当前 token 仍持有租约时，原子翻转 completed meta 并释放租约。 */ + async completeOwned(replayId: string, ownerToken: string, meta: ReplayMeta): Promise { + const redis = this.getRawRedis(); + if (!redis) return false; + try { + const result = await redis.eval( + LUA_COMPLETE_OWNED, + 2, + `cch:replay:owner:${replayId}`, + `cch:replay:meta:${replayId}`, + ownerToken, + resolveReplayTtlSeconds(), + JSON.stringify(meta) + ); + return result === 1; + } catch (error) { + logger.debug("[ReplayStore] fenced completion failed", { + replayId: replayId.slice(0, 12), + error: error instanceof Error ? error.message : String(error), + }); + return false; + } + } + // ===== PG 完成持久层 ===== /** diff --git a/src/app/v1/_lib/proxy/response-content-type.ts b/src/app/v1/_lib/proxy/response-content-type.ts new file mode 100644 index 000000000..3c68fd873 --- /dev/null +++ b/src/app/v1/_lib/proxy/response-content-type.ts @@ -0,0 +1,17 @@ +export function isJsonResponseContentType(contentType: string | null): boolean { + if (!contentType) return false; + + const mediaType = contentType.split(";", 1)[0]?.trim().toLowerCase() ?? ""; + return mediaType === "application/json" || mediaType.endsWith("+json"); +} + +export function isMalformedJsonResponseBody(contentType: string | null, body: string): boolean { + if (!isJsonResponseContentType(contentType)) return false; + + try { + JSON.parse(body); + return false; + } catch { + return true; + } +} diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index 9f0ae0b73..d143eaf82 100644 --- a/src/app/v1/_lib/proxy/response-handler.ts +++ b/src/app/v1/_lib/proxy/response-handler.ts @@ -61,7 +61,12 @@ import { } from "./demand-driven-response-pump"; import { isDiscoveryProtocolErrorPayload } from "./discovery-validity"; import { isClientAbortError, isTransportError } from "./errors"; -import { createReplaySpoolIfOwner, releaseReplayOwnership } from "./replay/replay-spool"; +import { + abortReplayOwnership, + createReplaySpoolIfOwner, + releaseReplayOwnership, +} from "./replay/replay-spool"; +import { isMalformedJsonResponseBody } from "./response-content-type"; import type { ProxySession } from "./session"; import { consumeDeferredStreamingFinalization, @@ -71,6 +76,11 @@ import { } from "./stream-finalization"; import { mapProviderTypeToFamily } from "./stream-gate/frame-classifier"; import { createShadowGateObserver, resolveStreamGateMode } from "./stream-gate/stream-content-gate"; +import { + createStreamProtocolObserver, + type StreamProtocolObservation, + type StreamProtocolObserver, +} from "./stream-gate/stream-protocol-observer"; const CLIENT_ABORT_DRAIN_MAX_MS = 60_000; const STREAM_STATS_MAX_BUFFER_BYTES = 10 * 1024 * 1024; @@ -78,6 +88,7 @@ const STREAM_STATS_HEAD_BYTES = 1024 * 1024; const STREAM_STATS_TAIL_BYTES = STREAM_STATS_MAX_BUFFER_BYTES - STREAM_STATS_HEAD_BYTES; const STREAM_STATS_TAIL_CHUNKS = 8192; const STREAM_STATS_TRUNCATED_MARKER = "\n\n: [cch_truncated]\n\n"; +const RESPONSE_TEXT_ENCODER = new TextEncoder(); type BoundedStreamTextSnapshot = { text: string; @@ -1534,6 +1545,7 @@ function finalizeDeferredStreamingFinalizationIfNeeded( streamEndedNormally: boolean, clientAborted: boolean, discoveryLeaseLifecycle: DiscoveryLeaseLifecycle, + protocolObservation: StreamProtocolObservation | null, abortReason?: string ): FinalizeDeferredStreamingResult { const meta = consumeDeferredStreamingFinalization(session); @@ -1788,12 +1800,21 @@ function finalizeDeferredStreamingFinalizationIfNeeded( const bodyDetected = shouldDetectFake200 ? detectUpstreamErrorFromSseOrJsonText(allContent) : ({ isError: false } as const); + const protocolFailure = protocolObservation?.failure ?? null; + const successfulHttpProtocolFailure = + streamEndedNormally && upstreamStatusCode >= 200 && upstreamStatusCode < 300 && protocolFailure; const detected = - shouldDetectFake200 && !bodyDetected.isError && completionInspection.hasProtocolError + !bodyDetected.isError && + ((shouldDetectFake200 && completionInspection.hasProtocolError) || + successfulHttpProtocolFailure) ? ({ isError: true, code: "UPSTREAM_PROTOCOL_ERROR", - detail: "Upstream stream emitted a protocol error event", + detail: protocolFailure + ? `Upstream stream emitted a ${protocolFailure.verdict} frame${ + protocolFailure.eventName ? ` (${protocolFailure.eventName})` : "" + }` + : "Upstream stream emitted a protocol error event", } as const) : bodyDetected; let clientAbortGateUsage: FinalizeDeferredStreamingResult["clientAbortGateUsage"]; @@ -1801,6 +1822,7 @@ function finalizeDeferredStreamingFinalizationIfNeeded( if (!clientAborted || upstreamStatusCode < 200 || upstreamStatusCode >= 300) { return false; } + if (protocolFailure) return false; const abortDetected = detectUpstreamErrorFromSseOrJsonText(allContent); if (abortDetected.isError) { @@ -2340,23 +2362,35 @@ export class ProxyResponseHandler { session: ProxySession, response: Response ): Promise { - // F2：stream 请求被上游以非流响应回答时不做 replay，立即让出 owner 租约 - releaseReplayOwnership(session); const messageContext = session.messageContext; const provider = session.provider; const discoveryLeaseLifecycle = startDiscoveryLeaseLifecycle(session); if (!provider) { + void abortReplayOwnership(session, "non_stream_missing_provider"); discardBeforeResponseBodySnapshot(session); releaseSessionAgent(session); void finalizeNonStreamDiscoveryResources(session, discoveryLeaseLifecycle); return response; } + const createBufferedReplaySpool = (targetResponse: Response) => { + if (!messageContext) { + void abortReplayOwnership(session, "non_stream_missing_message_context"); + return null; + } + if (targetResponse.status < 200 || targetResponse.status >= 300) { + void abortReplayOwnership(session, "non_stream_http_error"); + return null; + } + return createReplaySpoolIfOwner(session, targetResponse, "buffered"); + }; + const responseForLog = response.clone(); const statusCode = response.status; let finalResponse = response; let finalResponseBodyForSnapshot: string | null = null; + let responseTransformFailed = false; const persistNonStreamAfterSnapshot = (targetResponse: Response, body: string) => { if (!session.sessionId || !session.shouldPersistSessionDebugArtifacts()) { return; @@ -2401,6 +2435,8 @@ export class ProxyResponseHandler { const responseForStats = response.clone(); const statusCode = response.status; + const replaySpool = createBufferedReplaySpool(response); + let replayCompletionScheduled = false; const taskId = `non-stream-passthrough-${messageContext?.id || `unknown-${Date.now()}`}`; const statsAbortController = new AbortController(); @@ -2491,19 +2527,57 @@ export class ProxyResponseHandler { // 使用共享的统计处理方法 const duration = Date.now() - session.startTime; - let providerFailureScheduled = false; - const scheduleProviderFailure = () => { - if (!commitProviderFailure || providerFailureScheduled) return; - providerFailureScheduled = true; - return schedulePostTerminalSideEffects({ + const postTerminalSideEffects: Array<() => Promise> = []; + if (commitProviderFailure) { + postTerminalSideEffects.push(commitProviderFailure); + } + if (replaySpool) { + const replayMalformedJson = isMalformedJsonResponseBody( + response.headers.get("content-type"), + responseText + ); + const replayDetected = detectUpstreamErrorFromSseOrJsonText(responseText); + if ( + statusCode >= 200 && + statusCode < 300 && + !replayMalformedJson && + !replayDetected.isError + ) { + replaySpool.observe(RESPONSE_TEXT_ENCODER.encode(responseText)); + postTerminalSideEffects.push(() => + replaySpool.completeAfterBilling(messageContext?.id ?? null) + ); + } else { + await replaySpool.abort( + replayMalformedJson + ? "non_stream_malformed_json" + : replayDetected.isError + ? "non_stream_upstream_error" + : "non_stream_http_error" + ); + } + } + let postTerminalSideEffectsScheduled = false; + const scheduleCommittedSideEffects = () => { + if (postTerminalSideEffects.length === 0 || postTerminalSideEffectsScheduled) return; + postTerminalSideEffectsScheduled = true; + if (replaySpool && !replaySpool.isTerminal) { + replayCompletionScheduled = true; + } + const completion = schedulePostTerminalSideEffects({ taskId, providerId: provider.id, sessionId: session.sessionId, - commit: async (signal) => { - if (signal.aborted) return; - await commitProviderFailure(); - }, + commit: (signal) => runPostTerminalSideEffects(postTerminalSideEffects, signal), }); + if (replaySpool) { + void completion.finally(() => { + if (!replaySpool.isTerminal) { + void replaySpool.abort("post_terminal_effects_incomplete"); + } + }); + } + return completion; }; const finalizedUsage = await finalizeRequestStats( session, @@ -2513,7 +2587,7 @@ export class ProxyResponseHandler { errorMessageForFinalize, undefined, false, // Gemini 非流式透传 - scheduleProviderFailure + scheduleCommittedSideEffects ); emitProxyLangfuseTrace(session, { responseHeaders: response.headers, @@ -2620,6 +2694,9 @@ export class ProxyResponseHandler { } } } finally { + if (replaySpool && !replaySpool.isTerminal && !replayCompletionScheduled) { + await replaySpool.abort("non_stream_finalize_error"); + } cleanupTaskAbortBinding(); await finalizeNonStreamDiscoveryResources(session, discoveryLeaseLifecycle); releaseSessionAgent(session); @@ -2676,6 +2753,17 @@ export class ProxyResponseHandler { const transformed = GeminiAdapter.transformResponse(responseData, false); const transformedBody = JSON.stringify(transformed); + if (transformedBody === undefined) { + throw new Error("Gemini response transform produced no JSON body"); + } + const transformedJson = JSON.parse(transformedBody) as unknown; + if ( + !transformedJson || + typeof transformedJson !== "object" || + Array.isArray(transformedJson) + ) { + throw new Error("Gemini response transform produced an invalid JSON object"); + } logger.debug( "[ResponseHandler] Transformed Gemini non-stream response to client format", @@ -2688,19 +2776,28 @@ export class ProxyResponseHandler { // ⭐ 清理传输 headers（body 已从流转为 JSON 字符串） finalResponseBodyForSnapshot = transformedBody; + const transformedHeaders = cleanResponseHeaders(response.headers); + transformedHeaders.set("content-type", "application/json"); finalResponse = new Response(transformedBody, { status: response.status, statusText: response.statusText, - headers: cleanResponseHeaders(response.headers), + headers: transformedHeaders, }); } catch (error) { logger.error("[ResponseHandler] Failed to transform Gemini non-stream response:", error); + responseTransformFailed = true; finalResponse = response; finalResponseBodyForSnapshot = null; } } } + if (responseTransformFailed) { + await abortReplayOwnership(session, "non_stream_transform_error"); + } + const replaySpool = responseTransformFailed ? null : createBufferedReplaySpool(finalResponse); + let replayCompletionScheduled = false; + // 使用 AsyncTaskManager 管理后台处理任务 const taskId = `non-stream-${messageContext?.id || `unknown-${Date.now()}`}`; const abortController = new AbortController(); @@ -2757,12 +2854,21 @@ export class ProxyResponseHandler { const scheduleCommittedSideEffects = () => { if (postTerminalSideEffects.length === 0 || postTerminalSideEffectsScheduled) return; postTerminalSideEffectsScheduled = true; - return schedulePostTerminalSideEffects({ + const completion = schedulePostTerminalSideEffects({ taskId, providerId: provider.id, sessionId: session.sessionId, commit: (signal) => runPostTerminalSideEffects(postTerminalSideEffects, signal), }); + if (replaySpool) { + const activeReplaySpool = replaySpool; + void completion.finally(() => { + if (!activeReplaySpool.isTerminal) { + void activeReplaySpool.abort("post_terminal_effects_incomplete"); + } + }); + } + return completion; }; if (messageContext) { const duration = Date.now() - session.startTime; @@ -2817,6 +2923,7 @@ export class ProxyResponseHandler { // ⭐ 非流式：读取完整响应体（会等待所有数据下载完成） const responseText = await readResponseTextWithTaskActivity(responseForLog, taskId); + const clientVisibleResponseText = finalResponseBodyForSnapshot ?? responseText; // ⭐ 响应体读取完成：清除响应超时定时器 const sessionWithCleanup = session as typeof session & { @@ -2827,6 +2934,32 @@ export class ProxyResponseHandler { } let usageMetrics: UsageMetrics | null = null; const postTerminalSideEffects: Array<() => Promise> = []; + if (replaySpool) { + const replayMalformedJson = isMalformedJsonResponseBody( + finalResponse.headers.get("content-type"), + clientVisibleResponseText + ); + const replayDetected = detectUpstreamErrorFromSseOrJsonText(responseText); + if ( + statusCode >= 200 && + statusCode < 300 && + !replayMalformedJson && + !replayDetected.isError + ) { + replaySpool.observe(RESPONSE_TEXT_ENCODER.encode(clientVisibleResponseText)); + postTerminalSideEffects.push(() => + replaySpool.completeAfterBilling(messageContext?.id ?? null) + ); + } else { + await replaySpool.abort( + replayMalformedJson + ? "non_stream_malformed_json" + : replayDetected.isError + ? "non_stream_upstream_error" + : "non_stream_http_error" + ); + } + } const usageResult = parseUsageFromResponseText(responseText, provider.providerType); usageMetrics = usageResult.usageMetrics; @@ -2919,10 +3052,7 @@ export class ProxyResponseHandler { // after 快照复用本任务已经读取到的响应文本，避免再启动一个未受 // AsyncTaskManager 管理的 clone().text() 读取分支。 - persistNonStreamAfterSnapshot( - finalResponse, - finalResponseBodyForSnapshot ?? responseText - ); + persistNonStreamAfterSnapshot(finalResponse, clientVisibleResponseText); } if (billableUsageMetrics && messageContext) { @@ -3066,12 +3196,24 @@ export class ProxyResponseHandler { const scheduleCommittedSideEffects = () => { if (postTerminalSideEffects.length === 0 || postTerminalSideEffectsScheduled) return; postTerminalSideEffectsScheduled = true; - return schedulePostTerminalSideEffects({ + if (replaySpool && !replaySpool.isTerminal) { + replayCompletionScheduled = true; + } + const completion = schedulePostTerminalSideEffects({ taskId, providerId: provider.id, sessionId: session.sessionId, commit: (signal) => runPostTerminalSideEffects(postTerminalSideEffects, signal), }); + if (replaySpool) { + const activeReplaySpool = replaySpool; + void completion.finally(() => { + if (!activeReplaySpool.isTerminal) { + void activeReplaySpool.abort("post_terminal_effects_incomplete"); + } + }); + } + return completion; }; if (messageContext) { @@ -3242,6 +3384,9 @@ export class ProxyResponseHandler { }); } } finally { + if (replaySpool && !replaySpool.isTerminal && !replayCompletionScheduled) { + await replaySpool.abort("non_stream_finalize_error"); + } cleanupTaskAbortBinding(); cleanupClientAbortListener(); await finalizeNonStreamDiscoveryResources(session, discoveryLeaseLifecycle); @@ -3634,6 +3779,7 @@ export class ProxyResponseHandler { streamEndedNormally, clientAborted, discoveryLeaseLifecycle, + null, abortReason ); latestCommitSideEffects = finalized.commitSideEffects; @@ -3706,6 +3852,7 @@ export class ProxyResponseHandler { false, clientAborted, discoveryLeaseLifecycle, + null, abortReason ); latestCommitSideEffects = finalized.commitSideEffects; @@ -4099,6 +4246,7 @@ export class ProxyResponseHandler { }; let streamFinalizationPromise: Promise | null = null; + let replayProtocolObserver: StreamProtocolObserver | null = null; const finalizeStream = ( allContent: string, streamEndedNormally: boolean, @@ -4117,6 +4265,7 @@ export class ProxyResponseHandler { streamEndedNormally, clientAborted, discoveryLeaseLifecycle, + replayProtocolObserver?.finish() ?? null, abortReason ); latestStreamCommitSideEffects = finalized.commitSideEffects @@ -4428,7 +4577,8 @@ export class ProxyResponseHandler { const isReplayableSuccess = finalized.commitSideEffects !== undefined && effectiveStatusCode >= 200 && - effectiveStatusCode < 300; + effectiveStatusCode < 300 && + hasStreamCompletionMarker(allContent, session.originalFormat); if (isReplayableSuccess) { postTerminalSideEffects.push(async () => { try { @@ -4546,6 +4696,13 @@ export class ProxyResponseHandler { // write-behind 喂入 Redis 热层，供并发/断线的相同请求 attach 跟尾。 const replaySpool = createReplaySpoolIfOwner(session, response); if (replaySpool) { + const replayProtocolFamily = + session.getEndpointPolicy().kind === "raw_passthrough" + ? null + : mapProviderTypeToFamily(provider.providerType); + replayProtocolObserver = replayProtocolFamily + ? createStreamProtocolObserver(replayProtocolFamily) + : null; try { clientAbortDrainTimeoutMs = getEnvConfig().REPLAY_MAX_DETACHED_MS; } catch { @@ -4559,6 +4716,7 @@ export class ProxyResponseHandler { streamTextAccumulator.pushBytes(value); AsyncTaskManager.touch(taskId); shadowGateObserver?.observe(value); + replayProtocolObserver?.observe(value); replaySpool?.observe(value); logger.trace("ResponseHandler: Upstream stream chunk received", { diff --git a/src/app/v1/_lib/proxy/stream-gate/sse-frames.ts b/src/app/v1/_lib/proxy/stream-gate/sse-frames.ts index c728a1533..904fe27ef 100644 --- a/src/app/v1/_lib/proxy/stream-gate/sse-frames.ts +++ b/src/app/v1/_lib/proxy/stream-gate/sse-frames.ts @@ -19,11 +19,25 @@ export interface SseFrame { data: string; } +export interface SseFrameParserOptions { + maxBufferedCharacters?: number; +} + +export class SseFrameBufferLimitError extends Error { + constructor(maxBufferedCharacters: number) { + super(`SSE parser buffered data exceeded ${maxBufferedCharacters} characters`); + this.name = "SseFrameBufferLimitError"; + } +} + export class SseFrameParser { private readonly decoder = new TextDecoder("utf-8"); private lineTail = ""; private currentEvent: string | null = null; private dataLines: string[] = []; + private dataCharacters = 0; + + constructor(private readonly options: SseFrameParserOptions = {}) {} /** 喂入一个网络 chunk，返回其中完成的帧（可能为空数组）。 */ push(chunk: Uint8Array): SseFrame[] { @@ -62,6 +76,7 @@ export class SseFrameParser { const lines = buffer.split(/\r\n|\n|\r/); // 最后一段是未完成行，保留 this.lineTail = (lines.pop() ?? "") + (holdCr ? "\r" : ""); + this.assertBufferLimit(); for (const line of lines) { const frame = this.handleLine(line); if (frame) frames.push(frame); @@ -78,10 +93,15 @@ export class SseFrameParser { } if (line.startsWith("event:")) { this.currentEvent = line.slice(6).trim(); + this.assertBufferLimit(); return null; } if (line.startsWith("data:")) { - this.dataLines.push(line.slice(5).replace(/^\s/, "")); + const data = line.slice(5).replace(/^\s/, ""); + if (this.dataLines.length > 0) this.dataCharacters += 1; + this.dataCharacters += data.length; + this.dataLines.push(data); + this.assertBufferLimit(); return null; } // id: / retry: / 未知字段：忽略 @@ -96,8 +116,20 @@ export class SseFrameParser { } const data = this.dataLines.join("\n"); this.dataLines = []; + this.dataCharacters = 0; return { eventName: event, data }; } + + private assertBufferLimit(): void { + const maxBufferedCharacters = this.options.maxBufferedCharacters; + if (maxBufferedCharacters === undefined) return; + + const bufferedCharacters = + this.lineTail.length + (this.currentEvent?.length ?? 0) + this.dataCharacters; + if (bufferedCharacters > maxBufferedCharacters) { + throw new SseFrameBufferLimitError(maxBufferedCharacters); + } + } } function stripTrailingCr(line: string): string { diff --git a/src/app/v1/_lib/proxy/stream-gate/stream-protocol-observer.ts b/src/app/v1/_lib/proxy/stream-gate/stream-protocol-observer.ts new file mode 100644 index 000000000..d78c50c46 --- /dev/null +++ b/src/app/v1/_lib/proxy/stream-gate/stream-protocol-observer.ts @@ -0,0 +1,77 @@ +import { classifyFrame, type ProtocolFamily } from "./frame-classifier"; +import { type SseFrame, SseFrameParser } from "./sse-frames"; + +export const REPLAY_PROTOCOL_OBSERVER_MAX_BUFFER_CHARACTERS = 1024 * 1024; + +export interface StreamProtocolFailure { + verdict: "error" | "malformed"; + eventName: string | null; +} + +export interface StreamProtocolObservation { + sawContent: boolean; + sawTerminal: boolean; + failure: StreamProtocolFailure | null; +} + +export interface StreamProtocolObserver { + observe(chunk: Uint8Array): void; + finish(): StreamProtocolObservation; +} + +export function createStreamProtocolObserver(family: ProtocolFamily): StreamProtocolObserver { + const parser = new SseFrameParser({ + maxBufferedCharacters: REPLAY_PROTOCOL_OBSERVER_MAX_BUFFER_CHARACTERS, + }); + const observation: StreamProtocolObservation = { + sawContent: false, + sawTerminal: false, + failure: null, + }; + let finished = false; + let disabled = false; + + const disableWithMalformedFailure = (): void => { + disabled = true; + observation.failure ??= { verdict: "malformed", eventName: null }; + }; + + const record = (frame: SseFrame): void => { + const verdict = classifyFrame(family, frame.eventName, frame.data); + if (verdict === "content") observation.sawContent = true; + if (verdict === "terminal") observation.sawTerminal = true; + if ((verdict === "error" || verdict === "malformed") && !observation.failure) { + observation.failure = { verdict, eventName: frame.eventName }; + } + }; + + return { + observe(chunk: Uint8Array): void { + if (finished || disabled || chunk.byteLength === 0) return; + try { + for (const frame of parser.push(chunk)) record(frame); + } catch { + // 旁路观察器异常不得改变客户端流或计费热路径。 + disableWithMalformedFailure(); + } + }, + + finish(): StreamProtocolObservation { + if (!finished) { + finished = true; + if (!disabled) { + try { + for (const frame of parser.finish()) record(frame); + } catch { + disableWithMalformedFailure(); + } + } + } + return { + sawContent: observation.sawContent, + sawTerminal: observation.sawTerminal, + failure: observation.failure ? { ...observation.failure } : null, + }; + }, + }; +} diff --git a/src/lib/api-client/v1/openapi-types.gen.ts b/src/lib/api-client/v1/openapi-types.gen.ts index 0e14d5feb..4de6b384e 100644 --- a/src/lib/api-client/v1/openapi-types.gen.ts +++ b/src/lib/api-client/v1/openapi-types.gen.ts @@ -12234,7 +12234,7 @@ export interface operations { /** @description Public status aggregation interval in minutes. */ publicStatusAggregationIntervalMinutes: number; /** - * @description Stream content gate mode: buffer until the first valid content frame and fail over on error or empty streams (enforce), observe divergence only (shadow), or disable (off). + * @description Stream content gate mode for ordinary requests: buffer until the first valid content frame and fail over on error or empty streams (enforce), observe divergence only (shadow), or disable (off). Replay owners always retain the pre-content safety gate. * @enum {string} */ streamGateMode: "off" | "shadow" | "enforce"; @@ -12520,7 +12520,7 @@ export interface operations { /** @description Public status aggregation interval in minutes. */ publicStatusAggregationIntervalMinutes?: number; /** - * @description Stream content gate mode: buffer until the first valid content frame and fail over on error or empty streams (enforce), observe divergence only (shadow), or disable (off). + * @description Stream content gate mode for ordinary requests: buffer until the first valid content frame and fail over on error or empty streams (enforce), observe divergence only (shadow), or disable (off). Replay owners always retain the pre-content safety gate. * @enum {string} */ streamGateMode?: "off" | "shadow" | "enforce"; @@ -12681,7 +12681,7 @@ export interface operations { /** @description Public status aggregation interval in minutes. */ publicStatusAggregationIntervalMinutes: number; /** - * @description Stream content gate mode: buffer until the first valid content frame and fail over on error or empty streams (enforce), observe divergence only (shadow), or disable (off). + * @description Stream content gate mode for ordinary requests: buffer until the first valid content frame and fail over on error or empty streams (enforce), observe divergence only (shadow), or disable (off). Replay owners always retain the pre-content safety gate. * @enum {string} */ streamGateMode: "off" | "shadow" | "enforce"; diff --git a/src/lib/api/v1/schemas/system-config.ts b/src/lib/api/v1/schemas/system-config.ts index fd0887a41..c3c06dd70 100644 --- a/src/lib/api/v1/schemas/system-config.ts +++ b/src/lib/api/v1/schemas/system-config.ts @@ -215,7 +215,7 @@ export const SystemSettingsSchema = z streamGateMode: z .enum(["off", "shadow", "enforce"]) .describe( - "Stream content gate mode: buffer until the first valid content frame and fail over on error or empty streams (enforce), observe divergence only (shadow), or disable (off)." + "Stream content gate mode for ordinary requests: buffer until the first valid content frame and fail over on error or empty streams (enforce), observe divergence only (shadow), or disable (off). Replay owners always retain the pre-content safety gate." ), affinityIgnoreClientSessionId: z .boolean() diff --git a/tests/unit/proxy/proxy-forwarder-fake-200-html.test.ts b/tests/unit/proxy/proxy-forwarder-fake-200-html.test.ts index c4db9db75..02a5bf89b 100644 --- a/tests/unit/proxy/proxy-forwarder-fake-200-html.test.ts +++ b/tests/unit/proxy/proxy-forwarder-fake-200-html.test.ts @@ -194,11 +194,146 @@ function createSession(): ProxySession { return session as ProxySession; } +function attachReplayOwner(session: ProxySession): void { + session.request.message.stream = true; + session.replayState = { + role: "owner", + ownerToken: "owner-token", + identity: { + replayId: "0123456789abcdef0123456789abcdef", + verifier: "fedcba9876543210fedcba9876543210", + scopeTag: "0011223344556677", + keyId: 11, + userId: 22, + format: "claude", + model: "claude-test", + endpoint: "/v1/messages", + }, + }; +} + describe("ProxyForwarder - fake 200 HTML body", () => { beforeEach(() => { vi.clearAllMocks(); }); + test("Replay owner 收到 malformed buffered JSON 时在提交前切换供应商", async () => { + const provider1 = createProvider({ + id: 1, + name: "p1", + key: "k1", + maxRetryAttempts: 1, + firstByteTimeoutStreamingMs: 0, + }); + const provider2 = createProvider({ + id: 2, + name: "p2", + key: "k2", + maxRetryAttempts: 1, + firstByteTimeoutStreamingMs: 0, + }); + const session = createSession(); + session.setProvider(provider1); + attachReplayOwner(session); + + mocks.pickRandomProviderWithExclusion.mockResolvedValueOnce(provider2); + const doForward = vi.spyOn(ProxyForwarder as any, "doForward"); + const malformedBody = `{"padding":"${"x".repeat(40 * 1024)}","content":[`; + const okJson = JSON.stringify({ type: "message", content: [{ type: "text", text: "ok" }] }); + + doForward.mockResolvedValueOnce( + new Response(malformedBody, { + status: 200, + headers: { + "content-type": "application/problem+json; charset=utf-8", + "content-length": String(malformedBody.length), + }, + }) + ); + doForward.mockResolvedValueOnce( + new Response(okJson, { + status: 200, + headers: { + "content-type": "application/json", + "content-length": String(okJson.length), + }, + }) + ); + + const response = await ProxyForwarder.send(session); + + await expect(response.text()).resolves.toBe(okJson); + expect(doForward).toHaveBeenCalledTimes(2); + expect(mocks.recordFailure).toHaveBeenCalledWith( + provider1.id, + expect.objectContaining({ message: "MALFORMED_BUFFERED_JSON" }) + ); + expect(mocks.recordSuccess).toHaveBeenCalledWith(provider2.id); + expect(mocks.recordSuccess).not.toHaveBeenCalledWith(provider1.id); + }); + + test("Replay owner 的非 JSON buffered body 保持原样，不触发 strict JSON fallback", async () => { + const provider = createProvider({ + id: 1, + name: "p1", + key: "k1", + maxRetryAttempts: 1, + firstByteTimeoutStreamingMs: 0, + }); + const session = createSession(); + session.setProvider(provider); + attachReplayOwner(session); + const doForward = vi.spyOn(ProxyForwarder as any, "doForward"); + const body = "{not-json"; + + doForward.mockResolvedValueOnce( + new Response(body, { + status: 200, + headers: { "content-type": "text/plain", "content-length": String(body.length) }, + }) + ); + + const response = await ProxyForwarder.send(session); + + await expect(response.text()).resolves.toBe(body); + expect(doForward).toHaveBeenCalledOnce(); + expect(mocks.recordFailure).not.toHaveBeenCalled(); + expect(mocks.recordSuccess).toHaveBeenCalledWith(provider.id); + }); + + test("Replay JSON 声明体积超过缓存上限时释放 ownership 并保持响应透传", async () => { + const provider = createProvider({ + id: 1, + name: "p1", + key: "k1", + maxRetryAttempts: 1, + firstByteTimeoutStreamingMs: 0, + }); + const session = createSession(); + session.setProvider(provider); + attachReplayOwner(session); + const doForward = vi.spyOn(ProxyForwarder as any, "doForward"); + const body = '{"ok":true}'; + + doForward.mockResolvedValueOnce( + new Response(body, { + status: 200, + headers: { + "content-type": "application/json", + "content-length": "1000000000", + }, + }) + ); + + const response = await ProxyForwarder.send(session); + + await expect(response.text()).resolves.toBe(body); + expect(session.replayState).toBeNull(); + expect(doForward).toHaveBeenCalledOnce(); + expect(mocks.recordFailure).not.toHaveBeenCalled(); + expect(mocks.recordSuccess).toHaveBeenCalledWith(provider.id); + }); + test("200 + text/html 的 HTML 页面应视为失败并切换供应商", async () => { const provider1 = createProvider({ id: 1, name: "p1", key: "k1", maxRetryAttempts: 1 }); const provider2 = createProvider({ id: 2, name: "p2", key: "k2", maxRetryAttempts: 1 }); diff --git a/tests/unit/proxy/replay-guard.test.ts b/tests/unit/proxy/replay-guard.test.ts index 699b5c0d7..1bfde400a 100644 --- a/tests/unit/proxy/replay-guard.test.ts +++ b/tests/unit/proxy/replay-guard.test.ts @@ -194,6 +194,20 @@ describe("ProxyReplayGuard：放行路径", () => { expect(storeControl.deleteChunks).not.toHaveBeenCalled(); }); + it("buffered owner 不提供 attached_live，竞争请求按 miss 路径 fail-open", async () => { + const identity = expectedIdentity(); + storeControl.getMeta.mockResolvedValueOnce( + makeMeta(identity, { status: "owning", delivery: "buffered" }) + ); + + await expect(ProxyReplayGuard.ensure(makeSession())).resolves.toBeNull(); + + expect(storeControl.readChunks).not.toHaveBeenCalled(); + expect(storeControl.findCompleted).not.toHaveBeenCalled(); + expect(storeControl.tryClaimOwner).toHaveBeenCalledWith(identity.replayId, expect.any(String)); + expect(dbControl.rows).toHaveLength(0); + }); + it("claim 竞态输掉时放行且不带 replay 角色", async () => { storeControl.tryClaimOwner.mockResolvedValueOnce(false); const session = makeSession(); @@ -341,6 +355,36 @@ describe("ProxyReplayGuard：completed 全量重放", () => { expect(String(dbControl.rows[0].blockedReason)).toContain(identity.replayId.slice(0, 12)); }); + it("completed JSON 恢复语义 headers 和原始 body，不注入 SSE headers", async () => { + const identity = expectedIdentity(); + storeControl.getMeta.mockResolvedValueOnce( + makeMeta(identity, { + status: "completed", + statusCode: 201, + delivery: "buffered", + headers: { + connection: "keep-alive", + "content-length": "999", + "content-type": "application/json; charset=utf-8", + "cache-control": "private, max-age=30", + "x-provider-request-id": "req-json", + }, + }) + ); + storeControl.readChunks.mockResolvedValueOnce(['{"id":"resp_1","status":"completed"}']); + + const response = await ProxyReplayGuard.ensure(makeSession()); + + expect(response?.status).toBe(201); + expect(response?.headers.get("content-type")).toBe("application/json; charset=utf-8"); + expect(response?.headers.get("cache-control")).toBe("private, max-age=30"); + expect(response?.headers.get("x-provider-request-id")).toBe("req-json"); + expect(response?.headers.get("connection")).toBeNull(); + expect(response?.headers.get("content-length")).toBeNull(); + expect(response?.headers.get("x-cch-replay")).toBe("completed"); + await expect(response?.text()).resolves.toBe('{"id":"resp_1","status":"completed"}'); + }); + it("热层块已过期时落 PG 持久层重放", async () => { const identity = expectedIdentity(); storeControl.getMeta.mockResolvedValueOnce(makeMeta(identity, { status: "completed" })); diff --git a/tests/unit/proxy/replay-spool.test.ts b/tests/unit/proxy/replay-spool.test.ts index 35625383c..07dcdfe35 100644 --- a/tests/unit/proxy/replay-spool.test.ts +++ b/tests/unit/proxy/replay-spool.test.ts @@ -1,10 +1,12 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { ReplayIdentity } from "@/app/v1/_lib/proxy/replay/replay-identity"; import { + abortReplayOwnership, createReplaySpoolIfOwner, getActiveReplaySpoolCount, ReplaySpool, } from "@/app/v1/_lib/proxy/replay/replay-spool"; +import type { ReplayDelivery } from "@/app/v1/_lib/proxy/replay/replay-store"; import type { ProxySession } from "@/app/v1/_lib/proxy/session"; import { logger } from "@/lib/logger"; @@ -24,6 +26,7 @@ const envControl = vi.hoisted(() => ({ const storeControl = vi.hoisted(() => { const order: string[] = []; + let ownedChunkCount = 0; const store = { appendChunks: vi.fn(async (_replayId: string, values: string[]) => { order.push(`append:${values.join("|")}`); @@ -33,6 +36,25 @@ const storeControl = vi.hoisted(() => { order.push(`meta:${meta.status}`); return true; }), + writeOwned: vi.fn( + async ( + _replayId: string, + _ownerToken: string, + _meta: { status: string }, + values: string[] = [] + ) => { + order.push(`write:${values.join("|")}`); + ownedChunkCount += values.length; + return ownedChunkCount; + } + ), + completeOwned: vi.fn( + async (_replayId: string, _ownerToken: string, meta: { status: string }) => { + order.push(`meta:${meta.status}`); + order.push("release"); + return true; + } + ), renewOwnerLease: vi.fn(async () => { order.push("renew"); return true; @@ -49,8 +71,20 @@ const storeControl = vi.hoisted(() => { deleteChunks: vi.fn(async () => { order.push("deleteChunks"); }), + abortOwned: vi.fn(async (_replayId: string, _ownerToken: string, meta: { status: string }) => { + order.push(`meta:${meta.status}`); + order.push("deleteChunks"); + order.push("release"); + return true; + }), + }; + return { + order, + store, + resetOwnedChunkCount: () => { + ownedChunkCount = 0; + }, }; - return { order, store }; }); vi.mock("@/lib/logger", () => ({ @@ -95,8 +129,18 @@ const identity: ReplayIdentity = { const encoder = new TextEncoder(); -function makeSpool(statusCode = 200, contentType = "text/event-stream"): ReplaySpool { - return new ReplaySpool(identity, "owner-token", statusCode, contentType); +function makeSpool( + statusCode = 200, + contentType = "text/event-stream", + delivery: ReplayDelivery = "stream" +): ReplaySpool { + return new ReplaySpool( + identity, + "owner-token", + statusCode, + { "content-type": contentType }, + delivery + ); } async function drainWriteChain(spool: ReplaySpool): Promise { @@ -116,10 +160,25 @@ function sseResponse(status = 200, contentType: string | null = "text/event-stre } beforeEach(() => { + vi.clearAllMocks(); envControl.enableReplay = true; envControl.maxPayloadBytes = 8 * 1024 * 1024; envControl.maxConcurrentSpools = 64; storeControl.order.length = 0; + storeControl.resetOwnedChunkCount(); + storeControl.store.abortOwned.mockImplementation( + async (_replayId: string, _ownerToken: string, meta: { status: string }) => { + storeControl.order.push(`meta:${meta.status}`); + storeControl.order.push("deleteChunks"); + storeControl.order.push("release"); + return true; + } + ); + storeControl.store.completeOwned.mockClear(); + storeControl.store.releaseOwner.mockImplementation(async () => { + storeControl.order.push("release"); + }); + storeControl.store.writeOwned.mockClear(); vi.useFakeTimers(); }); @@ -129,23 +188,20 @@ afterEach(() => { }); describe("ReplaySpool：write-behind 批量冲刷", () => { - it("小块累积由 100ms 定时批量 RPUSH，并同步续 meta 心跳与 owner 租约", async () => { + it("小块累积由 100ms 定时通过 fenced write 原子追加并更新 meta/owner 租约", async () => { const spool = makeSpool(); spool.observe(encoder.encode("data: a\n\n")); spool.observe(encoder.encode("data: b\n\n")); - expect(storeControl.store.appendChunks).not.toHaveBeenCalled(); + expect(storeControl.store.writeOwned).not.toHaveBeenCalled(); await vi.advanceTimersByTimeAsync(100); await drainWriteChain(spool); - expect(storeControl.store.appendChunks).toHaveBeenCalledTimes(1); - expect(storeControl.store.appendChunks).toHaveBeenCalledWith(identity.replayId, [ - "data: a\n\n", - "data: b\n\n", - ]); - expect(storeControl.store.setMeta).toHaveBeenCalledWith( + expect(storeControl.store.writeOwned).toHaveBeenCalledTimes(1); + expect(storeControl.store.writeOwned).toHaveBeenCalledWith( identity.replayId, + "owner-token", expect.objectContaining({ status: "owning", verifier: identity.verifier, @@ -155,11 +211,8 @@ describe("ReplaySpool：write-behind 批量冲刷", () => { chunkCount: 2, byteSize: 18, heartbeatAt: expect.any(Number), - }) - ); - expect(storeControl.store.renewOwnerLease).toHaveBeenCalledWith( - identity.replayId, - "owner-token" + }), + ["data: a\n\n", "data: b\n\n"] ); await spool.abort("test_cleanup"); @@ -171,7 +224,7 @@ describe("ReplaySpool：write-behind 批量冲刷", () => { await drainWriteChain(spool); - expect(storeControl.store.appendChunks).toHaveBeenCalledTimes(1); + expect(storeControl.store.writeOwned).toHaveBeenCalledTimes(1); await spool.abort("test_cleanup"); }); @@ -182,27 +235,43 @@ describe("ReplaySpool：write-behind 批量冲刷", () => { await vi.advanceTimersByTimeAsync(200); await drainWriteChain(spool); - expect(storeControl.store.appendChunks).not.toHaveBeenCalled(); + expect(storeControl.store.writeOwned).not.toHaveBeenCalled(); await spool.abort("test_cleanup"); }); - it("appendChunks 返回 null（Redis 不可用）时放弃 spool", async () => { - storeControl.store.appendChunks.mockResolvedValueOnce(null); + it("没有响应 chunk 时仍按 15 秒间隔续租 owner", async () => { + const spool = makeSpool(); + + await vi.advanceTimersByTimeAsync(15_000); + await drainWriteChain(spool); + + expect(storeControl.store.renewOwnerLease).toHaveBeenCalledWith( + identity.replayId, + "owner-token" + ); + await spool.abort("test_cleanup"); + }); + + it("writeOwned 返回 null（Redis 不可用）时放弃 spool", async () => { + storeControl.store.writeOwned.mockResolvedValueOnce(null); const spool = makeSpool(); spool.observe(encoder.encode("data: a\n\n")); await vi.advanceTimersByTimeAsync(100); await drainWriteChain(spool); - expect(storeControl.store.deleteEntry).toHaveBeenCalledWith(identity.replayId); - expect(storeControl.store.releaseOwner).toHaveBeenCalledWith(identity.replayId, "owner-token"); + expect(storeControl.store.abortOwned).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + expect.objectContaining({ status: "aborted", abortReason: "redis_unavailable" }) + ); expect(storeControl.store.setMeta).not.toHaveBeenCalled(); expect(getActiveReplaySpoolCount()).toBe(0); }); - it("冲刷续接体异常（setMeta 抛错）时 disable：链不被 poisoned、条目删除、租约释放", async () => { - storeControl.store.setMeta.mockRejectedValueOnce(new Error("redis exploded")); + it("冲刷续接体异常时 disable：链不被 poisoned、条目删除、租约释放", async () => { + storeControl.store.writeOwned.mockRejectedValueOnce(new Error("redis exploded")); const spool = makeSpool(); spool.observe(encoder.encode("data: a\n\n")); @@ -210,8 +279,11 @@ describe("ReplaySpool：write-behind 批量冲刷", () => { // 链必须 resolve 而非 reject（unhandled rejection 防护） await expect(drainWriteChain(spool)).resolves.toBeUndefined(); - expect(storeControl.store.deleteEntry).toHaveBeenCalledWith(identity.replayId); - expect(storeControl.store.releaseOwner).toHaveBeenCalledWith(identity.replayId, "owner-token"); + expect(storeControl.store.abortOwned).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + expect.objectContaining({ status: "aborted", abortReason: "flush_error" }) + ); expect(spool.isTerminal).toBe(true); expect(getActiveReplaySpoolCount()).toBe(0); @@ -223,20 +295,24 @@ describe("ReplaySpool：write-behind 批量冲刷", () => { }); it("bootstrap 续接体异常同样 disable 而不 poison 链", async () => { - storeControl.store.setMeta.mockRejectedValueOnce(new Error("redis exploded")); + storeControl.store.writeOwned.mockRejectedValueOnce(new Error("redis exploded")); const spool = makeSpool(); spool.bootstrap(); await expect(drainWriteChain(spool)).resolves.toBeUndefined(); - expect(storeControl.store.deleteEntry).toHaveBeenCalledWith(identity.replayId); + expect(storeControl.store.abortOwned).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + expect.objectContaining({ status: "aborted", abortReason: "flush_error" }) + ); expect(spool.isTerminal).toBe(true); expect(getActiveReplaySpoolCount()).toBe(0); }); }); describe("ReplaySpool：续租丢失 halt", () => { - it("renewOwnerLease 返回 false 时停止 spool、释放自方租约，但绝不删条目", async () => { - storeControl.store.renewOwnerLease.mockResolvedValueOnce(false); + it("writeOwned 返回 false 时停止 spool、释放自方租约，但绝不删条目", async () => { + storeControl.store.writeOwned.mockResolvedValueOnce(false); const spool = makeSpool(); spool.observe(encoder.encode("data: a\n\n")); @@ -245,7 +321,7 @@ describe("ReplaySpool：续租丢失 halt", () => { // 新 owner 可能已在写同一 LIST：只 compare-delete 自己的租约，不碰 chunks/meta expect(storeControl.store.releaseOwner).toHaveBeenCalledWith(identity.replayId, "owner-token"); - expect(storeControl.store.deleteEntry).not.toHaveBeenCalled(); + expect(storeControl.store.abortOwned).not.toHaveBeenCalled(); expect(storeControl.store.deleteChunks).not.toHaveBeenCalled(); expect(spool.isTerminal).toBe(true); expect(getActiveReplaySpoolCount()).toBe(0); @@ -258,7 +334,7 @@ describe("ReplaySpool：续租丢失 halt", () => { }); it("halt 后 complete 为 no-op", async () => { - storeControl.store.renewOwnerLease.mockResolvedValueOnce(false); + storeControl.store.writeOwned.mockResolvedValueOnce(false); const spool = makeSpool(); spool.observe(encoder.encode("data: a\n\n")); await vi.advanceTimersByTimeAsync(100); @@ -266,8 +342,8 @@ describe("ReplaySpool：续租丢失 halt", () => { await spool.completeAfterBilling(1); expect(storeControl.store.persistCompleted).not.toHaveBeenCalled(); - const metaStatuses = storeControl.store.setMeta.mock.calls.map( - (call) => (call[1] as { status: string }).status + const metaStatuses = storeControl.store.writeOwned.mock.calls.map( + (call) => (call[2] as { status: string }).status ); expect(metaStatuses).not.toContain("completed"); }); @@ -283,21 +359,24 @@ describe("ReplaySpool：超尺寸自失效", () => { // 计数同步归还；存储清理顺着 writeChain 串行执行（避免与 in-flight append 竞态） expect(getActiveReplaySpoolCount()).toBe(0); await drainWriteChain(spool); - expect(storeControl.store.deleteEntry).toHaveBeenCalledWith(identity.replayId); - expect(storeControl.store.releaseOwner).toHaveBeenCalledWith(identity.replayId, "owner-token"); + expect(storeControl.store.abortOwned).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + expect.objectContaining({ status: "aborted", abortReason: "payload_too_large" }) + ); // 已失效：后续 observe 与 complete 均为 no-op spool.observe(encoder.encode("more")); await vi.advanceTimersByTimeAsync(200); await spool.completeAfterBilling(1); - expect(storeControl.store.appendChunks).not.toHaveBeenCalled(); + expect(storeControl.store.writeOwned).not.toHaveBeenCalled(); expect(storeControl.store.persistCompleted).not.toHaveBeenCalled(); expect(storeControl.store.setMeta).not.toHaveBeenCalled(); }); }); describe("ReplaySpool：completeAfterBilling 终态屏障", () => { - it("按 尾批冲刷 -> PG 持久化 -> completed meta -> 释放租约 顺序执行", async () => { + it("按 fenced 尾批冲刷 -> PG 持久化 -> completed meta 顺序执行", async () => { const spool = makeSpool(200, "text/event-stream; charset=utf-8"); spool.observe(encoder.encode("data: hello \n\n")); spool.observe(encoder.encode("data: world\n\n")); @@ -305,7 +384,7 @@ describe("ReplaySpool：completeAfterBilling 终态屏障", () => { await spool.completeAfterBilling(42); expect(storeControl.order).toEqual([ - "append:data: hello \n\n|data: world\n\n", + "write:data: hello \n\n|data: world\n\n", "persist", "meta:completed", "release", @@ -324,8 +403,9 @@ describe("ReplaySpool：completeAfterBilling 终态屏障", () => { byteSize: 27, sourceMessageRequestId: 42, }); - expect(storeControl.store.setMeta).toHaveBeenCalledWith( + expect(storeControl.store.completeOwned).toHaveBeenCalledWith( identity.replayId, + "owner-token", expect.objectContaining({ status: "completed", messageRequestId: 42, chunkCount: 2 }) ); expect(getActiveReplaySpoolCount()).toBe(0); @@ -338,13 +418,12 @@ describe("ReplaySpool：completeAfterBilling 终态屏障", () => { await spool.completeAfterBilling(7); - const metaStatuses = storeControl.store.setMeta.mock.calls.map( - (call) => (call[1] as { status: string }).status + expect(storeControl.store.completeOwned).not.toHaveBeenCalled(); + expect(storeControl.store.abortOwned).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + expect.objectContaining({ status: "aborted", abortReason: "complete_failed" }) ); - expect(metaStatuses).not.toContain("completed"); - expect(metaStatuses).toContain("aborted"); - expect(storeControl.store.deleteChunks).toHaveBeenCalledWith(identity.replayId); - expect(storeControl.store.releaseOwner).toHaveBeenCalledWith(identity.replayId, "owner-token"); expect(logger.warn).toHaveBeenCalledWith( "[ReplaySpool] complete failed, aborting entry", expect.objectContaining({ pgPersisted: false }) @@ -352,27 +431,25 @@ describe("ReplaySpool：completeAfterBilling 终态屏障", () => { expect(getActiveReplaySpoolCount()).toBe(0); }); - it("尾批冲刷返回 null（Redis 不可用）时终止为 aborted，绝不置 completed 也不写 PG", async () => { - storeControl.store.appendChunks.mockResolvedValueOnce(null); + it("尾批 fenced write 返回 null（Redis 不可用）时终止为 aborted，绝不置 completed 也不写 PG", async () => { + storeControl.store.writeOwned.mockResolvedValueOnce(null); const spool = makeSpool(); spool.observe(encoder.encode("data: a\n\n")); await spool.completeAfterBilling(5); expect(storeControl.store.persistCompleted).not.toHaveBeenCalled(); - const metaStatuses = storeControl.store.setMeta.mock.calls.map( - (call) => (call[1] as { status: string }).status + expect(storeControl.store.completeOwned).not.toHaveBeenCalled(); + expect(storeControl.store.abortOwned).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + expect.objectContaining({ status: "aborted", abortReason: "complete_failed" }) ); - expect(metaStatuses).not.toContain("completed"); - expect(metaStatuses).toContain("aborted"); - expect(storeControl.store.deleteChunks).toHaveBeenCalledWith(identity.replayId); - expect(storeControl.store.releaseOwner).toHaveBeenCalledWith(identity.replayId, "owner-token"); expect(getActiveReplaySpoolCount()).toBe(0); }); it("persist 成功但 completed 翻转失败：日志标记 pgPersisted=true，热层封死为 aborted", async () => { - // 首次 setMeta 即 complete 续接体里的 completed 写入（无 bootstrap、定时器未触发） - storeControl.store.setMeta.mockRejectedValueOnce(new Error("redis down")); + storeControl.store.completeOwned.mockResolvedValueOnce(false); const spool = makeSpool(); spool.observe(encoder.encode("data: a\n\n")); @@ -383,22 +460,52 @@ describe("ReplaySpool：completeAfterBilling 终态屏障", () => { "[ReplaySpool] complete failed, aborting entry", expect.objectContaining({ pgPersisted: true }) ); - const metaStatuses = storeControl.store.setMeta.mock.calls.map( - (call) => (call[1] as { status: string }).status + expect(storeControl.store.abortOwned).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + expect.objectContaining({ status: "aborted", abortReason: "complete_failed" }) ); - expect(metaStatuses).toContain("aborted"); expect(getActiveReplaySpoolCount()).toBe(0); }); + it("PG 持久化阻塞期间 heartbeat 独立续租，不被 writeChain 阻塞", async () => { + let resolvePersist!: () => void; + storeControl.store.persistCompleted.mockImplementationOnce( + () => + new Promise((resolve) => { + resolvePersist = resolve; + }) + ); + const spool = makeSpool(); + spool.observe(encoder.encode("data: a\n\n")); + + const completion = spool.completeAfterBilling(8); + await vi.advanceTimersByTimeAsync(0); + expect(storeControl.store.persistCompleted).toHaveBeenCalledTimes(1); + + await vi.advanceTimersByTimeAsync(30_000); + expect(storeControl.store.renewOwnerLease).toHaveBeenCalledTimes(2); + expect(storeControl.store.completeOwned).not.toHaveBeenCalled(); + + resolvePersist(); + await completion; + expect(storeControl.store.completeOwned).toHaveBeenCalledTimes(1); + }); + it("跨 chunk 截断的 UTF-8 序列在 complete 时冲刷解码尾部", async () => { const spool = makeSpool(); // "中" (0xE4 0xB8 0xAD) 只送前两字节：observe 阶段解码挂起，complete 时 flush 出替换字符 spool.observe(new Uint8Array([0xe4, 0xb8])); - expect(storeControl.store.appendChunks).not.toHaveBeenCalled(); + expect(storeControl.store.writeOwned).not.toHaveBeenCalled(); await spool.completeAfterBilling(9); - expect(storeControl.store.appendChunks).toHaveBeenCalledWith(identity.replayId, ["\uFFFD"]); + expect(storeControl.store.writeOwned).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + expect.objectContaining({ status: "owning", chunkCount: 1 }), + ["\uFFFD"] + ); expect(storeControl.store.persistCompleted).toHaveBeenCalledWith( expect.objectContaining({ payload: "\uFFFD", byteSize: 2 }) ); @@ -422,12 +529,11 @@ describe("ReplaySpool：abort 终态", () => { await spool.abort("upstream_error"); - expect(storeControl.store.setMeta).toHaveBeenCalledWith( + expect(storeControl.store.abortOwned).toHaveBeenCalledWith( identity.replayId, + "owner-token", expect.objectContaining({ status: "aborted", abortReason: "upstream_error" }) ); - expect(storeControl.store.deleteChunks).toHaveBeenCalledWith(identity.replayId); - expect(storeControl.store.releaseOwner).toHaveBeenCalledWith(identity.replayId, "owner-token"); expect(storeControl.order).toEqual(["meta:aborted", "deleteChunks", "release"]); expect(getActiveReplaySpoolCount()).toBe(0); }); @@ -442,7 +548,7 @@ describe("ReplaySpool：abort 终态", () => { await spool.completeAfterBilling(1); expect(storeControl.order).toEqual([]); - expect(storeControl.store.appendChunks).not.toHaveBeenCalled(); + expect(storeControl.store.writeOwned).not.toHaveBeenCalled(); expect(storeControl.store.persistCompleted).not.toHaveBeenCalled(); }); }); @@ -477,7 +583,7 @@ describe("createReplaySpoolIfOwner", () => { expect(session.replayState).toBeNull(); }); - it("非 2xx 或非 SSE 响应返回 null，并释放租约、清 replayState", () => { + it("非 2xx 或 stream delivery 收到非 SSE 响应时返回 null，并释放租约、清 replayState", () => { const non2xx = makeOwnerSession(); expect(createReplaySpoolIfOwner(non2xx, sseResponse(500))).toBeNull(); expect(non2xx.replayState).toBeNull(); @@ -490,6 +596,48 @@ describe("createReplaySpoolIfOwner", () => { expect(storeControl.store.releaseOwner).toHaveBeenCalledWith(identity.replayId, "owner-token"); }); + it("buffered delivery 捕获成功 JSON 的语义 headers，并过滤传输 headers", async () => { + const response = new Response(null, { + headers: { + connection: "keep-alive", + "content-encoding": "gzip", + "content-length": "123", + "content-type": "application/json; charset=utf-8", + "set-cookie": "session=secret", + "x-provider-request-id": "req-1", + }, + }); + const spool = createReplaySpoolIfOwner(makeOwnerSession(), response, "buffered"); + expect(spool).toBeInstanceOf(ReplaySpool); + + await drainWriteChain(spool as ReplaySpool); + expect(storeControl.store.writeOwned).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + expect.objectContaining({ + delivery: "buffered", + headers: { + "content-type": "application/json; charset=utf-8", + "x-provider-request-id": "req-1", + }, + }) + ); + + spool?.observe(encoder.encode('{"ok":true}')); + await spool?.completeAfterBilling(42); + expect(storeControl.store.persistCompleted).toHaveBeenCalledWith( + expect.objectContaining({ + statusCode: 200, + headers: { + "content-type": "application/json; charset=utf-8", + "x-provider-request-id": "req-1", + }, + payload: '{"ok":true}', + sourceMessageRequestId: 42, + }) + ); + }); + it("并发 spool 达上限时返回 null，并释放落选者租约", async () => { envControl.maxConcurrentSpools = 1; const first = createReplaySpoolIfOwner(makeOwnerSession(), sseResponse()); @@ -510,9 +658,15 @@ describe("createReplaySpoolIfOwner", () => { expect(getActiveReplaySpoolCount()).toBe(1); await drainWriteChain(spool as ReplaySpool); - expect(storeControl.store.setMeta).toHaveBeenCalledWith( + expect(storeControl.store.writeOwned).toHaveBeenCalledWith( identity.replayId, - expect.objectContaining({ status: "owning", chunkCount: 0, byteSize: 0 }) + "owner-token", + expect.objectContaining({ + status: "owning", + delivery: "stream", + chunkCount: 0, + byteSize: 0, + }) ); await spool?.abort("test_cleanup"); @@ -523,11 +677,101 @@ describe("createReplaySpoolIfOwner", () => { expect(spool).toBeInstanceOf(ReplaySpool); await drainWriteChain(spool as ReplaySpool); - expect(storeControl.store.setMeta).toHaveBeenCalledWith( + expect(storeControl.store.writeOwned).toHaveBeenCalledWith( identity.replayId, + "owner-token", expect.objectContaining({ headers: { "content-type": "text/event-stream" } }) ); await spool?.abort("test_cleanup"); }); }); + +describe("abortReplayOwnership", () => { + it("以 owner token 原子终止 pre-spool 条目，并在 await 前清空 session 状态", async () => { + const session = makeOwnerSession(); + + await abortReplayOwnership(session, "forward_failed"); + + expect(session.replayState).toBeNull(); + expect(storeControl.store.abortOwned).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + expect.objectContaining({ + status: "aborted", + verifier: identity.verifier, + scopeTag: identity.scopeTag, + statusCode: 502, + headers: {}, + format: identity.format, + model: identity.model, + chunkCount: 0, + byteSize: 0, + heartbeatAt: expect.any(Number), + abortReason: "forward_failed", + }) + ); + + await abortReplayOwnership(session, "duplicate"); + expect(storeControl.store.abortOwned).toHaveBeenCalledTimes(1); + expect(vi.getTimerCount()).toBe(0); + }); + + it("非 owner 会话不触碰 store", async () => { + const session = { replayState: null } as unknown as ProxySession; + + await abortReplayOwnership(session, "forward_failed"); + + expect(storeControl.store.abortOwned).not.toHaveBeenCalled(); + }); + + it("fenced abort 返回 false 时尝试释放 owner lease", async () => { + storeControl.store.abortOwned.mockResolvedValueOnce(false); + const session = makeOwnerSession(); + + await abortReplayOwnership(session, "forward_failed"); + + expect(storeControl.store.releaseOwner).toHaveBeenCalledWith(identity.replayId, "owner-token"); + expect(logger.warn).toHaveBeenCalledWith( + "[ReplaySpool] failed to abort pre-spool replay ownership", + expect.objectContaining({ + replayId: identity.replayId.slice(0, 12), + reason: "forward_failed", + }) + ); + }); + + it("cleanup 卡住时只同步等待 100ms，随后继续 detached cleanup", async () => { + let resolveAbort: ((value: boolean) => void) | undefined; + let resolveRelease: (() => void) | undefined; + const releaseCalled = new Promise((resolve) => { + resolveRelease = resolve; + }); + storeControl.store.abortOwned.mockImplementationOnce( + async () => + await new Promise((resolve) => { + resolveAbort = resolve; + }) + ); + storeControl.store.releaseOwner.mockImplementationOnce(async () => { + storeControl.order.push("release"); + resolveRelease?.(); + }); + const session = makeOwnerSession(); + let settled = false; + const aborting = abortReplayOwnership(session, "forward_failed").then(() => { + settled = true; + }); + + await vi.advanceTimersByTimeAsync(99); + expect(settled).toBe(false); + await vi.advanceTimersByTimeAsync(1); + await aborting; + expect(settled).toBe(true); + expect(session.replayState).toBeNull(); + + resolveAbort?.(false); + await releaseCalled; + expect(storeControl.store.releaseOwner).toHaveBeenCalledWith(identity.replayId, "owner-token"); + }); +}); diff --git a/tests/unit/proxy/replay-store.test.ts b/tests/unit/proxy/replay-store.test.ts index b9c22713c..83369b272 100644 --- a/tests/unit/proxy/replay-store.test.ts +++ b/tests/unit/proxy/replay-store.test.ts @@ -134,9 +134,22 @@ function createFakeRedis() { kv.set(key, value); return "OK"; }), - // 按脚本内容分发：RPUSH+EXPIRE（list 追加）/ compare-expire（续租）/ compare-delete（释放） + // 按脚本内容分发：fenced owner write / RPUSH+EXPIRE / terminal fencing / lease fencing eval: vi.fn( async (script: string, _numkeys: number, key: string, ...args: (string | number)[]) => { + if (script.includes("LLEN") && _numkeys === 3) { + const [metaKey, chunksKey, token, replayTtl, _ownerTtl, serializedMeta, ...values] = args; + if (kv.get(key) !== token) return -1; + const listKey = String(chunksKey); + const list = lists.get(listKey) ?? []; + list.push(...values.map(String)); + lists.set(listKey, list); + if (values.length > 0 && Number(replayTtl) > 0) { + listTtls.set(listKey, Number(replayTtl)); + } + kv.set(String(metaKey), String(serializedMeta)); + return list.length; + } if (script.includes("RPUSH")) { const [ttl, ...values] = args; const list = lists.get(key) ?? []; @@ -145,6 +158,21 @@ function createFakeRedis() { if (Number(ttl) > 0) listTtls.set(key, Number(ttl)); return list.length; } + if (script.includes("SETEX") && _numkeys === 3) { + const [metaKey, chunksKey, token, ttl, serializedMeta] = args; + if (kv.get(key) !== token) return 0; + kv.set(String(metaKey), String(serializedMeta)); + lists.delete(String(chunksKey)); + kv.delete(key); + return Number(ttl) > 0 ? 1 : 0; + } + if (script.includes("SETEX") && _numkeys === 2) { + const [metaKey, token, ttl, serializedMeta] = args; + if (kv.get(key) !== token) return 0; + kv.set(String(metaKey), String(serializedMeta)); + kv.delete(key); + return Number(ttl) > 0 ? 1 : 0; + } const token = args[0] as string; if (script.includes("EXPIRE")) { return kv.get(key) === token ? 1 : 0; @@ -379,6 +407,48 @@ describe("ReplayStore：owner 租约", () => { expect(currentRedis().kv.get("cch:replay:owner:r1")).toBe("tok-b"); }); + it("writeOwned 原子追加 chunks、更新 owning meta 并续租", async () => { + const store = new ReplayStore(); + const owningMeta = makeMeta({ chunkCount: 2, byteSize: 8 }); + await store.tryClaimOwner("r1", "tok-a"); + + await expect(store.writeOwned("r1", "tok-a", owningMeta, ["part-a", "part-b"])).resolves.toBe( + 2 + ); + + expect(currentRedis().eval).toHaveBeenLastCalledWith( + expect.stringContaining("LLEN"), + 3, + "cch:replay:owner:r1", + "cch:replay:meta:r1", + "cch:replay:chunks:r1", + "tok-a", + 600, + 45, + JSON.stringify(owningMeta), + "part-a", + "part-b" + ); + await expect(store.getMeta("r1")).resolves.toEqual(owningMeta); + await expect(store.readChunks("r1", 0)).resolves.toEqual(["part-a", "part-b"]); + expect(currentRedis().listTtls.get("cch:replay:chunks:r1")).toBe(600); + }); + + it("writeOwned 拒绝 stale token，且不污染新 owner 的 meta 或 chunks", async () => { + const store = new ReplayStore(); + const currentMeta = makeMeta({ verifier: "new-owner" }); + await store.tryClaimOwner("r1", "tok-new"); + await store.writeOwned("r1", "tok-new", currentMeta, ["new-data"]); + + await expect( + store.writeOwned("r1", "tok-old", makeMeta({ verifier: "stale-owner" }), ["stale-data"]) + ).resolves.toBe(false); + + await expect(store.getMeta("r1")).resolves.toEqual(currentMeta); + await expect(store.readChunks("r1", 0)).resolves.toEqual(["new-data"]); + expect(currentRedis().kv.get("cch:replay:owner:r1")).toBe("tok-new"); + }); + it("releaseOwner 是 compare-delete：token 不匹配不删，匹配才删", async () => { const store = new ReplayStore(); await store.tryClaimOwner("r1", "tok-a"); @@ -390,6 +460,60 @@ describe("ReplayStore：owner 租约", () => { expect(currentRedis().kv.has("cch:replay:owner:r1")).toBe(false); await expect(store.tryClaimOwner("r1", "tok-b")).resolves.toBe(true); }); + + it("abortOwned 在 token 匹配时原子写 aborted、删除 chunks 并释放租约", async () => { + const store = new ReplayStore(); + const abortedMeta = makeMeta({ status: "aborted", abortReason: "forward_failed" }); + await store.tryClaimOwner("r1", "tok-a"); + await store.setMeta("r1", makeMeta()); + await store.appendChunks("r1", ["partial"]); + + await expect(store.abortOwned("r1", "tok-a", abortedMeta)).resolves.toBe(true); + + expect(currentRedis().eval).toHaveBeenLastCalledWith( + expect.stringContaining("SETEX"), + 3, + "cch:replay:owner:r1", + "cch:replay:meta:r1", + "cch:replay:chunks:r1", + "tok-a", + 600, + JSON.stringify(abortedMeta) + ); + await expect(store.getMeta("r1")).resolves.toEqual(abortedMeta); + await expect(store.readChunks("r1", 0)).resolves.toEqual([]); + expect(currentRedis().kv.has("cch:replay:owner:r1")).toBe(false); + }); + + it("abortOwned 在 token 不匹配时不覆盖新 owner 的 meta、chunks 或租约", async () => { + const store = new ReplayStore(); + const currentMeta = makeMeta({ verifier: "new-owner" }); + await store.tryClaimOwner("r1", "tok-new"); + await store.setMeta("r1", currentMeta); + await store.appendChunks("r1", ["new-data"]); + + await expect( + store.abortOwned("r1", "tok-old", makeMeta({ status: "aborted", abortReason: "stale_owner" })) + ).resolves.toBe(false); + + await expect(store.getMeta("r1")).resolves.toEqual(currentMeta); + await expect(store.readChunks("r1", 0)).resolves.toEqual(["new-data"]); + expect(currentRedis().kv.get("cch:replay:owner:r1")).toBe("tok-new"); + }); + + it("completeOwned 仅在 token 匹配时原子写 completed meta 并释放租约", async () => { + const store = new ReplayStore(); + const completedMeta = makeMeta({ status: "completed", chunkCount: 2 }); + await store.tryClaimOwner("r1", "tok-a"); + + await expect(store.completeOwned("r1", "tok-other", completedMeta)).resolves.toBe(false); + expect(currentRedis().kv.get("cch:replay:owner:r1")).toBe("tok-a"); + await expect(store.getMeta("r1")).resolves.toBeNull(); + + await expect(store.completeOwned("r1", "tok-a", completedMeta)).resolves.toBe(true); + await expect(store.getMeta("r1")).resolves.toEqual(completedMeta); + expect(currentRedis().kv.has("cch:replay:owner:r1")).toBe(false); + }); }); describe("ReplayStore：PG 完成持久层", () => { diff --git a/tests/unit/proxy/response-handler-nonstream-terminal.test.ts b/tests/unit/proxy/response-handler-nonstream-terminal.test.ts index f72581739..ce9428fef 100644 --- a/tests/unit/proxy/response-handler-nonstream-terminal.test.ts +++ b/tests/unit/proxy/response-handler-nonstream-terminal.test.ts @@ -1,5 +1,6 @@ import { Context } from "hono"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { GeminiAdapter } from "@/app/v1/_lib/gemini/adapter"; 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"; @@ -19,6 +20,29 @@ const mocks = vi.hoisted(() => ({ trackerEnd: vi.fn(), })); +const replayControl = vi.hoisted(() => { + const state = { enabled: false, terminal: false }; + const spool = { + observe: vi.fn(), + completeAfterBilling: vi.fn(async () => { + state.terminal = true; + }), + abort: vi.fn(async () => { + state.terminal = true; + }), + }; + Object.defineProperty(spool, "isTerminal", { + get: () => state.terminal, + }); + return { + state, + spool, + create: vi.fn(() => (state.enabled ? spool : null)), + abortOwnership: vi.fn(async () => undefined), + releaseOwnership: vi.fn(), + }; +}); + vi.mock("@/app/v1/_lib/proxy/response-fixer", () => ({ ResponseFixer: { process: async (_session: ProxySession, response: Response) => response }, })); @@ -62,6 +86,11 @@ vi.mock("@/lib/logger", () => ({ vi.mock("@/lib/proxy-status-tracker", () => ({ ProxyStatusTracker: { getInstance: () => ({ endRequest: mocks.trackerEnd }) }, })); +vi.mock("@/app/v1/_lib/proxy/replay/replay-spool", () => ({ + abortReplayOwnership: replayControl.abortOwnership, + createReplaySpoolIfOwner: replayControl.create, + releaseReplayOwnership: replayControl.releaseOwnership, +})); vi.mock("@/repository/message", () => ({ addMessageRequestHedgeLoserCost: vi.fn(), updateMessageRequestCostWithBreakdown: vi.fn(), @@ -116,7 +145,7 @@ const MESSAGE = { user: USER, } satisfies MessageContext; -function createProvider(): Provider { +function createProvider(providerType: Provider["providerType"] = "claude"): Provider { return { activeTimeEnd: null, activeTimeStart: null, @@ -166,7 +195,7 @@ function createProvider(): Provider { name: "nonstream-terminal-provider", preserveClientIp: false, priority: 1, - providerType: "claude", + providerType, providerVendorId: null, proxyFallbackToDirect: false, proxyUrl: null, @@ -184,14 +213,17 @@ function createProvider(): Provider { } satisfies Provider; } -async function createSession(releaseAgent: () => void): Promise { +async function createSession( + releaseAgent: () => void, + options: { stream?: boolean; providerType?: Provider["providerType"] } = {} +): Promise { const request = new Request("https://hub.test/v1/messages", { - body: JSON.stringify({ messages: [], stream: false }), + body: JSON.stringify({ messages: [], stream: options.stream ?? false }), headers: { "content-type": "application/json" }, method: "POST", }); const session = await ProxySession.fromContext(new Context(request)); - session.setProvider(createProvider()); + session.setProvider(createProvider(options.providerType)); session.setMessageContext(MESSAGE); Object.defineProperty(session, "releaseAgent", { value: releaseAgent, writable: true }); return session; @@ -228,6 +260,8 @@ describe("ProxyResponseHandler.dispatch nonstream terminal behavior", () => { return true; }); mocks.recordFailure.mockResolvedValue(undefined); + replayControl.state.enabled = false; + replayControl.state.terminal = false; }); it("uses conditional persistence before recording a nonstream provider failure", async () => { @@ -279,4 +313,193 @@ describe("ProxyResponseHandler.dispatch nonstream terminal behavior", () => { expect(mocks.trackerEnd).toHaveBeenCalledWith(USER.id, MESSAGE.id); expect(releaseAgent).toHaveBeenCalledOnce(); }); + + it("stream 请求收到成功 JSON 时，在 durable terminal commit 后缓存客户端 body", async () => { + replayControl.state.enabled = true; + const releaseAgent = vi.fn(); + const session = await createSession(releaseAgent, { stream: true }); + const body = '{"id":"resp_1","status":"completed"}'; + + const returned = await ProxyResponseHandler.dispatch( + session, + new Response(body, { + headers: { + "content-type": "application/json", + "x-provider-request-id": "req-1", + }, + }) + ); + await expect(returned.text()).resolves.toBe(body); + const settlements = await settleTasks(); + + expect(settlements.every(({ status }) => status === "fulfilled")).toBe(true); + expect(replayControl.create).toHaveBeenCalledWith(session, expect.any(Response), "buffered"); + expect(replayControl.spool.observe).toHaveBeenCalledOnce(); + const observed = replayControl.spool.observe.mock.calls[0]?.[0] as Uint8Array; + expect(new TextDecoder().decode(observed)).toBe(body); + expect(replayControl.spool.completeAfterBilling).toHaveBeenCalledWith(MESSAGE.id); + expect(mocks.durable.mock.invocationCallOrder[0]).toBeLessThan( + replayControl.spool.completeAfterBilling.mock.invocationCallOrder[0] + ); + expect(replayControl.spool.abort).not.toHaveBeenCalled(); + }); + + it.each(["application/json", "application/problem+json; charset=utf-8"])( + "声明为 %s 的 malformed body 不得 completed Replay", + async (contentType) => { + replayControl.state.enabled = true; + const session = await createSession(vi.fn(), { stream: true }); + const body = '{"id":"truncated"'; + + const returned = await ProxyResponseHandler.dispatch( + session, + new Response(body, { headers: { "content-type": contentType } }) + ); + await expect(returned.text()).resolves.toBe(body); + await settleTasks(); + + expect(replayControl.spool.observe).not.toHaveBeenCalled(); + expect(replayControl.spool.completeAfterBilling).not.toHaveBeenCalled(); + expect(replayControl.spool.abort).toHaveBeenCalledWith("non_stream_malformed_json"); + } + ); + + it("非 JSON buffered body 不做 JSON parse，可正常 completed Replay", async () => { + replayControl.state.enabled = true; + const session = await createSession(vi.fn(), { stream: true }); + const body = "{not-json"; + + const returned = await ProxyResponseHandler.dispatch( + session, + new Response(body, { headers: { "content-type": "text/plain" } }) + ); + await expect(returned.text()).resolves.toBe(body); + await settleTasks(); + + expect(replayControl.spool.observe).toHaveBeenCalledOnce(); + expect(replayControl.spool.completeAfterBilling).toHaveBeenCalledWith(MESSAGE.id); + expect(replayControl.spool.abort).not.toHaveBeenCalled(); + }); + + it("terminal persistence 全部失败时 abort buffered Replay，绝不 completed", async () => { + replayControl.state.enabled = true; + mocks.durable.mockRejectedValueOnce(new Error("primary unavailable")); + mocks.conditional.mockRejectedValueOnce(new Error("fallback unavailable")); + const session = await createSession(vi.fn(), { stream: true }); + + const returned = await ProxyResponseHandler.dispatch( + session, + new Response('{"ok":true}', { + headers: { "content-type": "application/json" }, + }) + ); + await returned.text(); + const settlements = await settleTasks(); + + expect(settlements.some(({ status }) => status === "rejected")).toBe(true); + expect(replayControl.spool.completeAfterBilling).not.toHaveBeenCalled(); + expect(replayControl.spool.abort).toHaveBeenCalledWith("non_stream_finalize_error"); + }); + + it("Gemini 转换路径缓存转换后的客户端 body，而不是 upstream body", async () => { + replayControl.state.enabled = true; + const session = await createSession(vi.fn(), { stream: true, providerType: "gemini" }); + session.setOriginalFormat("claude"); + const upstreamBody = + '{"candidates":[{"content":{"role":"model","parts":[{"text":"hello"}]},"finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash"}'; + + const returned = await ProxyResponseHandler.dispatch( + session, + new Response(upstreamBody, { + headers: { "content-type": "application/json" }, + }) + ); + const returnedBody = await returned.text(); + await settleTasks(); + + expect(returned.headers.get("content-type")).toContain("application/json"); + + const observed = replayControl.spool.observe.mock.calls[0]?.[0] as Uint8Array; + const observedBody = new TextDecoder().decode(observed); + expect(observedBody).toBe(returnedBody); + expect(observedBody).not.toBe(upstreamBody); + expect(JSON.parse(observedBody)).toMatchObject({ + object: "chat.completion", + model: "gemini-2.5-flash", + choices: [ + { + index: 0, + message: { role: "assistant", content: "hello" }, + finish_reason: "stop", + }, + ], + }); + expect(replayControl.spool.completeAfterBilling).toHaveBeenCalledWith(MESSAGE.id); + }); + + it("Gemini 转换失败时返回原响应，但 abort Replay ownership 且不创建缓存", async () => { + replayControl.state.enabled = true; + const session = await createSession(vi.fn(), { stream: true, providerType: "gemini" }); + session.setOriginalFormat("claude"); + const upstreamBody = '{"candidates":[]}'; + vi.spyOn(GeminiAdapter, "transformResponse").mockImplementationOnce(() => { + throw new Error("transform failed"); + }); + + const returned = await ProxyResponseHandler.dispatch( + session, + new Response(upstreamBody, { + headers: { "content-type": "application/json" }, + }) + ); + await expect(returned.text()).resolves.toBe(upstreamBody); + await settleTasks(); + + expect(replayControl.abortOwnership).toHaveBeenCalledWith( + session, + "non_stream_transform_error" + ); + expect(replayControl.create).not.toHaveBeenCalled(); + expect(replayControl.spool.observe).not.toHaveBeenCalled(); + expect(replayControl.spool.completeAfterBilling).not.toHaveBeenCalled(); + }); + + it("Gemini passthrough 路径缓存未转换的客户端 body", async () => { + replayControl.state.enabled = true; + const session = await createSession(vi.fn(), { stream: true, providerType: "gemini" }); + session.setOriginalFormat("gemini"); + const upstreamBody = + '{"candidates":[{"content":{"role":"model","parts":[{"text":"hello"}]},"finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash"}'; + + const returned = await ProxyResponseHandler.dispatch( + session, + new Response(upstreamBody, { + headers: { "content-type": "application/json" }, + }) + ); + await expect(returned.text()).resolves.toBe(upstreamBody); + await settleTasks(); + + const observed = replayControl.spool.observe.mock.calls[0]?.[0] as Uint8Array; + expect(new TextDecoder().decode(observed)).toBe(upstreamBody); + expect(replayControl.spool.completeAfterBilling).toHaveBeenCalledWith(MESSAGE.id); + }); + + it("Gemini passthrough 的 malformed JSON 会 abort Replay 且不发布 chunks", async () => { + replayControl.state.enabled = true; + const session = await createSession(vi.fn(), { stream: true, providerType: "gemini" }); + session.setOriginalFormat("gemini"); + const upstreamBody = '{"candidates":['; + + const returned = await ProxyResponseHandler.dispatch( + session, + new Response(upstreamBody, { headers: { "content-type": "application/json" } }) + ); + await expect(returned.text()).resolves.toBe(upstreamBody); + await settleTasks(); + + expect(replayControl.spool.observe).not.toHaveBeenCalled(); + expect(replayControl.spool.completeAfterBilling).not.toHaveBeenCalled(); + expect(replayControl.spool.abort).toHaveBeenCalledWith("non_stream_malformed_json"); + }); }); diff --git a/tests/unit/proxy/response-handler-stream-terminal.test.ts b/tests/unit/proxy/response-handler-stream-terminal.test.ts index efe5ddd8e..260fa97a3 100644 --- a/tests/unit/proxy/response-handler-stream-terminal.test.ts +++ b/tests/unit/proxy/response-handler-stream-terminal.test.ts @@ -2,6 +2,7 @@ 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 { setDeferredStreamingFinalization } from "@/app/v1/_lib/proxy/stream-finalization"; import type { Key } from "@/types/key"; import type { Provider } from "@/types/provider"; import type { User } from "@/types/user"; @@ -15,6 +16,9 @@ const mocks = vi.hoisted(() => ({ >(), tasks: Array.from>([]), trackerEnd: vi.fn(), + replayObserve: vi.fn(), + replayComplete: vi.fn(async () => {}), + replayAbort: vi.fn(async () => {}), })); vi.mock("@/app/v1/_lib/proxy/response-fixer", () => ({ @@ -54,6 +58,19 @@ vi.mock("@/lib/logger", () => ({ vi.mock("@/lib/proxy-status-tracker", () => ({ ProxyStatusTracker: { getInstance: () => ({ endRequest: mocks.trackerEnd }) }, })); +vi.mock("@/app/v1/_lib/proxy/replay/replay-spool", () => ({ + abortReplayOwnership: vi.fn(async () => undefined), + createReplaySpoolIfOwner: (session: ProxySession) => + session.replayState?.role === "owner" + ? { + abort: mocks.replayAbort, + completeAfterBilling: mocks.replayComplete, + isTerminal: false, + observe: mocks.replayObserve, + } + : null, + releaseReplayOwnership: vi.fn(), +})); vi.mock("@/repository/message", () => ({ addMessageRequestHedgeLoserCost: vi.fn(), updateMessageRequestCostWithBreakdown: vi.fn(), @@ -201,15 +218,18 @@ async function settleTasks(): Promise { } } -function sseResponse(body: BodyInit): Response { - return new Response(body, { headers: { "content-type": "text/event-stream" } }); +function sseResponse(body: BodyInit, status = 200): Response { + return new Response(body, { status, headers: { "content-type": "text/event-stream" } }); } describe("ProxyResponseHandler.dispatch stream terminal behavior", () => { beforeEach(() => { mocks.tasks.length = 0; vi.clearAllMocks(); - mocks.durable.mockResolvedValue(true); + mocks.durable.mockImplementation(async (_id, _details, options) => { + await options?.onCommitted?.(); + return true; + }); }); it("persists a naturally completed stream and releases its transport", async () => { @@ -277,4 +297,113 @@ describe("ProxyResponseHandler.dispatch stream terminal behavior", () => { expect(cancelSource).toHaveBeenCalledOnce(); expect(releaseAgent).toHaveBeenCalledOnce(); }); + + it("aborts Replay when a late protocol error is outside the bounded text snapshot", async () => { + const { session } = await createSession({}); + session.setProvider({ ...createProvider(), providerType: "codex" }); + session.originalFormat = "response"; + session.replayState = { + role: "owner", + ownerToken: "owner-token", + identity: { + replayId: "replay-id", + verifier: "verifier", + scopeTag: "scope-tag", + keyId: KEY.id, + userId: USER.id, + format: "response", + model: "gpt-test", + endpoint: "/v1/responses", + }, + }; + setDeferredStreamingFinalization(session, { + providerId: session.provider?.id ?? 0, + providerName: session.provider?.name ?? "provider", + providerPriority: session.provider?.priority ?? 0, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: null, + endpointUrl: session.provider?.url ?? "", + upstreamStatusCode: 200, + bindingIntent: "none", + }); + + const headFiller = "x".repeat(2 * 1024 * 1024); + const tailFiller = "y".repeat(10 * 1024 * 1024); + const body = [ + 'event: response.output_text.delta\ndata: {"type":"response.output_text.delta","delta":"ok"}\n\n', + `event: response.in_progress\ndata: ${JSON.stringify({ type: "response.in_progress", headFiller })}\n\n`, + 'event: response.failed\ndata: {"type":"response.failed","response":{"status":"failed"}}\n\n', + `event: response.in_progress\ndata: ${JSON.stringify({ type: "response.in_progress", tailFiller })}\n\n`, + 'event: response.completed\ndata: {"type":"response.completed","response":{"status":"completed"}}\n\n', + ].join(""); + + const returned = await ProxyResponseHandler.dispatch(session, sseResponse(body)); + await returned.text(); + await settleTasks(); + + expect(mocks.replayAbort).toHaveBeenCalledWith( + expect.stringContaining("UPSTREAM_PROTOCOL_ERROR") + ); + expect(mocks.replayComplete).not.toHaveBeenCalled(); + expect(mocks.durable).toHaveBeenCalledWith( + MESSAGE.id, + expect.objectContaining({ statusCode: 502 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + }); + + it("aborts Replay when a successful non-200 2xx stream contains a late protocol error", async () => { + const { session } = await createSession({}); + session.setProvider({ ...createProvider(), providerType: "codex" }); + session.originalFormat = "response"; + session.replayState = { + role: "owner", + ownerToken: "owner-token", + identity: { + replayId: "replay-id-201", + verifier: "verifier", + scopeTag: "scope-tag", + keyId: KEY.id, + userId: USER.id, + format: "response", + model: "gpt-test", + endpoint: "/v1/responses", + }, + }; + setDeferredStreamingFinalization(session, { + providerId: session.provider?.id ?? 0, + providerName: session.provider?.name ?? "provider", + providerPriority: session.provider?.priority ?? 0, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: null, + endpointUrl: session.provider?.url ?? "", + upstreamStatusCode: 201, + bindingIntent: "none", + }); + const body = [ + 'event: response.output_text.delta\ndata: {"type":"response.output_text.delta","delta":"ok"}\n\n', + 'event: response.completed\ndata: {"type":"response.completed","response":{"status":"completed"}}\n\n', + 'event: response.failed\ndata: {"type":"response.failed","response":{"status":"failed"}}\n\n', + ].join(""); + + const returned = await ProxyResponseHandler.dispatch(session, sseResponse(body, 201)); + await returned.text(); + await settleTasks(); + + expect(mocks.replayAbort).toHaveBeenCalledWith( + expect.stringContaining("UPSTREAM_PROTOCOL_ERROR") + ); + expect(mocks.replayComplete).not.toHaveBeenCalled(); + expect(mocks.durable).toHaveBeenCalledWith( + MESSAGE.id, + expect.objectContaining({ statusCode: 502 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + }); }); diff --git a/tests/unit/proxy/stream-gate-forwarder-integration.test.ts b/tests/unit/proxy/stream-gate-forwarder-integration.test.ts index 61adc58f0..263a264a6 100644 --- a/tests/unit/proxy/stream-gate-forwarder-integration.test.ts +++ b/tests/unit/proxy/stream-gate-forwarder-integration.test.ts @@ -188,6 +188,73 @@ const MESSAGE_STOP_FRAME = sseFrame("message_stop", { type: "message_stop" }); // failover 后获胜供应商的正常内容流 const WINNER_FRAMES = [MESSAGE_START_FRAME, CONTENT_DELTA_FRAME, MESSAGE_STOP_FRAME]; +function createOpenAiResponsesOverloadFrames(): string[] { + const requestEcho = "x".repeat(500_000); + const error = { + type: "server_error", + code: "server_is_overloaded", + message: "Our servers are currently overloaded. Please try again later.", + }; + + return [ + sseFrame("response.created", { + type: "response.created", + response: { id: "resp_failed", status: "in_progress", instructions: requestEcho }, + }), + sseFrame("response.in_progress", { + type: "response.in_progress", + response: { id: "resp_failed", status: "in_progress", instructions: requestEcho }, + }), + sseFrame("error", { type: "error", ...error }), + sseFrame("response.failed", { + type: "response.failed", + response: { id: "resp_failed", status: "failed", error }, + }), + ]; +} + +const OPENAI_RESPONSES_WINNER_FRAMES = [ + sseFrame("response.output_text.delta", { + type: "response.output_text.delta", + delta: "winner", + }), + sseFrame("response.completed", { + type: "response.completed", + response: { id: "resp_winner", status: "completed" }, + }), +]; + +type ReplayGateCase = { + name: string; + providerType: Provider["providerType"]; + endpoint: "/v1/messages" | "/v1/responses"; + format: "claude" | "response"; + failedFrames: () => string[]; + winnerFrames: string[]; + failedMarkers: string[]; +}; + +const REPLAY_GATE_CASES: ReplayGateCase[] = [ + { + name: "OpenAI Responses server overload", + providerType: "codex", + endpoint: "/v1/responses", + format: "response", + failedFrames: createOpenAiResponsesOverloadFrames, + winnerFrames: OPENAI_RESPONSES_WINNER_FRAMES, + failedMarkers: ["server_is_overloaded", "resp_failed"], + }, + { + name: "Anthropic overloaded_error", + providerType: "claude", + endpoint: "/v1/messages", + format: "claude", + failedFrames: () => [PING_FRAME, ERROR_FRAME], + winnerFrames: WINNER_FRAMES, + failedMarkers: ["overloaded_error"], + }, +]; + function createProvider(overrides: Partial = {}): Provider { return { id: 1, @@ -297,6 +364,28 @@ function createSession(clientAbortSignal: AbortSignal | null = null): ProxySessi return session as ProxySession; } +function attachReplayOwner(session: ProxySession, testCase: ReplayGateCase): void { + Object.assign(session, { + requestUrl: new URL(`https://example.com${testCase.endpoint}`), + originalFormat: testCase.format, + endpointPolicy: resolveEndpointPolicy(testCase.endpoint), + }); + session.replayState = { + role: "owner", + ownerToken: "owner-token", + identity: { + replayId: "replay-id", + verifier: "verifier", + scopeTag: "scope-tag", + keyId: 1, + userId: 1, + format: testCase.format, + model: "claude-test", + endpoint: testCase.endpoint, + }, + }; +} + function createSseResponse(frames: string[]): Response { const encoder = new TextEncoder(); const stream = new ReadableStream({ @@ -480,5 +569,59 @@ describe("F1 stream content gate x ProxyForwarder sequential path", () => { expect(mocks.pickRandomProviderWithExclusion).not.toHaveBeenCalled(); expect(mocks.recordFailure).not.toHaveBeenCalled(); }); + + test.each(REPLAY_GATE_CASES)( + "Replay owner 仍拦截首内容前的 $name，并切换到成功供应商", + async (testCase) => { + const provider1 = createProvider({ + id: 1, + name: "replay-p1", + providerType: testCase.providerType, + }); + const provider2 = createProvider({ + id: 2, + name: "replay-p2", + providerType: testCase.providerType, + }); + const session = createSession(); + session.setProvider(provider1); + attachReplayOwner(session, testCase); + + mocks.pickRandomProviderWithExclusion.mockResolvedValueOnce(provider2); + const doForward = spyOnDoForward(); + doForward.mockImplementationOnce(async () => createSseResponse(testCase.failedFrames())); + doForward.mockImplementationOnce(async () => createSseResponse(testCase.winnerFrames)); + + const response = await ProxyForwarder.send(session); + const text = await response.text(); + + expect(doForward).toHaveBeenCalledTimes(2); + expect((doForward.mock.calls[1] as unknown[])[1]).toMatchObject({ id: provider2.id }); + expect(text).toBe(testCase.winnerFrames.join("")); + for (const marker of testCase.failedMarkers) { + expect(text).not.toContain(marker); + } + expect(mocks.recordFailure).toHaveBeenCalledWith(provider1.id, expect.any(Error)); + expect(session.provider?.id).toBe(provider2.id); + } + ); + + test("Replay owner 在所有 precommit attempt 失败后立即释放所有权", async () => { + const provider = createProvider({ id: 1, name: "replay-only", providerType: "codex" }); + const session = createSession(); + session.setProvider(provider); + attachReplayOwner(session, REPLAY_GATE_CASES[0]); + + mocks.pickRandomProviderWithExclusion.mockResolvedValueOnce(null); + const doForward = spyOnDoForward(); + doForward.mockImplementationOnce(async () => + createSseResponse(createOpenAiResponsesOverloadFrames()) + ); + + await expect(ProxyForwarder.send(session)).rejects.toThrow(); + + expect(doForward).toHaveBeenCalledTimes(1); + expect(session.replayState).toBeNull(); + }); }); }); diff --git a/tests/unit/proxy/stream-gate-protocol-observer.test.ts b/tests/unit/proxy/stream-gate-protocol-observer.test.ts new file mode 100644 index 000000000..a7e8ff1cc --- /dev/null +++ b/tests/unit/proxy/stream-gate-protocol-observer.test.ts @@ -0,0 +1,158 @@ +import { describe, expect, test } from "vitest"; +import { + createStreamProtocolObserver, + REPLAY_PROTOCOL_OBSERVER_MAX_BUFFER_CHARACTERS, +} from "@/app/v1/_lib/proxy/stream-gate/stream-protocol-observer"; + +const encoder = new TextEncoder(); + +function observeAtEveryBoundary(stream: string): void { + const bytes = encoder.encode(stream); + for (let split = 0; split <= bytes.length; split++) { + const observer = createStreamProtocolObserver("openai-responses"); + observer.observe(bytes.subarray(0, split)); + observer.observe(bytes.subarray(split)); + expect(observer.finish()).toEqual({ + sawContent: true, + sawTerminal: true, + failure: null, + }); + } +} + +describe("StreamProtocolObserver", () => { + test("在任意网络 chunk boundary 下都识别 content 与成功 terminal", () => { + observeAtEveryBoundary( + [ + 'event: response.output_text.delta\r\ndata: {"type":"response.output_text.delta","delta":"你好"}\r\n\r\n', + 'event: response.completed\r\ndata: {"type":"response.completed","response":{"status":"completed"}}\r\n\r\n', + ].join("") + ); + }); + + test("首内容后继续识别失败帧，并在 EOF 冲刷无空行结尾的 terminal", () => { + const observer = createStreamProtocolObserver("openai-responses"); + const stream = [ + 'event: response.output_text.delta\ndata: {"type":"response.output_text.delta","delta":"ok"}\n\n', + 'event: response.failed\ndata: {"type":"response.failed","response":{"status":"failed"}}\n\n', + 'event: response.completed\ndata: {"type":"response.completed","response":{"status":"completed"}}', + ].join(""); + const bytes = encoder.encode(stream); + + observer.observe(bytes.subarray(0, 17)); + observer.observe(bytes.subarray(17, bytes.length - 11)); + observer.observe(bytes.subarray(bytes.length - 11)); + const result = observer.finish(); + + expect(result.sawContent).toBe(true); + expect(result.sawTerminal).toBe(true); + expect(result.failure).toEqual({ + verdict: "error", + eventName: "response.failed", + }); + }); + + test("把 malformed data 记录为终态失败", () => { + const observer = createStreamProtocolObserver("anthropic"); + observer.observe( + encoder.encode( + [ + 'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"ok"}}\n\n', + "event: message_stop\ndata: {not-json}\n\n", + ].join("") + ) + ); + + expect(observer.finish()).toEqual({ + sawContent: true, + sawTerminal: false, + failure: { verdict: "malformed", eventName: "message_stop" }, + }); + }); + + test("成功 Anthropic tool_use/message_stop 可 completed", () => { + const observer = createStreamProtocolObserver("anthropic"); + observer.observe( + encoder.encode( + [ + 'event: content_block_start\ndata: {"type":"content_block_start","content_block":{"type":"tool_use","name":"lookup"}}\n\n', + 'event: message_stop\ndata: {"type":"message_stop"}\n\n', + ].join("") + ) + ); + + expect(observer.finish()).toEqual({ + sawContent: true, + sawTerminal: true, + failure: null, + }); + }); + + test("缺失 terminal 时保留 sawContent，但不能证明流已成功完成", () => { + const observer = createStreamProtocolObserver("openai-chat"); + observer.observe(encoder.encode('data: {"choices":[{"delta":{"content":"partial"}}]}\n\n')); + + expect(observer.finish()).toEqual({ + sawContent: true, + sawTerminal: false, + failure: null, + }); + }); + + test("OpenAI Chat 的 [DONE] 是成功 terminal", () => { + const observer = createStreamProtocolObserver("openai-chat"); + observer.observe( + encoder.encode( + ['data: {"choices":[{"delta":{"content":"ok"}}]}\n\n', "data: [DONE]\n\n"].join("") + ) + ); + + expect(observer.finish()).toEqual({ + sawContent: true, + sawTerminal: true, + failure: null, + }); + }); + + test("finish 冲刷没有尾部空行的 malformed frame", () => { + const observer = createStreamProtocolObserver("openai-responses"); + observer.observe(encoder.encode("event: response.completed\ndata: not-json")); + + expect(observer.finish()).toEqual({ + sawContent: false, + sawTerminal: false, + failure: { verdict: "malformed", eventName: "response.completed" }, + }); + }); + + test("超长未终止 SSE 行会 fail closed，而不是被视为干净流", () => { + const observer = createStreamProtocolObserver("openai-responses"); + observer.observe( + encoder.encode(`data: ${"x".repeat(REPLAY_PROTOCOL_OBSERVER_MAX_BUFFER_CHARACTERS + 1)}`) + ); + + expect(observer.finish()).toEqual({ + sawContent: false, + sawTerminal: false, + failure: { verdict: "malformed", eventName: null }, + }); + }); + + test("首内容后 observer 超限仍保留 content，并阻止 Replay completed", () => { + const observer = createStreamProtocolObserver("openai-responses"); + observer.observe( + encoder.encode( + 'event: response.output_text.delta\ndata: {"type":"response.output_text.delta","delta":"ok"}\n\n' + ) + ); + observer.observe( + encoder.encode(`data: ${"x".repeat(REPLAY_PROTOCOL_OBSERVER_MAX_BUFFER_CHARACTERS + 1)}`) + ); + + expect(observer.finish()).toEqual({ + sawContent: true, + sawTerminal: false, + failure: { verdict: "malformed", eventName: null }, + }); + }); +}); diff --git a/tests/unit/proxy/stream-gate-sse-frames.test.ts b/tests/unit/proxy/stream-gate-sse-frames.test.ts index 5db9d824d..80339ba71 100644 --- a/tests/unit/proxy/stream-gate-sse-frames.test.ts +++ b/tests/unit/proxy/stream-gate-sse-frames.test.ts @@ -1,5 +1,9 @@ import { describe, expect, it } from "vitest"; -import { parseSseBody, SseFrameParser } from "@/app/v1/_lib/proxy/stream-gate/sse-frames"; +import { + parseSseBody, + SseFrameBufferLimitError, + SseFrameParser, +} from "@/app/v1/_lib/proxy/stream-gate/sse-frames"; function collectAll(parser: SseFrameParser, chunks: Uint8Array[]) { const frames = chunks.flatMap((chunk) => parser.push(chunk)); @@ -95,4 +99,26 @@ describe("SseFrameParser", () => { } expect(collectAll(parser, chunks)).toEqual(expected); }); + + it("rejects an unterminated line that exceeds the configured retained buffer", () => { + const parser = new SseFrameParser({ maxBufferedCharacters: 16 }); + + expect(() => parser.push(new TextEncoder().encode(`data: ${"x".repeat(20)}`))).toThrow( + SseFrameBufferLimitError + ); + }); + + it("counts accumulated data lines before an event is dispatched", () => { + const parser = new SseFrameParser({ maxBufferedCharacters: 10 }); + + expect(() => parser.push(new TextEncoder().encode("data: 12345\ndata: 67890\n"))).toThrow( + SseFrameBufferLimitError + ); + }); + + it("keeps the existing unlimited behavior when no retained buffer limit is configured", () => { + const data = "x".repeat(1024 * 1024 + 1); + + expect(parseSseBody(`data: ${data}\n\n`)).toEqual([{ eventName: null, data }]); + }); });