diff --git a/.env.example b/.env.example index 0d6137a7c..9ee6a1b7d 100644 --- a/.env.example +++ b/.env.example @@ -171,6 +171,13 @@ FETCH_HEADERS_TIMEOUT=600000 FETCH_BODY_TIMEOUT=600000 MAX_RETRY_ATTEMPTS_DEFAULT=2 # 单供应商最大尝试次数(含首次调用),范围 1-10,留空使用默认值 2 +# 客户端断开后的 detached stream 共享带权进程级资源预算。 +# Replay owner 申请较重的 replay lease;预算不足时降级为轻量 metering, +# 两者都无法准入时才终止上游并按 499 结算。 +DETACHED_STREAM_MAX_CONCURRENCY=64 +DETACHED_STREAM_BUDGET_BYTES=67108864 +DETACHED_STREAM_METERING_RESERVE_BYTES=16777216 + # 入站压缩请求体(content-encoding: zstd/gzip/deflate/br)解压上限(字节) # 功能说明:/v1、/v1beta 代理路径不受 proxyClientMaxBodySize 钳制,这两项是入站解压的内存/CPU 兜底。 # - MAX_DECOMPRESSED_REQUEST_BYTES:解压输出上限,防御解压炸弹,超过按 413 拒绝。默认 100MB。 diff --git a/.vscode/settings.json b/.vscode/settings.json index 0a14f6b6e..e997d2405 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,3 +1,8 @@ { - "chatgpt.openOnStartup": true + "chatgpt.openOnStartup": true, + "i18n-ally.localesPaths": [ + "messages", + "src/i18n", + "src/app/[locale]/dashboard/sessions/[sessionId]/messages" + ] } diff --git a/CHANGELOG.md b/CHANGELOG.md index aa6aca9c7..7b0f69a67 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,12 @@ ### 修复 +- 修复上游响应流发生 error 后 Node/Undici body 未完成销毁的问题:Node-to-Web adapter 和 demand-driven pump + 现在在源流错误终态显式取消、销毁底层流,并为异步 destroy error 保留有界保护;同时将 raw body 的兜底错误监听改为一次性监听, + 避免 HTTP/2 reset、客户端断开和竞速取消路径长期保留 socket 与 ArrayBuffer backing store (#1430) +- 修复客户端断开后的后台计费 drain 继续累加完整响应正文导致的高并发内存放大:断线后切换到有界计量观察器, + 仅保留 usage、终止标记、模型与协议错误等结算证据,拿到终态即取消上游;Replay owner 在预算内继续保存完整 + 客户端可见流,Replay 预算不足时降级到 metering,新增共享进程级并发与带权保留容量预算,覆盖通用流与 Gemini 透传路径 (#1430) - 修复 Replay owner 在客户端断线后保留完整流正文和 300 秒传输资源导致的内存失控:限制 Redis write-behind backlog,Replay 失效后按断线起点恢复 60 秒 drain,并为 Redis session response body 增加默认 5 MiB 的可配置存储上限,避免大 SSE 正文及 before/after 快照放大内存和持久化压力; diff --git a/src/app/v1/_lib/proxy/client-abort-metering.test.ts b/src/app/v1/_lib/proxy/client-abort-metering.test.ts new file mode 100644 index 000000000..5367b490b --- /dev/null +++ b/src/app/v1/_lib/proxy/client-abort-metering.test.ts @@ -0,0 +1,281 @@ +import { describe, expect, it } from "vitest"; +import { + CLIENT_ABORT_METER_MAX_RETAINED_BYTES, + createClientAbortMeteringObserver, +} from "./client-abort-metering"; + +const encoder = new TextEncoder(); + +describe("createClientAbortMeteringObserver", () => { + it("keeps only compact Responses accounting evidence", () => { + const observer = createClientAbortMeteringObserver("response"); + observer.observe( + encoder.encode( + `event: response.output_text.delta\ndata: ${JSON.stringify({ + type: "response.output_text.delta", + delta: "x".repeat(32 * 1024), + })}\n\n` + ) + ); + const result = observer.observe( + encoder.encode( + `event: response.completed\ndata: ${JSON.stringify({ + type: "response.completed", + response: { + id: "resp_1", + model: "gpt-test", + output: [{ content: [{ text: "discard me" }] }], + usage: { input_tokens: 10, output_tokens: 5 }, + }, + })}\n\n` + ) + ); + + const snapshot = observer.finish(); + expect(result.billingComplete).toBe(true); + expect(snapshot.text).toContain("response.completed"); + expect(snapshot.text).toContain('"input_tokens":10'); + expect(snapshot.text).not.toContain("discard me"); + expect(snapshot.retainedBytes).toBeLessThanOrEqual(CLIENT_ABORT_METER_MAX_RETAINED_BYTES); + }); + + it("retains Claude initial and terminal usage until message_stop", () => { + const observer = createClientAbortMeteringObserver("claude"); + observer.observe( + encoder.encode( + `event: message_start\ndata: ${JSON.stringify({ + type: "message_start", + message: { model: "claude-test", usage: { input_tokens: 20, output_tokens: 1 } }, + })}\n\n` + ) + ); + expect( + observer.observe( + encoder.encode( + `event: message_delta\ndata: ${JSON.stringify({ + type: "message_delta", + usage: { output_tokens: 7 }, + })}\n\n` + ) + ).billingComplete + ).toBe(false); + expect( + observer.observe(encoder.encode(`event: message_stop\ndata: {"type":"message_stop"}\n\n`)) + .billingComplete + ).toBe(true); + + const snapshot = observer.finish(); + expect(snapshot.text).toContain("message_start"); + expect(snapshot.text).toContain("message_delta"); + expect(snapshot.text).toContain("message_stop"); + }); + + it("skips an oversized content frame and resumes at the next frame boundary", () => { + const observer = createClientAbortMeteringObserver("response"); + const result = observer.observe( + encoder.encode( + `event: response.output_text.delta\ndata: ${JSON.stringify({ + type: "response.output_text.delta", + delta: "x".repeat(70 * 1024), + })}\n\nevent: response.completed\ndata: ${JSON.stringify({ + type: "response.completed", + response: { usage: { input_tokens: 10, output_tokens: 5 } }, + })}\n\n` + ) + ); + + const snapshot = observer.finish(); + expect(result.billingComplete).toBe(true); + expect(snapshot.skippedOversizedFrames).toBe(1); + expect(snapshot.text).toContain("response.completed"); + }); + + it("requires terminal usage rather than a marker alone", () => { + const observer = createClientAbortMeteringObserver("openai"); + expect(observer.observe(encoder.encode("data: [DONE]\n\n")).billingComplete).toBe(false); + expect(observer.finish().billingComplete).toBe(false); + }); + + it("combines an OpenAI usage chunk with a later done marker across arbitrary splits", () => { + const observer = createClientAbortMeteringObserver("openai"); + const text = `data: ${JSON.stringify({ + id: "chatcmpl_1", + choices: [], + usage: { prompt_tokens: 12, completion_tokens: 4 }, + })}\r\n\r\ndata: [DONE]\r\n\r\n`; + const bytes = encoder.encode(text); + for (let offset = 0; offset < bytes.length; offset += 7) { + observer.observe(bytes.subarray(offset, offset + 7)); + } + + const snapshot = observer.finish(); + expect(snapshot.billingComplete).toBe(true); + expect(snapshot.text).toContain('"prompt_tokens":12'); + expect(snapshot.text).toContain("[DONE]"); + }); + + it("uses the last Gemini NDJSON usage and finishReason as terminal evidence", () => { + const observer = createClientAbortMeteringObserver("gemini"); + observer.observe( + encoder.encode( + `${JSON.stringify({ + candidates: [{ content: { parts: [{ text: "discard" }] } }], + usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 1 }, + })}\n` + ) + ); + const result = observer.observe( + encoder.encode( + `${JSON.stringify({ + candidates: [{ finishReason: "STOP" }], + usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 8 }, + })}\n` + ) + ); + + const snapshot = observer.finish(); + expect(result.billingComplete).toBe(true); + expect(snapshot.text).toContain('"candidatesTokenCount":8'); + expect(snapshot.text).not.toContain("discard"); + }); + + it("retains compact protocol errors without retaining content", () => { + const observer = createClientAbortMeteringObserver("response"); + observer.observe( + encoder.encode( + `event: error\ndata: ${JSON.stringify({ + type: "response.error", + error: { code: "upstream_failed", message: "failure" }, + debug: "x".repeat(32 * 1024), + })}\n\n` + ) + ); + + const snapshot = observer.finish(); + expect(snapshot.billingComplete).toBe(false); + expect(snapshot.text).toContain("upstream_failed"); + expect(snapshot.text).not.toContain('"debug"'); + }); + + it("compacts extended usage, metadata, cache, and signature evidence", () => { + const observer = createClientAbortMeteringObserver("response"); + observer.observe( + encoder.encode( + `event: response.in_progress\ndata: ${JSON.stringify({ + id: "resp_extended", + model: "gpt-extended", + prompt_cache_key: "cache-key", + service_tier: "priority", + status: "in_progress", + type: "response.in_progress", + message: { + id: "message-1", + model: "gpt-message", + usage: { input_tokens: 1 }, + }, + delta: { + type: "signature_delta", + stop_reason: "end_turn", + signature: "signed-model", + usage: { output_tokens: 2 }, + }, + usage: { + input_tokens: 10, + output_tokens: 3, + cache_creation_input_tokens: 2, + cache_creation_5m_input_tokens: 1, + cache_creation_1h_input_tokens: 1, + cache_read_input_tokens: 4, + input_tokens_details: { cached_tokens: 4, cache_write_tokens: 2 }, + prompt_tokens_details: { cached_tokens: 4, cache_write_tokens: 2 }, + cache_creation: { + ephemeral_5m_input_tokens: 1, + ephemeral_1h_input_tokens: 1, + }, + candidatesTokensDetails: [ + null, + {}, + { modality: "TEXT", tokenCount: 2 }, + { tokenCount: 1 }, + ], + promptTokensDetails: [{ modality: "IMAGE", tokenCount: 3 }], + }, + usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 3 }, + choices: [null, {}, { finish_reason: "stop" }], + candidates: [null, {}, { finishReason: "STOP" }], + ignored: "not retained", + })}\n\n` + ) + ); + observer.observe( + encoder.encode( + `event: response.completed\ndata: ${JSON.stringify({ + type: "response.completed", + response: { + id: "resp_extended", + model: "gpt-extended", + service_tier: "priority", + usage: { input_tokens: 10, output_tokens: 3 }, + }, + })}\n\n` + ) + ); + + const snapshot = observer.finish(); + expect(snapshot.billingComplete).toBe(true); + expect(snapshot.text).toContain('"prompt_cache_key":"cache-key"'); + expect(snapshot.text).toContain('"signature":"signed-model"'); + expect(snapshot.text).toContain('"cache_write_tokens":2'); + expect(snapshot.text).toContain('"modality":"IMAGE"'); + expect(snapshot.text).not.toContain("not retained"); + }); + + it("handles comments, multi-line data, bare JSON tails, and malformed frames", () => { + const observer = createClientAbortMeteringObserver("gemini-cli"); + observer.observe(new Uint8Array()); + observer.observe( + encoder.encode( + ': keepalive\rretry: 1000\revent: message\rdata: {"usageMetadata":\rdata: {"promptTokenCount":10,"candidatesTokenCount":2}}\r\r' + ) + ); + observer.observe(encoder.encode("data: true\n\n")); + observer.observe(encoder.encode("data: not-json\n\n")); + observer.observe(encoder.encode("data: still-not-json\n\n")); + observer.observe( + encoder.encode( + JSON.stringify({ + candidates: [{ finishReason: "STOP" }], + usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 2 }, + }) + ) + ); + + const snapshot = observer.finish(); + expect(snapshot.billingComplete).toBe(true); + expect(snapshot.protocolFailure).toEqual({ + afterContent: false, + verdict: "malformed", + eventName: null, + }); + }); + + it("recovers after an oversized bare JSON line and ignores post-finish input", () => { + const observer = createClientAbortMeteringObserver("gemini"); + observer.observe(encoder.encode(`{"ignored":"${"x".repeat(70 * 1024)}"}\n`)); + observer.observe( + encoder.encode( + `${JSON.stringify({ + candidates: [{ finishReason: "STOP" }], + usageMetadata: { promptTokenCount: 4, candidatesTokenCount: 2 }, + })}` + ) + ); + const first = observer.finish(); + observer.observe(encoder.encode('{"error":true}\n')); + const second = observer.finish(); + + expect(first.billingComplete).toBe(true); + expect(first.skippedOversizedFrames).toBe(1); + expect(second).toEqual(first); + }); +}); diff --git a/src/app/v1/_lib/proxy/client-abort-metering.ts b/src/app/v1/_lib/proxy/client-abort-metering.ts new file mode 100644 index 000000000..ea62e399b --- /dev/null +++ b/src/app/v1/_lib/proxy/client-abort-metering.ts @@ -0,0 +1,517 @@ +import type { ClientFormat } from "./format-mapper"; + +export const CLIENT_ABORT_METER_MAX_RETAINED_BYTES = 64 * 1024; +export const CLIENT_ABORT_METER_MAX_FRAME_BYTES = 64 * 1024; + +type EvidenceSlot = + | "error" + | "initial-usage" + | "latest-usage" + | "metadata" + | "signature" + | "terminal"; + +export interface ClientAbortMeteringSnapshot { + text: string; + billingComplete: boolean; + retainedBytes: number; + skippedOversizedFrames: number; + protocolFailure: { + afterContent: boolean; + verdict: "error" | "malformed"; + eventName: string | null; + } | null; +} + +export interface ClientAbortMeteringObserver { + observe(chunk: Uint8Array): { billingComplete: boolean }; + finish(): ClientAbortMeteringSnapshot; +} + +interface ParsedFrame { + eventName: string | null; + data: string; +} + +const USAGE_NUMBER_FIELDS = [ + "cachedContentTokenCount", + "cache_creation_1h_input_tokens", + "cache_creation_5m_input_tokens", + "cache_creation_input_tokens", + "cache_read_input_tokens", + "candidatesTokenCount", + "claude_cache_creation_1_h_tokens", + "claude_cache_creation_5_m_tokens", + "completion_tokens", + "input_tokens", + "output_tokens", + "promptTokenCount", + "prompt_tokens", + "thoughtsTokenCount", +] as const; + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function copyFiniteNumberFields( + source: Record, + target: Record, + fields: readonly string[] +): void { + for (const field of fields) { + const value = source[field]; + if (typeof value === "number" && Number.isFinite(value)) target[field] = value; + } +} + +function compactTokenDetails(value: unknown): unknown { + if (!Array.isArray(value)) return undefined; + const details = value.slice(0, 16).flatMap((entry) => { + if (!isRecord(entry)) return []; + const compact: Record = {}; + if (typeof entry.modality === "string") compact.modality = entry.modality.slice(0, 32); + if (typeof entry.tokenCount === "number" && Number.isFinite(entry.tokenCount)) { + compact.tokenCount = entry.tokenCount; + } + return Object.keys(compact).length > 0 ? [compact] : []; + }); + return details.length > 0 ? details : undefined; +} + +function compactUsage(value: unknown): Record | null { + if (!isRecord(value)) return null; + const compact: Record = {}; + copyFiniteNumberFields(value, compact, USAGE_NUMBER_FIELDS); + + for (const field of ["input_tokens_details", "prompt_tokens_details"] as const) { + const details = value[field]; + if (!isRecord(details)) continue; + const compactDetails: Record = {}; + copyFiniteNumberFields(details, compactDetails, ["cached_tokens", "cache_write_tokens"]); + if (Object.keys(compactDetails).length > 0) compact[field] = compactDetails; + } + + if (isRecord(value.cache_creation)) { + const cacheCreation: Record = {}; + copyFiniteNumberFields(value.cache_creation, cacheCreation, [ + "ephemeral_1h_input_tokens", + "ephemeral_5m_input_tokens", + ]); + if (Object.keys(cacheCreation).length > 0) compact.cache_creation = cacheCreation; + } + + for (const field of ["candidatesTokensDetails", "promptTokensDetails"] as const) { + const details = compactTokenDetails(value[field]); + if (details) compact[field] = details; + } + + return Object.keys(compact).length > 0 ? compact : null; +} + +function compactError(value: unknown): unknown { + if (typeof value === "string") return value.slice(0, 1024); + if (!isRecord(value)) return value === true ? true : undefined; + const compact: Record = {}; + for (const field of ["code", "message", "type"] as const) { + const fieldValue = value[field]; + if (typeof fieldValue === "string") compact[field] = fieldValue.slice(0, 1024); + } + return Object.keys(compact).length > 0 ? compact : true; +} + +function compactPayload(value: Record): Record { + const compact: Record = {}; + for (const field of [ + "id", + "model", + "prompt_cache_key", + "service_tier", + "status", + "type", + ] as const) { + const fieldValue = value[field]; + if (typeof fieldValue === "string") compact[field] = fieldValue.slice(0, 256); + } + if (value.failed === true) compact.failed = true; + if (value.error !== undefined) compact.error = compactError(value.error); + + for (const field of ["usage", "usageMetadata"] as const) { + const usage = compactUsage(value[field]); + if (usage) compact[field] = usage; + } + + if (isRecord(value.message)) { + const message: Record = {}; + for (const field of ["id", "model"] as const) { + const fieldValue = value.message[field]; + if (typeof fieldValue === "string") message[field] = fieldValue.slice(0, 256); + } + const usage = compactUsage(value.message.usage); + if (usage) message.usage = usage; + if (Object.keys(message).length > 0) compact.message = message; + } + + if (isRecord(value.delta)) { + const delta: Record = {}; + if (typeof value.delta.type === "string") delta.type = value.delta.type.slice(0, 128); + if (typeof value.delta.stop_reason === "string") { + delta.stop_reason = value.delta.stop_reason.slice(0, 128); + } + if (typeof value.delta.signature === "string") { + delta.signature = value.delta.signature.slice(0, 8192); + } + const usage = compactUsage(value.delta.usage); + if (usage) delta.usage = usage; + if (Object.keys(delta).length > 0) compact.delta = delta; + } + + if (isRecord(value.response)) compact.response = compactPayload(value.response); + + if (Array.isArray(value.choices)) { + compact.choices = value.choices.slice(0, 16).map((choice) => { + if (!isRecord(choice) || typeof choice.finish_reason !== "string") return {}; + return { finish_reason: choice.finish_reason.slice(0, 128) }; + }); + } + + if (Array.isArray(value.candidates)) { + compact.candidates = value.candidates.slice(0, 16).map((candidate) => { + if (!isRecord(candidate) || typeof candidate.finishReason !== "string") return {}; + return { finishReason: candidate.finishReason.slice(0, 128) }; + }); + } + + return compact; +} + +function positiveUsage(value: unknown): boolean { + const usage = compactUsage(value); + if (!usage) return false; + const stack: unknown[] = [usage]; + while (stack.length > 0) { + const current = stack.pop(); + if (typeof current === "number" && current > 0) return true; + if (Array.isArray(current)) stack.push(...current); + else if (isRecord(current)) stack.push(...Object.values(current)); + } + return false; +} + +function findUsage(value: Record): boolean { + if (positiveUsage(value.usage) || positiveUsage(value.usageMetadata)) return true; + if (isRecord(value.message) && positiveUsage(value.message.usage)) return true; + if (isRecord(value.delta) && positiveUsage(value.delta.usage)) return true; + return isRecord(value.response) && findUsage(value.response); +} + +function isProtocolError(value: Record): boolean { + return ( + value.error !== undefined || + value.failed === true || + value.type === "error" || + value.type === "response.error" || + value.type === "response.failed" || + (isRecord(value.response) && value.response.error !== undefined) + ); +} + +function hasOpenAiCompletion(value: Record): boolean { + return ( + Array.isArray(value.choices) && + value.choices.some( + (choice) => + isRecord(choice) && + typeof choice.finish_reason === "string" && + choice.finish_reason.trim().length > 0 + ) + ); +} + +function hasGeminiCompletion(value: Record): boolean { + const payload = isRecord(value.response) ? value.response : value; + return ( + Array.isArray(payload.candidates) && + payload.candidates.some( + (candidate) => + isRecord(candidate) && + typeof candidate.finishReason === "string" && + candidate.finishReason.trim().length > 0 + ) + ); +} + +class BoundedEventFramer { + private readonly decoder = new TextDecoder("utf-8"); + private line = ""; + private lineOverflow = false; + private overflowedRawJsonLine = false; + private pendingCr = false; + private eventName: string | null = null; + private dataLines: string[] = []; + private frameCharacters = 0; + private droppingFrame = false; + skippedOversizedFrames = 0; + + constructor( + private readonly maxFrameCharacters: number, + private readonly onFrame: (frame: ParsedFrame) => void + ) {} + + push(chunk: Uint8Array): void { + if (chunk.byteLength === 0) return; + this.consume(this.decoder.decode(chunk, { stream: true })); + } + + finish(): void { + this.consume(this.decoder.decode()); + if (this.line.length > 0 || this.lineOverflow) this.consumeLine(); + this.flushFrame(); + } + + private consume(text: string): void { + for (const character of text) { + if (this.pendingCr) { + this.pendingCr = false; + if (character === "\n") continue; + } + if (character === "\r") { + this.consumeLine(); + this.pendingCr = true; + continue; + } + if (character === "\n") { + this.consumeLine(); + continue; + } + if (this.lineOverflow) continue; + if (this.line.length >= this.maxFrameCharacters) { + this.overflowedRawJsonLine = + this.eventName === null && + this.dataLines.length === 0 && + this.line.trimStart().startsWith("{"); + this.line = ""; + this.lineOverflow = true; + this.dropCurrentFrame(); + continue; + } + this.line += character; + } + } + + private consumeLine(): void { + const line = this.line; + const overflowed = this.lineOverflow; + const overflowedRawJsonLine = this.overflowedRawJsonLine; + this.line = ""; + this.lineOverflow = false; + this.overflowedRawJsonLine = false; + + if (overflowed && overflowedRawJsonLine) { + this.droppingFrame = false; + this.resetFrame(); + return; + } + + if (line.length === 0 && !overflowed) { + if (this.droppingFrame) { + this.droppingFrame = false; + this.resetFrame(); + } else { + this.flushFrame(); + } + return; + } + if (this.droppingFrame || overflowed) return; + if (line.startsWith(":")) return; + if (line.startsWith("event:")) { + this.eventName = line.slice(6).trim().slice(0, 256); + this.frameCharacters += line.length; + this.enforceFrameLimit(); + return; + } + if (line.startsWith("data:")) { + const data = line.slice(5).replace(/^\s/, ""); + this.dataLines.push(data); + this.frameCharacters += data.length; + this.enforceFrameLimit(); + return; + } + + const candidate = line.trim(); + if (this.eventName === null && this.dataLines.length === 0 && candidate.startsWith("{")) { + if (candidate.length <= this.maxFrameCharacters) { + this.onFrame({ eventName: null, data: candidate }); + } else { + this.skippedOversizedFrames += 1; + } + } + } + + private enforceFrameLimit(): void { + if (this.frameCharacters <= this.maxFrameCharacters) return; + this.dropCurrentFrame(); + } + + private dropCurrentFrame(): void { + if (!this.droppingFrame) this.skippedOversizedFrames += 1; + this.droppingFrame = true; + this.resetFrame(); + } + + private flushFrame(): void { + if (this.droppingFrame || this.dataLines.length === 0) { + this.resetFrame(); + return; + } + this.onFrame({ eventName: this.eventName, data: this.dataLines.join("\n") }); + this.resetFrame(); + } + + private resetFrame(): void { + this.eventName = null; + this.dataLines = []; + this.frameCharacters = 0; + } +} + +export function createClientAbortMeteringObserver( + format: ClientFormat +): ClientAbortMeteringObserver { + const evidence = new Map(); + const encoder = new TextEncoder(); + let terminalSeen = false; + let terminalUsageSeen = false; + let protocolFailure: ClientAbortMeteringSnapshot["protocolFailure"] = null; + let finished = false; + + const retainedBytes = () => + [...new Set(evidence.values())].reduce( + (total, value) => total + encoder.encode(value).length, + 0 + ); + + const setEvidence = (slot: EvidenceSlot, value: string): void => { + const previous = evidence.get(slot); + evidence.set(slot, value); + if (retainedBytes() <= CLIENT_ABORT_METER_MAX_RETAINED_BYTES) return; + evidence.delete(slot); + if (previous !== undefined) evidence.set(slot, previous); + }; + + const recordFrame = (frame: ParsedFrame): void => { + const trimmed = frame.data.trim(); + if (trimmed === "[DONE]") { + if (format === "openai") terminalSeen = true; + setEvidence("terminal", "data: [DONE]\n\n"); + return; + } + + let parsed: unknown; + try { + parsed = JSON.parse(trimmed) as unknown; + } catch { + protocolFailure ??= { + afterContent: terminalSeen, + verdict: "malformed", + eventName: frame.eventName, + }; + return; + } + if (!isRecord(parsed)) return; + + const hasUsage = findUsage(parsed); + const protocolError = isProtocolError(parsed) || frame.eventName === "error"; + const type = typeof parsed.type === "string" ? parsed.type : null; + const terminal = (() => { + switch (format) { + case "response": + return ( + (type === "response.completed" || type === "response.done") && + (frame.eventName === null || frame.eventName === "message" || frame.eventName === type) + ); + case "claude": + return ( + type === "message_stop" && + (frame.eventName === null || + frame.eventName === "message" || + frame.eventName === "message_stop") + ); + case "openai": + return hasOpenAiCompletion(parsed); + case "gemini": + case "gemini-cli": + return hasGeminiCompletion(parsed); + } + })(); + + if (terminal) terminalSeen = true; + if ( + hasUsage && + (format !== "claude" || type === "message_delta" || frame.eventName === "message_delta") + ) { + terminalUsageSeen = true; + } + + const compact = compactPayload(parsed); + const compactData = JSON.stringify(compact); + const normalized = `${frame.eventName ? `event: ${frame.eventName}\n` : ""}data: ${compactData}\n\n`; + + if (protocolError) { + protocolFailure ??= { + afterContent: terminalSeen, + verdict: "error", + eventName: frame.eventName, + }; + setEvidence("error", normalized); + } + if (terminal) setEvidence("terminal", normalized); + if (hasUsage) { + const isInitialClaudeUsage = + format === "claude" && (type === "message_start" || frame.eventName === "message_start"); + setEvidence(isInitialClaudeUsage ? "initial-usage" : "latest-usage", normalized); + } + if ( + isRecord(parsed.delta) && + parsed.delta.type === "signature_delta" && + typeof parsed.delta.signature === "string" + ) { + setEvidence("signature", normalized); + } + if ( + typeof parsed.model === "string" || + typeof parsed.prompt_cache_key === "string" || + typeof parsed.service_tier === "string" || + (isRecord(parsed.message) && typeof parsed.message.model === "string") || + (isRecord(parsed.response) && + (typeof parsed.response.model === "string" || + typeof parsed.response.service_tier === "string")) + ) { + setEvidence("metadata", normalized); + } + }; + + const framer = new BoundedEventFramer(CLIENT_ABORT_METER_MAX_FRAME_BYTES, recordFrame); + const isBillingComplete = () => terminalSeen && terminalUsageSeen; + + return { + observe(chunk): { billingComplete: boolean } { + if (!finished) framer.push(chunk); + return { billingComplete: isBillingComplete() }; + }, + finish(): ClientAbortMeteringSnapshot { + if (!finished) { + finished = true; + framer.finish(); + } + const text = [...new Set(evidence.values())].join(""); + return { + text, + billingComplete: isBillingComplete(), + retainedBytes: encoder.encode(text).length, + skippedOversizedFrames: framer.skippedOversizedFrames, + protocolFailure: protocolFailure ? { ...protocolFailure } : null, + }; + }, + }; +} diff --git a/src/app/v1/_lib/proxy/demand-driven-response-pump.test.ts b/src/app/v1/_lib/proxy/demand-driven-response-pump.test.ts index 965b35d62..5f0655462 100644 --- a/src/app/v1/_lib/proxy/demand-driven-response-pump.test.ts +++ b/src/app/v1/_lib/proxy/demand-driven-response-pump.test.ts @@ -11,6 +11,14 @@ function nextTurn(): Promise { return new Promise((resolve) => setImmediate(resolve)); } +function createDeferred() { + let resolve!: (value: T | PromiseLike) => void; + const promise = new Promise((promiseResolve) => { + resolve = promiseResolve; + }); + return { promise, resolve }; +} + function trackReaderRelease(source: ReadableStream) { const reader = source.getReader(); const releaseLock = vi.spyOn(reader, "releaseLock"); @@ -403,4 +411,62 @@ describe("createDemandDrivenResponsePump", () => { }); expect(releaseLock).toHaveBeenCalledTimes(1); }); + + it("cancels the source when the primed read fails", async () => { + const sourceError = new Error("source-read-failed"); + const cancel = vi.fn(() => Promise.resolve()); + const releaseLock = vi.fn(); + const reader = { + read: vi.fn(() => Promise.reject(sourceError)), + cancel, + releaseLock, + } as unknown as ReadableStreamDefaultReader; + const source = { + getReader: vi.fn(() => reader), + } as unknown as ReadableStream; + + const pump = createDemandDrivenResponsePump({ source, onChunk: vi.fn() }); + + await expect(pump.completion).resolves.toMatchObject({ + streamEndedNormally: false, + error: sourceError, + }); + expect(cancel).toHaveBeenCalledWith(sourceError); + expect(releaseLock).toHaveBeenCalledOnce(); + }); + + it("finishes a metering drain without reporting a source error", async () => { + const cancelGate = createDeferred(); + const cancel = vi.fn(() => cancelGate.promise); + const releaseLock = vi.fn(); + const reader = { + read: vi.fn(() => new Promise>(() => {})), + cancel, + releaseLock, + } as unknown as ReadableStreamDefaultReader; + const source = { + getReader: vi.fn(() => reader), + } as unknown as ReadableStream; + const pump = createDemandDrivenResponsePump({ source, onChunk: vi.fn() }); + let teardownSettled = false; + void pump.teardown.then(() => { + teardownSettled = true; + }); + + pump.startDrain("client detached"); + pump.finishDrain("terminal usage captured"); + + await expect(pump.completion).resolves.toEqual({ + streamEndedNormally: false, + clientAborted: true, + error: null, + }); + expect(cancel).toHaveBeenCalledOnce(); + expect(teardownSettled).toBe(false); + + cancelGate.resolve(); + await pump.teardown; + expect(teardownSettled).toBe(true); + expect(releaseLock).toHaveBeenCalledOnce(); + }); }); diff --git a/src/app/v1/_lib/proxy/demand-driven-response-pump.ts b/src/app/v1/_lib/proxy/demand-driven-response-pump.ts index 137db3dd9..1f2a96ba7 100644 --- a/src/app/v1/_lib/proxy/demand-driven-response-pump.ts +++ b/src/app/v1/_lib/proxy/demand-driven-response-pump.ts @@ -16,7 +16,9 @@ export interface DemandDrivenResponsePumpOptions { export interface DemandDrivenResponsePump { stream: ReadableStream; completion: Promise; + teardown: Promise; startDrain: (reason?: unknown) => void; + finishDrain: (reason?: unknown) => void; cancelSource: (reason?: unknown) => void; errorClient: (error: Error) => void; getState: () => DemandDrivenResponsePumpState; @@ -43,9 +45,13 @@ export function createDemandDrivenResponsePump( let readerReleased = false; let pendingChunkDeadlineId: ReturnType | null = null; let resolveCompletion: (completion: DemandDrivenResponsePumpCompletion) => void = () => {}; + let resolveTeardown = () => {}; const completion = new Promise((resolve) => { resolveCompletion = resolve; }); + const teardown = new Promise((resolve) => { + resolveTeardown = resolve; + }); const releaseReader = () => { if (readerReleased) return; @@ -91,7 +97,11 @@ export function createDemandDrivenResponsePump( clientController = null; state = "closed"; resolveCompletion({ streamEndedNormally, clientAborted, error }); - void cancelPromise?.then(undefined, recordSourceCancelFailure); + if (cancelPromise) { + void cancelPromise.then(undefined, recordSourceCancelFailure).finally(resolveTeardown); + } else { + resolveTeardown(); + } }; const finishWithError = (error: unknown) => { @@ -104,7 +114,10 @@ export function createDemandDrivenResponsePump( // The downstream may have cancelled concurrently. } } - settle(false, normalized); + // A rejected source read bypasses the Web stream cancel algorithm. Keep + // source ownership explicit so adapters can release the underlying Node + // stream, socket, and native backing store on every terminal error. + settle(false, normalized, normalized); }; const finishNormally = () => { @@ -146,6 +159,12 @@ export function createDemandDrivenResponsePump( settle(false, normalized, normalized); }; + const finishDrain = (reason?: unknown) => { + if (settled || state !== "draining") return; + const normalized = reason == null ? new Error("Background drain complete") : toError(reason); + settle(false, null, normalized); + }; + const armPendingChunkDeadline = () => { clearPendingChunkDeadline(); pendingChunkDeadlineId = setTimeout(() => { @@ -264,7 +283,9 @@ export function createDemandDrivenResponsePump( return { stream, completion, + teardown, startDrain, + finishDrain, cancelSource, errorClient(error) { if (settled || state !== "client-active") return; diff --git a/src/app/v1/_lib/proxy/detached-stream-budget.test.ts b/src/app/v1/_lib/proxy/detached-stream-budget.test.ts new file mode 100644 index 000000000..cfe38c290 --- /dev/null +++ b/src/app/v1/_lib/proxy/detached-stream-budget.test.ts @@ -0,0 +1,73 @@ +import { describe, expect, it } from "vitest"; +import { DetachedStreamBudget } from "./detached-stream-budget"; + +function createBudget( + overrides: Partial["limits"]> = {} +) { + return new DetachedStreamBudget(() => ({ + maxConcurrency: 4, + maxReservedBytes: 1024, + meteringReserveBytes: 256, + ...overrides, + })); +} + +describe("DetachedStreamBudget", () => { + it("enforces concurrency and releases weighted leases idempotently", () => { + const budget = createBudget({ maxConcurrency: 1 }); + const first = budget.tryAcquire("metering", 256); + expect(first.acquired).toBe(true); + expect(budget.tryAcquire("metering", 256)).toEqual({ + acquired: false, + reason: "concurrency_exhausted", + }); + + if (!first.acquired) throw new Error("expected first lease"); + first.lease.release(); + first.lease.release(); + expect(budget.snapshot()).toMatchObject({ + activeStreams: 0, + reservedBytes: 0, + activeByKind: { metering: 0, replay: 0 }, + reservedByKind: { metering: 0, replay: 0 }, + }); + }); + + it("reserves headroom for metering when admitting Replay owners", () => { + const budget = createBudget(); + const replay = budget.tryAcquire("replay", 768); + expect(replay.acquired).toBe(true); + expect(budget.tryAcquire("replay", 1)).toEqual({ + acquired: false, + reason: "metering_reserve", + }); + expect(budget.tryAcquire("metering", 256).acquired).toBe(true); + }); + + it("enforces the aggregate memory budget across lease kinds", () => { + const budget = createBudget({ meteringReserveBytes: 0 }); + expect(budget.tryAcquire("replay", 768).acquired).toBe(true); + expect(budget.tryAcquire("metering", 257)).toEqual({ + acquired: false, + reason: "memory_budget_exhausted", + }); + }); + + it("tracks Replay and metering reservations independently", () => { + const budget = createBudget({ meteringReserveBytes: 0 }); + expect(budget.tryAcquire("replay", 512).acquired).toBe(true); + expect(budget.tryAcquire("metering", 256).acquired).toBe(true); + expect(budget.snapshot()).toMatchObject({ + activeStreams: 2, + reservedBytes: 768, + activeByKind: { metering: 1, replay: 1 }, + reservedByKind: { metering: 256, replay: 512 }, + }); + }); + + it("rejects invalid reservations without mutating state", () => { + const budget = createBudget(); + expect(() => budget.tryAcquire("metering", 0)).toThrow(RangeError); + expect(budget.snapshot().activeStreams).toBe(0); + }); +}); diff --git a/src/app/v1/_lib/proxy/detached-stream-budget.ts b/src/app/v1/_lib/proxy/detached-stream-budget.ts new file mode 100644 index 000000000..942100fde --- /dev/null +++ b/src/app/v1/_lib/proxy/detached-stream-budget.ts @@ -0,0 +1,130 @@ +import { getEnvConfig } from "@/lib/config/env.schema"; + +export type DetachedStreamLeaseKind = "metering" | "replay"; + +export interface DetachedStreamBudgetLimits { + maxConcurrency: number; + maxReservedBytes: number; + meteringReserveBytes: number; +} + +export interface DetachedStreamBudgetSnapshot { + activeStreams: number; + reservedBytes: number; + activeByKind: Record; + reservedByKind: Record; + limits: DetachedStreamBudgetLimits; +} + +export interface DetachedStreamLease { + readonly kind: DetachedStreamLeaseKind; + readonly reservedBytes: number; + release(): void; +} + +export type DetachedStreamAcquireResult = + | { acquired: true; lease: DetachedStreamLease } + | { + acquired: false; + reason: "concurrency_exhausted" | "memory_budget_exhausted" | "metering_reserve"; + }; + +function createKindCounters(): Record { + return { metering: 0, replay: 0 }; +} + +export class DetachedStreamBudget { + private activeStreams = 0; + private reservedBytes = 0; + private readonly activeByKind = createKindCounters(); + private readonly reservedByKind = createKindCounters(); + + constructor(private readonly resolveLimits: () => DetachedStreamBudgetLimits) {} + + tryAcquire(kind: DetachedStreamLeaseKind, reservedBytes: number): DetachedStreamAcquireResult { + if (!Number.isSafeInteger(reservedBytes) || reservedBytes <= 0) { + throw new RangeError("Detached stream reservation must be a positive safe integer"); + } + + const limits = this.resolveLimits(); + if (this.activeStreams >= limits.maxConcurrency) { + return { acquired: false, reason: "concurrency_exhausted" }; + } + + const nextReservedBytes = this.reservedBytes + reservedBytes; + if (nextReservedBytes > limits.maxReservedBytes) { + return { acquired: false, reason: "memory_budget_exhausted" }; + } + + const effectiveMeteringReserve = Math.min( + limits.maxReservedBytes, + Math.max(0, limits.meteringReserveBytes) + ); + if ( + kind === "replay" && + nextReservedBytes > limits.maxReservedBytes - effectiveMeteringReserve + ) { + return { acquired: false, reason: "metering_reserve" }; + } + + this.activeStreams += 1; + this.reservedBytes = nextReservedBytes; + this.activeByKind[kind] += 1; + this.reservedByKind[kind] += reservedBytes; + let released = false; + + return { + acquired: true, + lease: { + kind, + reservedBytes, + release: () => { + if (released) return; + released = true; + this.activeStreams = Math.max(0, this.activeStreams - 1); + this.reservedBytes = Math.max(0, this.reservedBytes - reservedBytes); + this.activeByKind[kind] = Math.max(0, this.activeByKind[kind] - 1); + this.reservedByKind[kind] = Math.max(0, this.reservedByKind[kind] - reservedBytes); + }, + }, + }; + } + + snapshot(): DetachedStreamBudgetSnapshot { + return { + activeStreams: this.activeStreams, + reservedBytes: this.reservedBytes, + activeByKind: { ...this.activeByKind }, + reservedByKind: { ...this.reservedByKind }, + limits: { ...this.resolveLimits() }, + }; + } +} + +const DETACHED_STREAM_BUDGET_SYMBOL = Symbol.for("cch.detachedStreamBudget"); + +function getDetachedStreamBudget(): DetachedStreamBudget { + const globalState = globalThis as typeof globalThis & { + [DETACHED_STREAM_BUDGET_SYMBOL]?: DetachedStreamBudget; + }; + globalState[DETACHED_STREAM_BUDGET_SYMBOL] ??= new DetachedStreamBudget(() => { + const env = getEnvConfig(); + return { + maxConcurrency: env.DETACHED_STREAM_MAX_CONCURRENCY ?? 64, + maxReservedBytes: env.DETACHED_STREAM_BUDGET_BYTES ?? 64 * 1024 * 1024, + meteringReserveBytes: env.DETACHED_STREAM_METERING_RESERVE_BYTES ?? 16 * 1024 * 1024, + }; + }); + return globalState[DETACHED_STREAM_BUDGET_SYMBOL]; +} + +export function acquireDetachedStreamLease( + kind: DetachedStreamLeaseKind, + reservedBytes: number +): DetachedStreamAcquireResult { + return getDetachedStreamBudget().tryAcquire(kind, reservedBytes); +} + +export function getDetachedStreamBudgetSnapshot(): DetachedStreamBudgetSnapshot { + return getDetachedStreamBudget().snapshot(); +} diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 84bc26942..04069b8c3 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -8290,7 +8290,7 @@ export class ProxyForwarder { // ⭐ 立即为 undici body 添加错误处理,防止 uncaughtException // 必须在任何其他操作之前设置,否则 ECONNRESET 等错误会导致 uncaughtException const rawBody = undiciRes.body as Readable; - rawBody.on("error", (err) => { + rawBody.once("error", (err) => { const code = (err as NodeJS.ErrnoException).code; // 客户端/上游断连是高频路径事件,降级为 debug 以减少噪音 // 集合需与下方 streamPipeline 回调中的 isExpectedDisconnect 保持一致 @@ -8307,6 +8307,17 @@ export class ProxyForwarder { error: err.message, errorCode: code, }); + // Undici normally auto-destroys BodyReadable instances, but custom + // dispatchers and HTTP/2 abort races can surface an error before that + // teardown reaches the wrapper. Keep this cleanup idempotent so the + // raw body cannot remain paused with its socket/backing store retained. + if (!rawBody.destroyed) { + try { + rawBody.destroy(err); + } catch { + // ignore + } + } }); // 构建响应头 diff --git a/src/app/v1/_lib/proxy/node-stream-to-web.test.ts b/src/app/v1/_lib/proxy/node-stream-to-web.test.ts index 6ef589232..da5e3b9ef 100644 --- a/src/app/v1/_lib/proxy/node-stream-to-web.test.ts +++ b/src/app/v1/_lib/proxy/node-stream-to-web.test.ts @@ -188,6 +188,7 @@ describe("nodeStreamToWebStreamSafe", () => { queueMicrotask(() => node.emit("error", boom)); await expect(reader.read()).rejects.toThrow("boom"); + await new Promise((resolve) => setImmediate(resolve)); // After error settles, listeners must be detached expect(node.listenerCount("data")).toBe(0); expect(node.listenerCount("end")).toBe(0); @@ -195,6 +196,29 @@ describe("nodeStreamToWebStreamSafe", () => { expect(node.listenerCount("error")).toBe(0); }); + it("destroys the underlying source when it emits an error", async () => { + const node = new Readable({ + read() { + // Keep the source open until the explicit error below. + }, + }); + const web = nodeStreamToWebStreamSafe(node, 1, "test"); + const reader = web.getReader(); + const uncaughtSpy = vi.fn(); + process.once("uncaughtException", uncaughtSpy); + + const boom = new Error("destroy-after-error"); + const pendingRead = reader.read(); + node.emit("error", boom); + + await expect(pendingRead).rejects.toBe(boom); + await new Promise((resolve) => setImmediate(resolve)); + + process.removeListener("uncaughtException", uncaughtSpy); + expect(node.destroyed).toBe(true); + expect(uncaughtSpy).not.toHaveBeenCalled(); + }); + it("rejects when the source closes after conversion without reaching EOF", async () => { const node = new Readable({ read() { diff --git a/src/app/v1/_lib/proxy/node-stream-to-web.ts b/src/app/v1/_lib/proxy/node-stream-to-web.ts index 82b3bbff8..d30ebeb83 100644 --- a/src/app/v1/_lib/proxy/node-stream-to-web.ts +++ b/src/app/v1/_lib/proxy/node-stream-to-web.ts @@ -120,6 +120,17 @@ export function nodeStreamToWebStreamSafe( errorName: err.name, }); detach(nodeStream); + if (!nodeStream.destroyed) { + // A source error settles the Web stream but does not invoke its + // cancel algorithm. Destroy the Node/Undici body explicitly so a + // paused transport cannot retain its socket and backing store. + installPendingDestroyErrorGuard(nodeStream); + try { + nodeStream.destroy(err); + } catch { + // ignore + } + } try { controller.error(err); } catch { diff --git a/src/app/v1/_lib/proxy/replay/replay-spool.ts b/src/app/v1/_lib/proxy/replay/replay-spool.ts index 4582a06d4..12a54fa64 100644 --- a/src/app/v1/_lib/proxy/replay/replay-spool.ts +++ b/src/app/v1/_lib/proxy/replay/replay-spool.ts @@ -45,6 +45,7 @@ function serializeDurablePersistence(operation: () => Promise): Promise export interface ReplaySpoolOptions { onInactive?: () => void; + onTerminal?: () => void; } export function getActiveReplaySpoolCount(): number { @@ -453,6 +454,13 @@ export class ReplaySpool { this.released = true; this.clearOwnerHeartbeat(); activeSpoolCount = Math.max(0, activeSpoolCount - 1); + try { + this.options.onTerminal?.(); + } catch (error) { + logger.debug("[ReplaySpool] terminal callback failed", { + error: error instanceof Error ? error.message : String(error), + }); + } } private clearFlushTimer(): void { diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index d06c75c65..973f7f7b5 100644 --- a/src/app/v1/_lib/proxy/response-handler.ts +++ b/src/app/v1/_lib/proxy/response-handler.ts @@ -55,10 +55,19 @@ import type { GeminiResponse } from "../gemini/types"; import { extractActualResponseModelForProvider, extractJsonChunks } from "./actual-response-model"; import { recordAffinityWinner, tombstoneAffinityOnFailure } from "./affinity/affinity-recorder"; import { bindClientAbortListener } from "./client-abort-listener"; +import { + CLIENT_ABORT_METER_MAX_RETAINED_BYTES, + createClientAbortMeteringObserver, +} from "./client-abort-metering"; import { createDemandDrivenResponsePump, type DemandDrivenResponsePump, } from "./demand-driven-response-pump"; +import { + acquireDetachedStreamLease, + type DetachedStreamLease, + getDetachedStreamBudgetSnapshot, +} from "./detached-stream-budget"; import { isDiscoveryProtocolErrorPayload } from "./discovery-validity"; import { isClientAbortError, isTransportError } from "./errors"; import { @@ -82,6 +91,10 @@ import { } from "./stream-gate/stream-protocol-observer"; const CLIENT_ABORT_DRAIN_MAX_MS = 60_000; +const CLIENT_ABORT_DRAIN_RESERVATION_BYTES = + 3 * 1024 * 1024 + CLIENT_ABORT_METER_MAX_RETAINED_BYTES; +const REPLAY_DRAIN_FIXED_OVERHEAD_BYTES = 5 * 1024 * 1024; +const GEMINI_STREAM_TRANSFORM_MAX_BUFFER_CHARACTERS = 1024 * 1024; const STREAM_STATS_MAX_BUFFER_BYTES = 10 * 1024 * 1024; const STREAM_STATS_HEAD_BYTES = 1024 * 1024; const STREAM_STATS_TAIL_BYTES = STREAM_STATS_MAX_BUFFER_BYTES - STREAM_STATS_HEAD_BYTES; @@ -93,6 +106,14 @@ function getSessionRequestOwnerKeyId(session: ProxySession): number | undefined return session.authState?.key?.id ?? session.messageContext?.key?.id ?? undefined; } +function resolveReplayDrainReservationBytes(): number { + const payloadBytes = getEnvConfig().REPLAY_MAX_PAYLOAD_BYTES; + // ReplaySpool keeps bounded write-back state, then terminal persistence reads + // the Redis chunks and joins one payload string. Reserve the payload three + // times for chunk strings, the joined string, and UTF-16 expansion. + return REPLAY_DRAIN_FIXED_OVERHEAD_BYTES + payloadBytes * 3; +} + type BoundedStreamTextSnapshot = { text: string; truncated: boolean; @@ -734,6 +755,17 @@ export class BoundedStreamTextAccumulator { return this.finishedSnapshot; } + discardRetainedBytes(): void { + this.headChunks.length = 0; + this.tailChunks.length = 0; + this.tailChunkBytes.length = 0; + this.headBufferedBytes = 0; + this.tailBufferedBytes = 0; + this.tailHead = 0; + this.tailMode = false; + this.finishedSnapshot = null; + } + private createSnapshotText(): string { if (!this.tailMode) { return this.decodeChunks(this.headChunks, 0, this.headBufferedBytes); @@ -3532,9 +3564,10 @@ export class ProxyResponseHandler { session.getEndpointPolicy().kind === "raw_passthrough" ? null : mapProviderTypeToFamily(provider.providerType); - const streamProtocolObserver = nativeStreamProtocolFamily + let streamProtocolObserver = nativeStreamProtocolFamily ? createStreamProtocolObserver(nativeStreamProtocolFamily) : null; + const clientAbortMeter = createClientAbortMeteringObserver(session.originalFormat); let protocolObservedBeforeProcessing = false; // --- GEMINI STREAM HANDLING --- @@ -3558,7 +3591,7 @@ export class ProxyResponseHandler { releaseReplayOwnership(session); // F1 shadow 遥测:enforce 已在 forwarder 作用于该流量,shadow 观察同样不留盲区 - const passthroughShadowObserver = (() => { + let passthroughShadowObserver = (() => { if (resolveStreamGateMode() !== "shadow") return null; if (session.getEndpointPolicy().kind === "raw_passthrough") return null; const family = mapProviderTypeToFamily(provider.providerType); @@ -3583,6 +3616,8 @@ export class ProxyResponseHandler { let abortPassthroughTransport = (_reason: Error) => {}; let passthroughPump: DemandDrivenResponsePump; let passthroughDrainTimeoutId: ReturnType | null = null; + let passthroughClientDetached = false; + let passthroughDrainLease: DetachedStreamLease | null = null; const clearPassthroughDrainTimeout = () => { if (passthroughDrainTimeoutId) { clearTimeout(passthroughDrainTimeoutId); @@ -3590,8 +3625,43 @@ export class ProxyResponseHandler { } }; const startPassthroughDrain = (reason?: unknown) => { + if (passthroughPump.getState() === "closed") return; + if (passthroughClientDetached) { + passthroughPump.startDrain(reason); + return; + } + passthroughClientDetached = true; + const admission = acquireDetachedStreamLease( + "metering", + CLIENT_ABORT_DRAIN_RESERVATION_BYTES + ); passthroughPump.startDrain(reason); + streamTextAccumulator.discardRetainedBytes(); + streamProtocolObserver = null; + passthroughShadowObserver = null; + if (!admission.acquired) { + const rejection = new Error(`client_abort_drain_${admission.reason}`); + logger.warn("ResponseHandler: Client abort drain rejected by pool", { + taskId: `stream-passthrough-${messageContext.id}`, + providerId: provider.id, + messageId: messageContext.id, + reason: admission.reason, + budget: getDetachedStreamBudgetSnapshot(), + }); + abortPassthroughTransport(rejection); + passthroughPump.cancelSource(rejection); + return; + } + passthroughDrainLease = admission.lease; + void passthroughPump.teardown.finally(() => { + passthroughDrainLease?.release(); + passthroughDrainLease = null; + }); observePassthroughDrainStart(); + if (clientAbortMeter.observe(new Uint8Array()).billingComplete) { + passthroughPump.finishDrain(new Error("client_abort_metering_complete")); + return; + } if (passthroughDrainTimeoutId) return; passthroughDrainTimeoutId = setTimeout(() => { passthroughDrainTimeoutId = null; @@ -3605,10 +3675,14 @@ export class ProxyResponseHandler { source: response.body, onReadStart: () => observePassthroughReadStart(), onChunk: (value) => { + const metering = clientAbortMeter.observe(value); passthroughShadowObserver?.observe(value); streamProtocolObserver?.observe(value); - streamTextAccumulator.pushBytes(value); + if (!passthroughClientDetached) streamTextAccumulator.pushBytes(value); observePassthroughChunk(value); + if (passthroughClientDetached && metering.billingComplete) { + passthroughPump.finishDrain(new Error("client_abort_metering_complete")); + } }, onClientCancel: (reason) => { startPassthroughDrain(reason); @@ -3746,6 +3820,18 @@ export class ProxyResponseHandler { }; const flushAndSnapshot = (): BoundedStreamTextSnapshot => { + if (passthroughClientDetached) { + const metering = clientAbortMeter.finish(); + const snapshot: BoundedStreamTextSnapshot = { + text: metering.text, + truncated: true, + totalBytes: streamTextAccumulator.totalByteCount, + bufferedBytes: metering.retainedBytes, + chunkCount: streamTextAccumulator.chunkCount, + }; + lastStreamTextSnapshot = snapshot; + return snapshot; + } const snapshot = streamTextAccumulator.finish(); lastStreamTextSnapshot = snapshot; return snapshot; @@ -3846,7 +3932,15 @@ export class ProxyResponseHandler { streamEndedNormally, clientAborted, discoveryLeaseLifecycle, - streamProtocolObserver?.finish() ?? null, + streamProtocolObserver?.finish() ?? + (passthroughClientDetached + ? { + sawContent: false, + sawTerminal: clientAbortMeter.finish().billingComplete, + observationIncomplete: clientAbortMeter.finish().skippedOversizedFrames > 0, + failure: clientAbortMeter.finish().protocolFailure, + } + : null), abortReason ); latestCommitSideEffects = finalized.commitSideEffects; @@ -3919,7 +4013,14 @@ export class ProxyResponseHandler { false, clientAborted, discoveryLeaseLifecycle, - null, + passthroughClientDetached + ? { + sawContent: false, + sawTerminal: clientAbortMeter.finish().billingComplete, + observationIncomplete: clientAbortMeter.finish().skippedOversizedFrames > 0, + failure: clientAbortMeter.finish().protocolFailure, + } + : null, abortReason ); latestCommitSideEffects = finalized.commitSideEffects; @@ -4035,6 +4136,10 @@ export class ProxyResponseHandler { const decoder = new TextDecoder(); const text = decoder.decode(chunk, { stream: true }); buffer += text; + if (buffer.length > GEMINI_STREAM_TRANSFORM_MAX_BUFFER_CHARACTERS) { + buffer = ""; + throw new Error("Gemini stream line exceeded transform buffer limit"); + } const lines = buffer.split("\n"); // Keep the last line in buffer as it might be incomplete @@ -4091,6 +4196,14 @@ export class ProxyResponseHandler { // 让上游响应在客户端断开后继续被缓存直至完成;非 replay 请求维持 60s 现状。 let clientAbortDrainTimeoutMs = CLIENT_ABORT_DRAIN_MAX_MS; let responsePump: DemandDrivenResponsePump | null = null; + let clientAbortDrainLease: DetachedStreamLease | null = null; + let clientAbortReplayLease: DetachedStreamLease | null = null; + let clientAbortDrainMode: "metering" | "replay" | "rejected" | null = null; + let clientAbortFinalizing = false; + let streamReplayCompletionScheduled = false; + let shadowGateObserver: ReturnType | null = null; + let replaySpool: ReturnType = null; + let downgradeDetachedReplay = () => {}; // 提升 idleTimeoutId 到外部作用域,以便客户端断开时能清除 let idleTimeoutId: NodeJS.Timeout | null = null; @@ -4139,11 +4252,14 @@ export class ProxyResponseHandler { clientAbortDrainTimeoutId.unref?.(); }; const capInactiveReplayDrainWindow = () => { - if (clientAbortDrainTimeoutMs <= CLIENT_ABORT_DRAIN_MAX_MS) return; - clientAbortDrainTimeoutMs = CLIENT_ABORT_DRAIN_MAX_MS; - if (clientAbortDrainStartedAt === null) return; - const elapsedMs = Date.now() - clientAbortDrainStartedAt; - scheduleClientAbortDrainTimeout(CLIENT_ABORT_DRAIN_MAX_MS - elapsedMs); + if (clientAbortDrainTimeoutMs > CLIENT_ABORT_DRAIN_MAX_MS) { + clientAbortDrainTimeoutMs = CLIENT_ABORT_DRAIN_MAX_MS; + if (clientAbortDrainStartedAt !== null) { + const elapsedMs = Date.now() - clientAbortDrainStartedAt; + scheduleClientAbortDrainTimeout(CLIENT_ABORT_DRAIN_MAX_MS - elapsedMs); + } + } + downgradeDetachedReplay(); }; const clearIdleTimer = () => { if (idleTimeoutId) { @@ -4208,15 +4324,111 @@ export class ProxyResponseHandler { }; let cleanupClientAbortListener = () => {}; let clientDetachHandled = false; + const releaseDetachedLease = (lease = clientAbortDrainLease) => { + lease?.release(); + if (clientAbortDrainLease === lease) clientAbortDrainLease = null; + }; + const releaseDetachedReplayLease = (lease = clientAbortReplayLease) => { + lease?.release(); + if (clientAbortReplayLease === lease) clientAbortReplayLease = null; + }; + const rejectDetachedDrain = (reason: string) => { + clientAbortDrainMode = "rejected"; + releaseDetachedLease(); + const rejection = new Error(`client_abort_drain_${reason}`); + logger.warn("ResponseHandler: Detached stream rejected by budget", { + taskId, + providerId: provider.id, + messageId: messageContext.id, + reason, + budget: getDetachedStreamBudgetSnapshot(), + }); + responsePump?.startDrain(rejection); + try { + const sessionWithController = session as typeof session & { + responseController?: AbortController; + }; + sessionWithController.responseController?.abort(rejection); + } catch { + // The pump cancellation below remains authoritative. + } + responsePump?.cancelSource(rejection); + }; + const acquireMeteringDrain = (replayAbortReason?: string): boolean => { + const previousLease = clientAbortDrainLease; + clientAbortDrainMode = "metering"; + releaseDetachedLease(previousLease); + if (replayAbortReason && replaySpool && !replaySpool.isTerminal) { + void replaySpool.abort(replayAbortReason); + } + + const admission = acquireDetachedStreamLease( + "metering", + CLIENT_ABORT_DRAIN_RESERVATION_BYTES + ); + if (!admission.acquired) { + rejectDetachedDrain(admission.reason); + return false; + } + + const lease = admission.lease; + clientAbortDrainLease = lease; + const activePump = responsePump; + if (activePump) { + void activePump.teardown.finally(() => releaseDetachedLease(lease)); + } + return true; + }; + downgradeDetachedReplay = () => { + if (!clientDetachHandled || clientAbortDrainMode !== "replay" || clientAbortFinalizing) { + return; + } + logger.info("ResponseHandler: Detached Replay became inactive, using metering drain", { + taskId, + providerId: provider.id, + messageId: messageContext.id, + }); + acquireMeteringDrain(); + }; const handleClientAbort = (reason?: unknown) => { if (responsePump?.getState() === "closed") return; - responsePump?.startDrain(reason ?? "client_detached"); - if (clientDetachHandled) return; + if (clientDetachHandled) { + responsePump?.startDrain(reason ?? "client_detached"); + return; + } clientDetachHandled = true; + const activeReplaySpool = replaySpool && !replaySpool.isTerminal ? replaySpool : null; + if (activeReplaySpool) { + const replayAdmission = acquireDetachedStreamLease( + "replay", + resolveReplayDrainReservationBytes() + ); + if (replayAdmission.acquired) { + clientAbortDrainMode = "replay"; + clientAbortReplayLease = replayAdmission.lease; + } else { + logger.info("ResponseHandler: Detached Replay budget unavailable, using metering drain", { + taskId, + providerId: provider.id, + messageId: messageContext.id, + reason: replayAdmission.reason, + budget: getDetachedStreamBudgetSnapshot(), + }); + if (!acquireMeteringDrain(`detached_replay_${replayAdmission.reason}`)) return; + } + } else if (!acquireMeteringDrain()) { + return; + } + + responsePump?.startDrain(reason ?? "client_detached"); + streamTextAccumulator.discardRetainedBytes(); + streamProtocolObserver = null; + shadowGateObserver = null; logger.debug("ResponseHandler: Client disconnected, cleaning up", { taskId, providerId: provider.id, messageId: messageContext.id, + drainMode: clientAbortDrainMode, }); // Do not cancel internal accounting on pure client disconnect. Transfer // ownership to the bounded background drain so terminal usage can still @@ -4227,6 +4439,9 @@ export class ProxyResponseHandler { } clientAbortDrainStartedAt = Date.now(); scheduleClientAbortDrainTimeout(clientAbortDrainTimeoutMs); + if (clientAbortMeter.observe(new Uint8Array()).billingComplete) { + responsePump?.finishDrain(new Error("client_abort_metering_complete")); + } }; // 统计/结算只保留有界的“头 + 尾”文本快照,避免长流式响应把进程堆撑满。 @@ -4239,6 +4454,17 @@ export class ProxyResponseHandler { // 静默一直等到 60s drain 总上限。 const flushAndJoin = (): string => { + if (clientDetachHandled) { + const metering = clientAbortMeter.finish(); + lastStreamTextSnapshot = { + text: metering.text, + truncated: true, + totalBytes: streamTextAccumulator.totalByteCount, + bufferedBytes: metering.retainedBytes, + chunkCount: streamTextAccumulator.chunkCount, + }; + return metering.text; + } const snapshot = streamTextAccumulator.finish(); lastStreamTextSnapshot = snapshot; return snapshot.text; @@ -4340,10 +4566,21 @@ export class ProxyResponseHandler { abortReason?: string ): Promise => { if (streamFinalizationPromise) return streamFinalizationPromise; + if (clientDetachHandled) clientAbortFinalizing = true; streamFinalizationPromise = (async () => { const finalizationDeadlineAtMs = Date.now() + STREAM_FINALIZATION_MAX_MS; const awaitFinalization = (promise: Promise): Promise => raceWithDeadline(promise, finalizationDeadlineAtMs, "stream_finalization_timeout"); + const detachedProtocolObservation: StreamProtocolObservation | null = (() => { + if (!clientDetachHandled || streamProtocolObserver) return null; + const metering = clientAbortMeter.finish(); + return { + sawContent: false, + sawTerminal: metering.billingComplete, + observationIncomplete: metering.skippedOversizedFrames > 0, + failure: metering.protocolFailure, + }; + })(); const finalized = finalizeDeferredStreamingFinalizationIfNeeded( session, allContent, @@ -4351,7 +4588,7 @@ export class ProxyResponseHandler { streamEndedNormally, clientAborted, discoveryLeaseLifecycle, - streamProtocolObserver?.finish() ?? null, + streamProtocolObserver?.finish() ?? detachedProtocolObservation, abortReason ); latestStreamCommitSideEffects = finalized.commitSideEffects @@ -4641,6 +4878,9 @@ export class ProxyResponseHandler { // F2 终态屏障:replay completed 只能出现在计费落库(onCommitted)之后; // 任何失败终态(假 200/中断/非 2xx)立即 abort,绝不被已完成重放命中。 if (replaySpool) { + const activeReplaySpool = replaySpool; + const detachedReplayLease = + clientAbortDrainMode === "replay" ? clientAbortReplayLease : null; const isReplayableSuccess = finalized.commitSideEffects !== undefined && effectiveStatusCode >= 200 && @@ -4648,19 +4888,23 @@ export class ProxyResponseHandler { !finalized.replayIneligibleReason && hasStreamCompletionMarker(allContent, session.originalFormat); if (isReplayableSuccess) { + streamReplayCompletionScheduled = true; postTerminalSideEffects.push(async () => { try { - await replaySpool.completeAfterBilling(messageContext.id); + await activeReplaySpool.completeAfterBilling(messageContext.id); } catch (err) { logger.warn("[ResponseHandler] Replay spool completion failed:", { error: err }); } }); } else { - void replaySpool.abort( - finalized.replayIneligibleReason ?? - streamErrorMessage ?? - `status_${effectiveStatusCode}` - ); + streamReplayCompletionScheduled = true; + void activeReplaySpool + .abort( + finalized.replayIneligibleReason ?? + streamErrorMessage ?? + `status_${effectiveStatusCode}` + ) + .finally(() => releaseDetachedReplayLease(detachedReplayLease)); } } @@ -4744,7 +4988,12 @@ export class ProxyResponseHandler { // 调用方原样 reject(既有传播语义不变)。 streamFinalizationPromise.catch(() => { if (replaySpool && !replaySpool.isTerminal) { - void replaySpool.abort("finalize_error"); + streamReplayCompletionScheduled = true; + const detachedReplayLease = + clientAbortDrainMode === "replay" ? clientAbortReplayLease : null; + void replaySpool + .abort("finalize_error") + .finally(() => releaseDetachedReplayLease(detachedReplayLease)); } }); return streamFinalizationPromise; @@ -4752,7 +5001,7 @@ export class ProxyResponseHandler { // F1 shadow 模式:旁路逐帧分类,记录「首非空字节 vs 首有效内容」的分歧与延迟差, // 不缓冲、不 failover,仅用于 enforce 灰度前评估误判率。 - const shadowGateObserver = (() => { + shadowGateObserver = (() => { if (resolveStreamGateMode() !== "shadow") return null; if (session.getEndpointPolicy().kind === "raw_passthrough") return null; const family = mapProviderTypeToFamily(provider.providerType); @@ -4766,8 +5015,11 @@ export class ProxyResponseHandler { // F2 owner spool:guard 阶段已抢到 owner 租约的请求,把客户端可见字节 // write-behind 喂入 Redis 热层,供并发/断线的相同请求 attach 跟尾。 - const replaySpool = createReplaySpoolIfOwner(session, response, "stream", { + replaySpool = createReplaySpoolIfOwner(session, response, "stream", { onInactive: capInactiveReplayDrainWindow, + onTerminal: () => { + releaseDetachedReplayLease(); + }, }); if (replaySpool) { try { @@ -4780,20 +5032,25 @@ export class ProxyResponseHandler { const observeChunk = (value: Uint8Array) => { const chunkSize = value.length; clearIdleTimer(); - streamTextAccumulator.pushBytes(value); + const metering = clientAbortMeter.observe(value); AsyncTaskManager.touch(taskId); - shadowGateObserver?.observe(value); - const protocolFailure = protocolObservedBeforeProcessing - ? null - : (streamProtocolObserver?.observe(value) ?? null); - if (protocolFailure && replaySpool && !replaySpool.isTerminal) { - void replaySpool.abort( - `stream_protocol_${protocolFailure.verdict}_${ - protocolFailure.afterContent ? "after" : "before" - }_content` - ); + if (!clientDetachHandled) { + streamTextAccumulator.pushBytes(value); + shadowGateObserver?.observe(value); + const protocolFailure = protocolObservedBeforeProcessing + ? null + : (streamProtocolObserver?.observe(value) ?? null); + if (protocolFailure && replaySpool && !replaySpool.isTerminal) { + void replaySpool.abort( + `stream_protocol_${protocolFailure.verdict}_${ + protocolFailure.afterContent ? "after" : "before" + }_content` + ); + } + if (!replaySpool?.isTerminal) replaySpool?.observe(value); + } else if (clientAbortDrainMode === "replay" && replaySpool && !replaySpool.isTerminal) { + replaySpool.observe(value); } - if (!replaySpool?.isTerminal) replaySpool?.observe(value); logger.trace("ResponseHandler: Upstream stream chunk received", { taskId, @@ -4814,6 +5071,9 @@ export class ProxyResponseHandler { }); } } + if (clientDetachHandled && metering.billingComplete) { + responsePump?.finishDrain(new Error("client_abort_metering_complete")); + } }; responsePump = createDemandDrivenResponsePump({ @@ -5151,6 +5411,20 @@ export class ProxyResponseHandler { clearClientAbortDrainTimer(); clearIdleTimer(); // 清除静默期计时器(防止泄漏) releaseSessionAgent(session); + if ( + clientAbortDrainMode === "replay" && + clientAbortReplayLease && + !streamReplayCompletionScheduled + ) { + const detachedReplayLease = clientAbortReplayLease; + if (replaySpool && !replaySpool.isTerminal) { + void replaySpool + .abort("stream_task_finalized_without_replay_terminal") + .finally(() => releaseDetachedReplayLease(detachedReplayLease)); + } else { + releaseDetachedReplayLease(detachedReplayLease); + } + } if (!streamCommitSideEffectsScheduled) { void (async () => { await latestStreamFinalizeAttemptResources?.(); diff --git a/src/lib/config/env.schema.ts b/src/lib/config/env.schema.ts index 2eb8a6097..d2dbfdf96 100644 --- a/src/lib/config/env.schema.ts +++ b/src/lib/config/env.schema.ts @@ -192,6 +192,21 @@ export const EnvSchema = z.object({ // 超时后主动断开该输家连接,仅用已收到的内容尝试计费(通常计不出 -> 跳过)。 HEDGE_LOSER_DRAIN_TIMEOUT_MS: z.coerce.number().int().min(1000).default(120_000), + // 客户端断线后的 detached stream 使用进程级带权预算。 + DETACHED_STREAM_MAX_CONCURRENCY: z.coerce.number().int().min(1).max(4096).default(64), + DETACHED_STREAM_BUDGET_BYTES: z.coerce + .number() + .int() + .min(64 * 1024) + .max(1024 * 1024 * 1024) + .default(64 * 1024 * 1024), + DETACHED_STREAM_METERING_RESERVE_BYTES: z.coerce + .number() + .int() + .min(64 * 1024) + .max(1024 * 1024 * 1024) + .default(16 * 1024 * 1024), + // ===== CCHP 网关移植功能开关 ===== // 流式内容门控:off=关闭;shadow=旁路分类只记录分歧;enforce=首个有效内容帧前缓冲+failover STREAM_GATE_MODE: z.enum(["off", "shadow", "enforce"]).default("enforce"), diff --git a/tests/configs/detached-stream-budget.config.mts b/tests/configs/detached-stream-budget.config.mts new file mode 100644 index 000000000..4562695cb --- /dev/null +++ b/tests/configs/detached-stream-budget.config.mts @@ -0,0 +1,20 @@ +import { createCoverageConfig } from "../vitest.base.mts"; + +export default createCoverageConfig({ + name: "detached-stream-budget", + environment: "node", + testFiles: [ + "src/app/v1/_lib/proxy/client-abort-metering.test.ts", + "src/app/v1/_lib/proxy/detached-stream-budget.test.ts", + ], + sourceFiles: [ + "src/app/v1/_lib/proxy/client-abort-metering.ts", + "src/app/v1/_lib/proxy/detached-stream-budget.ts", + ], + thresholds: { + branches: 80, + functions: 80, + lines: 80, + statements: 80, + }, +}); diff --git a/tests/load/issue-1408-replay-oom/sample-container.sh b/tests/load/issue-1408-replay-oom/sample-container.sh index f372b81b1..efca09ebc 100755 --- a/tests/load/issue-1408-replay-oom/sample-container.sh +++ b/tests/load/issue-1408-replay-oom/sample-container.sh @@ -18,10 +18,10 @@ container_cgroup_metric() { metric="$2" pid=$(docker inspect -f '{{.State.Pid}}' "$container" 2>/dev/null || true) case "$pid" in - *[!0-9]* | "" | 0) return ;; + *[!0-9]* | "" | 0) return 0 ;; esac cgroup_path=$(awk -F: '$1 == "0" { print $3 }' "/proc/$pid/cgroup" 2>/dev/null || true) - [ -n "$cgroup_path" ] || return + [ -n "$cgroup_path" ] || return 0 metric_path="/sys/fs/cgroup${cgroup_path}/${metric}" [ -r "$metric_path" ] && tr -d '\n' <"$metric_path" } diff --git a/tests/unit/lib/env-detached-stream-budget.test.ts b/tests/unit/lib/env-detached-stream-budget.test.ts new file mode 100644 index 000000000..deba36098 --- /dev/null +++ b/tests/unit/lib/env-detached-stream-budget.test.ts @@ -0,0 +1,28 @@ +import { describe, expect, it } from "vitest"; +import { EnvSchema } from "@/lib/config/env.schema"; + +describe("EnvSchema - detached stream budget", () => { + it("uses bounded defaults", () => { + const env = EnvSchema.parse({}); + expect(env.DETACHED_STREAM_MAX_CONCURRENCY).toBe(64); + expect(env.DETACHED_STREAM_BUDGET_BYTES).toBe(64 * 1024 * 1024); + expect(env.DETACHED_STREAM_METERING_RESERVE_BYTES).toBe(16 * 1024 * 1024); + }); + + it("parses explicit budget limits", () => { + const env = EnvSchema.parse({ + DETACHED_STREAM_MAX_CONCURRENCY: "8", + DETACHED_STREAM_BUDGET_BYTES: String(512 * 1024), + DETACHED_STREAM_METERING_RESERVE_BYTES: String(128 * 1024), + }); + expect(env.DETACHED_STREAM_MAX_CONCURRENCY).toBe(8); + expect(env.DETACHED_STREAM_BUDGET_BYTES).toBe(512 * 1024); + expect(env.DETACHED_STREAM_METERING_RESERVE_BYTES).toBe(128 * 1024); + }); + + it("rejects a budget smaller than one metering reservation", () => { + expect(() => + EnvSchema.parse({ DETACHED_STREAM_BUDGET_BYTES: String(64 * 1024 - 1) }) + ).toThrow(); + }); +}); diff --git a/tests/unit/proxy/replay-spool.test.ts b/tests/unit/proxy/replay-spool.test.ts index f07e67408..7e7f796a1 100644 --- a/tests/unit/proxy/replay-spool.test.ts +++ b/tests/unit/proxy/replay-spool.test.ts @@ -984,6 +984,55 @@ describe("ReplaySpool:isTerminal", () => { }); }); +describe("ReplaySpool:terminal callback", () => { + it("runs exactly once after completed cleanup", async () => { + const onTerminal = vi.fn(); + const spool = new ReplaySpool( + identity, + "owner-token", + 200, + { "content-type": "text/event-stream" }, + "stream", + { onTerminal } + ); + spool.observe(encoder.encode("data: complete\n\n")); + + await spool.completeAfterBilling(1); + await spool.completeAfterBilling(1); + + expect(onTerminal).toHaveBeenCalledTimes(1); + }); + + it("runs exactly once after aborted or disabled cleanup", async () => { + const abortedTerminal = vi.fn(); + const aborted = new ReplaySpool( + identity, + "owner-token", + 200, + { "content-type": "text/event-stream" }, + "stream", + { onTerminal: abortedTerminal } + ); + await aborted.abort("test_abort"); + await aborted.abort("test_abort_again"); + expect(abortedTerminal).toHaveBeenCalledTimes(1); + + envControl.maxPayloadBytes = 4; + const disabledTerminal = vi.fn(); + const disabled = new ReplaySpool( + identity, + "owner-token-disabled", + 200, + { "content-type": "text/event-stream" }, + "stream", + { onTerminal: disabledTerminal } + ); + disabled.observe(encoder.encode("12345678")); + await drainWriteChain(disabled); + expect(disabledTerminal).toHaveBeenCalledTimes(1); + }); +}); + describe("createReplaySpoolIfOwner", () => { it("非 owner 会话返回 null(无租约可释放)", () => { const session = { replayState: null } as unknown as ProxySession; diff --git a/tests/unit/proxy/response-handler-client-abort-drain.test.ts b/tests/unit/proxy/response-handler-client-abort-drain.test.ts index caf4026ae..92f32b2e1 100644 --- a/tests/unit/proxy/response-handler-client-abort-drain.test.ts +++ b/tests/unit/proxy/response-handler-client-abort-drain.test.ts @@ -1,5 +1,9 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { resolveEndpointPolicy } from "@/app/v1/_lib/proxy/endpoint-policy"; +import { + acquireDetachedStreamLease, + getDetachedStreamBudgetSnapshot, +} from "@/app/v1/_lib/proxy/detached-stream-budget"; import { BoundedStreamTextAccumulator, ProxyResponseHandler, @@ -406,6 +410,45 @@ function createPullTrackedResponsesSse(): { }; } +function createMeteringTerminalResponsesSse(): { + response: Response; + cancel: ReturnType; +} { + const encoder = new TextEncoder(); + const chunks = [ + `event: response.output_text.delta\ndata: ${JSON.stringify({ + type: "response.output_text.delta", + delta: "x".repeat(128 * 1024), + })}\n\n`, + `event: response.completed\ndata: ${JSON.stringify({ + type: "response.completed", + response: { + id: "resp_metered", + model: "gpt-5.4-mini-2026-03-17", + usage: { input_tokens: 463, output_tokens: 11 }, + }, + })}\n\n`, + ]; + let index = 0; + const cancel = vi.fn(); + return { + response: new Response( + new ReadableStream({ + pull(controller) { + const chunk = chunks[index++]; + if (chunk) controller.enqueue(encoder.encode(chunk)); + }, + cancel, + }), + { + status: 200, + headers: { "content-type": "text/event-stream" }, + } + ), + cancel, + }; +} + function createControllableTransportErrorResponsesSse(): { response: Response; fail: () => void; @@ -2115,6 +2158,83 @@ describe("ProxyResponseHandler stream client abort finalization", () => { ); }); + it("stops a detached source as soon as compact terminal usage is captured", async () => { + const clientController = new AbortController(); + const session = createSession(clientController.signal); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + const metered = createMeteringTerminalResponsesSse(); + + await ProxyResponseHandler.dispatch(session, metered.response); + clientController.abort(new Error("client detached")); + await drainAsyncTasks(); + + expect(metered.cancel).toHaveBeenCalledWith( + expect.objectContaining({ message: "client_abort_metering_complete" }) + ); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ + statusCode: 200, + inputTokens: 463, + outputTokens: 11, + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + expect(getDetachedStreamBudgetSnapshot().activeStreams).toBe(0); + }); + + it("rejects a new detached drain immediately when the process budget is exhausted", async () => { + const budget = getDetachedStreamBudgetSnapshot(); + const reservation = acquireDetachedStreamLease("metering", budget.limits.maxReservedBytes); + if (!reservation.acquired) throw new Error("expected test budget reservation"); + const clientController = new AbortController(); + const upstreamController = new AbortController(); + try { + const session = createSession(clientController.signal); + Object.assign(session, { responseController: upstreamController }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + await ProxyResponseHandler.dispatch( + session, + createHangingResponsesSse(upstreamController.signal) + ); + clientController.abort(new Error("client detached")); + await drainAsyncTasks(); + + expect(upstreamController.signal.aborted).toBe(true); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 123, + expect.objectContaining({ statusCode: 499, errorMessage: "CLIENT_ABORTED" }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + } finally { + reservation.lease.release(); + } + expect(getDetachedStreamBudgetSnapshot().activeStreams).toBe(0); + }); + it.each([ { bindingIntent: "create" as const, providerId: null }, { bindingIntent: "renew" as const, providerId: 1 }, diff --git a/tests/unit/proxy/response-handler-stream-terminal.test.ts b/tests/unit/proxy/response-handler-stream-terminal.test.ts index 3718a28c3..76d8c5b14 100644 --- a/tests/unit/proxy/response-handler-stream-terminal.test.ts +++ b/tests/unit/proxy/response-handler-stream-terminal.test.ts @@ -1,6 +1,10 @@ import { Context } from "hono"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { ProxyResponseHandler } from "@/app/v1/_lib/proxy/response-handler"; +import { + acquireDetachedStreamLease, + getDetachedStreamBudgetSnapshot, +} from "@/app/v1/_lib/proxy/detached-stream-budget"; 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"; @@ -20,6 +24,8 @@ const mocks = vi.hoisted(() => ({ replayComplete: vi.fn(async () => {}), replayAbort: vi.fn(async () => {}), replayInactive: null as (() => void) | null, + replayTerminal: null as (() => void) | null, + replayTerminalState: false, })); vi.mock("@/app/v1/_lib/proxy/response-fixer", () => ({ @@ -65,14 +71,17 @@ vi.mock("@/app/v1/_lib/proxy/replay/replay-spool", () => ({ session: ProxySession, _response: Response, _delivery: string, - options: { onInactive?: () => void } = {} + options: { onInactive?: () => void; onTerminal?: () => void } = {} ) => { if (session.replayState?.role !== "owner") return null; mocks.replayInactive = options.onInactive ?? null; + mocks.replayTerminal = options.onTerminal ?? null; return { abort: mocks.replayAbort, completeAfterBilling: mocks.replayComplete, - isTerminal: false, + get isTerminal() { + return mocks.replayTerminalState; + }, observe: mocks.replayObserve, }; }, @@ -246,11 +255,38 @@ function setReplayOwner(session: ProxySession, suffix: string): void { }; } +function setReplayFinalization(session: ProxySession): void { + 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", + }); +} + describe("ProxyResponseHandler.dispatch stream terminal behavior", () => { beforeEach(() => { mocks.tasks.length = 0; mocks.replayInactive = null; + mocks.replayTerminal = null; + mocks.replayTerminalState = false; vi.clearAllMocks(); + mocks.replayAbort.mockImplementation(async () => { + mocks.replayTerminalState = true; + mocks.replayInactive?.(); + mocks.replayTerminal?.(); + }); + mocks.replayComplete.mockImplementation(async () => { + mocks.replayTerminalState = true; + mocks.replayTerminal?.(); + }); mocks.durable.mockImplementation(async (_id, _details, options) => { await options?.onCommitted?.(); return true; @@ -302,6 +338,81 @@ describe("ProxyResponseHandler.dispatch stream terminal behavior", () => { expect(releaseAgent).toHaveBeenCalledOnce(); }); + it("keeps an admitted detached Replay complete and replayable", async () => { + const chunks = [ + 'event: message_start\ndata: {"type":"message_start","message":{"usage":{"input_tokens":10,"output_tokens":1}}}\n\n', + 'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"type":"text_delta","text":"kept"}}\n\n', + 'event: message_delta\ndata: {"type":"message_delta","usage":{"output_tokens":4}}\n\n', + 'event: message_stop\ndata: {"type":"message_stop"}\n\n', + ]; + let index = 0; + const source = new ReadableStream({ + pull(controller) { + const chunk = chunks[index++]; + if (chunk) controller.enqueue(new TextEncoder().encode(chunk)); + }, + }); + const { session } = await createSession({}); + setReplayOwner(session, "admitted-detached"); + setReplayFinalization(session); + + const returned = await ProxyResponseHandler.dispatch(session, sseResponse(source)); + const reader = returned.body?.getReader(); + await reader?.read(); + await reader?.cancel(new Error("client disconnected")); + await settleTasks(); + + expect(mocks.replayAbort).not.toHaveBeenCalled(); + expect(mocks.replayComplete).toHaveBeenCalledWith(MESSAGE.id); + const replayedText = mocks.replayObserve.mock.calls + .map(([chunk]) => new TextDecoder().decode(chunk as Uint8Array)) + .join(""); + expect(replayedText).toContain('"text":"kept"'); + expect(replayedText).toContain("message_stop"); + expect(getDetachedStreamBudgetSnapshot().activeStreams).toBe(0); + }); + + it("downgrades a detached Replay to metering when Replay headroom is exhausted", async () => { + const blocker = acquireDetachedStreamLease("replay", 20 * 1024 * 1024); + if (!blocker.acquired) throw new Error("expected Replay budget blocker"); + try { + const chunks = [ + 'event: message_start\ndata: {"type":"message_start","message":{"usage":{"input_tokens":10,"output_tokens":1}}}\n\n', + 'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"type":"text_delta","text":"not replayed"}}\n\n', + 'event: message_delta\ndata: {"type":"message_delta","usage":{"output_tokens":4}}\n\n', + 'event: message_stop\ndata: {"type":"message_stop"}\n\n', + ]; + let index = 0; + const source = new ReadableStream({ + pull(controller) { + const chunk = chunks[index++]; + if (chunk) controller.enqueue(new TextEncoder().encode(chunk)); + }, + }); + const { session } = await createSession({}); + setReplayOwner(session, "metering-fallback"); + setReplayFinalization(session); + + const returned = await ProxyResponseHandler.dispatch(session, sseResponse(source)); + const reader = returned.body?.getReader(); + await reader?.read(); + const observedBeforeDetach = mocks.replayObserve.mock.calls.length; + await reader?.cancel(new Error("client disconnected")); + await settleTasks(); + + expect(mocks.replayAbort).toHaveBeenCalledWith("detached_replay_metering_reserve"); + expect(mocks.replayObserve.mock.calls.length).toBe(observedBeforeDetach); + expect(mocks.durable).toHaveBeenCalledWith( + MESSAGE.id, + expect.objectContaining({ statusCode: 200, inputTokens: 10, outputTokens: 4 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + } finally { + blocker.lease.release(); + } + expect(getDetachedStreamBudgetSnapshot().activeStreams).toBe(0); + }); + it("caps a detached Replay drain at 60 seconds after the spool becomes inactive", async () => { vi.useFakeTimers(); const previousReplayDetachedMs = process.env.REPLAY_MAX_DETACHED_MS; @@ -324,10 +435,24 @@ describe("ProxyResponseHandler.dispatch stream terminal behavior", () => { await reader?.cancel(new Error("client disconnected")); expect(mocks.replayInactive).toEqual(expect.any(Function)); + expect(getDetachedStreamBudgetSnapshot().activeByKind).toEqual({ + metering: 0, + replay: 1, + }); await vi.advanceTimersByTimeAsync(59_999); expect(responseController.signal.aborted).toBe(false); + mocks.replayTerminalState = true; mocks.replayInactive?.(); + expect(getDetachedStreamBudgetSnapshot().activeByKind).toEqual({ + metering: 1, + replay: 1, + }); + mocks.replayTerminal?.(); + expect(getDetachedStreamBudgetSnapshot().activeByKind).toEqual({ + metering: 1, + replay: 0, + }); await vi.advanceTimersByTimeAsync(1); expect(responseController.signal.aborted).toBe(true); @@ -336,6 +461,7 @@ describe("ProxyResponseHandler.dispatch stream terminal behavior", () => { ); expect(cancelSource).toHaveBeenCalledOnce(); await settleTasks(); + expect(getDetachedStreamBudgetSnapshot().activeStreams).toBe(0); } finally { if (previousReplayDetachedMs === undefined) { delete process.env.REPLAY_MAX_DETACHED_MS;