fix: replay buffered success streams incrementally

2026-07-08 21:27:39 +08:00
parent f9a71663f6
commit 26070fb7d5
3 changed files with 139 additions and 39 deletions
@@ -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`
@@ -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;
@@ -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`, {