diff --git a/README.md b/README.md index 6b1cbc1..6cf0c74 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,7 @@ - 上游若返回明确的容量错误(默认匹配错误文案 `Selected model is at capacity. Please try a different model.`,以及 `stream disconnected before completion: Concurrency limit exceeded for account, please retry later`),也会自动重试;重试耗尽后转成本地 `502` - 除了 `429/503` JSON 错误响应,也会识别 `200` 但返回体本质是错误、以及流式失败事件里携带同样文案的情况 - 流式命中时默认先缓存并判断;一旦命中 `516`,统一返回 `502` +- 流式成功响应在严格检查模式下仍会先缓存完成,但会按分块 SSE 回放给 Codex,而不是单次大包 `res.end(...)` - 默认同时拦截 root 路径和 `/v1` 路径: - `/responses` - `/chat/completions` diff --git a/gateway.mjs b/gateway.mjs index a7095e0..71cbe0b 100644 --- a/gateway.mjs +++ b/gateway.mjs @@ -2879,21 +2879,81 @@ function cloneResponseHeaders(sourceHeaders) { return headers; } -function writeCapturedResponse(res, delivery) { +const CAPTURED_STREAM_REPLAY_CHUNKS_PER_TICK = 16; + +function cloneBufferList(chunks) { + if (!Array.isArray(chunks) || chunks.length === 0) { + return null; + } + return chunks + .filter((chunk) => chunk && chunk.length > 0) + .map((chunk) => (Buffer.isBuffer(chunk) ? Buffer.from(chunk) : Buffer.from(chunk))); +} + +async function writeBufferedStreamChunks(res, chunks) { + let writesSinceYield = 0; + for (const chunk of chunks || []) { + if (!chunk || chunk.length === 0 || res.destroyed || res.writableEnded) { + continue; + } + const accepted = res.write(chunk); + writesSinceYield += 1; + if (!accepted) { + await new Promise((resolve) => { + let settled = false; + const finish = () => { + if (settled) { + return; + } + settled = true; + res.off("drain", onDrain); + res.off("close", onClose); + res.off("error", onClose); + resolve(); + }; + const onDrain = () => finish(); + const onClose = () => finish(); + res.once("drain", onDrain); + res.once("close", onClose); + res.once("error", onClose); + }); + writesSinceYield = 0; + continue; + } + if (writesSinceYield >= CAPTURED_STREAM_REPLAY_CHUNKS_PER_TICK) { + writesSinceYield = 0; + await new Promise((resolve) => setImmediate(resolve)); + } + } +} + +async function writeCapturedResponse(res, delivery) { if (!delivery) { throw new Error("missing captured response delivery"); } copyHeadersToClient(delivery.headers, res); res.writeHead(delivery.status_code); + if (Array.isArray(delivery.stream_chunks) && delivery.stream_chunks.length > 0) { + await writeBufferedStreamChunks(res, delivery.stream_chunks); + if (!res.writableEnded) { + res.end(); + } + return; + } res.end(delivery.body); } -function buildCapturedDelivery(statusCode, headers, body) { - return { +function buildCapturedDelivery(statusCode, headers, body, options = {}) { + const delivery = { status_code: statusCode, headers: cloneResponseHeaders(headers), body: Buffer.isBuffer(body) ? body : Buffer.from(body || ""), }; + const streamChunks = cloneBufferList(options.stream_chunks); + if (streamChunks && streamChunks.length > 0) { + delivery.stream_chunks = streamChunks; + } + return delivery; } function createAbortReason(code, message) { @@ -3680,11 +3740,15 @@ async function handleStreaming({ persistStreamingProgress(runtime, requestEntry, { force: true }, new Date()); } if (strict502Mode) { - const bufferedBody = Buffer.concat(bufferedChunks); + const replayChunks = cloneBufferList(bufferedChunks) || []; + const bufferedBody = Buffer.concat(replayChunks); if (!captureOnly) { copyHeadersToClient(upstreamResponse.headers, res); res.writeHead(upstreamResponse.status); - res.end(bufferedBody); + await writeBufferedStreamChunks(res, replayChunks); + if (!res.writableEnded) { + res.end(); + } } return { inspected: true, @@ -3698,7 +3762,9 @@ async function handleStreaming({ response_bytes_received: requestEntry.response_bytes_received, stream_chunk_count: requestEntry.stream_chunk_count, delivery: captureOnly - ? buildCapturedDelivery(upstreamResponse.status, upstreamResponse.headers, bufferedBody) + ? buildCapturedDelivery(upstreamResponse.status, upstreamResponse.headers, bufferedBody, { + stream_chunks: replayChunks, + }) : null, }; } else { @@ -4513,7 +4579,7 @@ async function proxyRequest(runtime, req, res) { const { delivery, total_upstream_attempts: observedUpstreamAttempts, ...resultFields } = result; if (delivery && !res.headersSent) { - writeCapturedResponse(res, delivery); + await writeCapturedResponse(res, delivery); } if (Number.isInteger(observedUpstreamAttempts)) { requestEntry.upstream_attempt_count = observedUpstreamAttempts; diff --git a/scripts/test-gateway-e2e.mjs b/scripts/test-gateway-e2e.mjs index 3029108..a36780f 100644 --- a/scripts/test-gateway-e2e.mjs +++ b/scripts/test-gateway-e2e.mjs @@ -299,11 +299,19 @@ function startFakeUpstream(port) { return; } if (parsed.stream) { - createSseResponse(res, [ - `data: ${JSON.stringify({ type: "response.output_text.delta", delta: "hello", response_id: "resp_stream", thread_id: parsed.thread_id || "thread_stream", retry_attempt: reasoningAttempt })}\n\n`, + const deltaChunkCount = Number.isInteger(parsed.test_stream_delta_chunks) && parsed.test_stream_delta_chunks > 0 + ? parsed.test_stream_delta_chunks + : 1; + const deltaTextBase = parsed.test_stream_delta_text || "hello"; + const streamChunks = Array.from({ length: deltaChunkCount }, (_, index) => { + const deltaText = deltaChunkCount === 1 ? deltaTextBase : `${deltaTextBase}-${index + 1}`; + return `data: ${JSON.stringify({ type: "response.output_text.delta", delta: deltaText, response_id: "resp_stream", thread_id: parsed.thread_id || "thread_stream", retry_attempt: reasoningAttempt })}\n\n`; + }); + streamChunks.push( `data: {"response":{"usage":{"output_tokens_details":{"reasoning_tokens":${reasoning}}}}}\n\n`, "data: [DONE]\n\n", - ], parsed.test_reasoning_response_delay_ms ?? parsed.test_stream_chunk_delay_ms ?? 20, { + ); + createSseResponse(res, streamChunks, parsed.test_reasoning_response_delay_ms ?? parsed.test_stream_chunk_delay_ms ?? 20, { headers: reasoningAttempt ? { "x-upstream-reasoning-attempt": `${reasoningAttempt}` } : {}, @@ -430,40 +438,49 @@ function startGateway(configPath, logPath) { }; } -async function readSseUntilClose(url, requestBody) { - const response = await fetch(url, { - method: "POST", - headers: { "content-type": "application/json" }, - body: JSON.stringify(requestBody), - }); +async function readSseUntilClose(url, requestBody) { + const startedAt = Date.now(); + const response = await fetch(url, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(requestBody), + }); - const reader = response.body.getReader(); - const decoder = new TextDecoder("utf8"); - let text = ""; - let closedByError = false; - - while (true) { - try { - const { done, value } = await reader.read(); - if (done) { - break; - } - text += decoder.decode(value, { stream: true }); - } catch (error) { - closedByError = true; - text += `\n[[reader-error:${error?.name || "unknown"}]]`; - break; + const reader = response.body.getReader(); + const decoder = new TextDecoder("utf8"); + let text = ""; + let closedByError = false; + let readCount = 0; + let firstChunkDelayMs = null; + + while (true) { + try { + const { done, value } = await reader.read(); + if (done) { + break; + } + readCount += 1; + if (firstChunkDelayMs === null) { + firstChunkDelayMs = Date.now() - startedAt; + } + text += decoder.decode(value, { stream: true }); + } catch (error) { + closedByError = true; + text += `\n[[reader-error:${error?.name || "unknown"}]]`; + break; } } text += decoder.decode(); - return { - status: response.status, - headers: response.headers, - text, - closedByError, - }; -} + return { + status: response.status, + headers: response.headers, + text, + closedByError, + readCount, + firstChunkDelayMs, + }; +} async function run() { assert( @@ -1052,6 +1069,22 @@ async function run() { assert(okStream.status === 200, `${streamPath} 128 首状态异常: ${okStream.status}`); assert(okStream.text.includes("[DONE]"), `${streamPath} 流式 128 未完整结束`); assert(!okStream.closedByError, `${streamPath} 流式 128 不应异常断开`); + + if (streamPath === "/responses" || streamPath === "/v1/responses") { + const replayedStream = await readSseUntilClose( + `http://127.0.0.1:${gatewayPort}${streamPath}`, + { + stream: true, + test_reasoning_tokens: 128, + test_stream_delta_chunks: 48, + test_stream_delta_text: "chunk", + test_stream_chunk_delay_ms: 2, + }, + ); + assert(replayedStream.status === 200, `${streamPath} 回放流首状态异常: ${replayedStream.status}`); + assert(replayedStream.text.includes("chunk-48"), `${streamPath} 回放流未保留尾部 delta`); + assert(replayedStream.readCount > 1, `${streamPath} 成功流不应退化为单块回放`); + } } const streamProgressPromise = fetch(`http://127.0.0.1:${gatewayPort}/responses`, {