shujakuin/codex-retry-gateway
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: replay buffered success streams incrementally

Browse Source
This commit is contained in:
yunyaozhou
2026-07-08 21:27:39 +08:00
parent f9a71663f6
commit 26070fb7d5
3 changed files with 139 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1
View File
@@ -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`
gateway.mjs
+73 -7
View File
@@ -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;
scripts/test-gateway-e2e.mjs
+36 -3
View File
@@ -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}` }
: {},
@@ -431,6 +439,7 @@ function startGateway(configPath, logPath) {
}
async function readSseUntilClose(url, requestBody) {
const startedAt = Date.now();
const response = await fetch(url, {
method: "POST",
headers: { "content-type": "application/json" },
@@ -441,6 +450,8 @@ async function readSseUntilClose(url, requestBody) {
const decoder = new TextDecoder("utf8");
let text = "";
let closedByError = false;
let readCount = 0;
let firstChunkDelayMs = null;
while (true) {
try {
@@ -448,6 +459,10 @@ async function readSseUntilClose(url, requestBody) {
if (done) {
break;
}
readCount += 1;
if (firstChunkDelayMs === null) {
firstChunkDelayMs = Date.now() - startedAt;
}
text += decoder.decode(value, { stream: true });
} catch (error) {
closedByError = true;
@@ -462,6 +477,8 @@ async function readSseUntilClose(url, requestBody) {
headers: response.headers,
text,
closedByError,
readCount,
firstChunkDelayMs,
};
}
@@ -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`, {