fix: replay buffered success streams incrementally

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;