retry stream capacity errors in gateway

scripts/test-gateway-e2e.mjs

@@ -80,6 +80,21 @@ function createTerminatedSseResponse(res, chunks, destroyDelayMs = 20) {
   }, destroyDelayMs);
 }
 
+function createCapacityErrorSseResponse(
+  res,
+  message = "Selected model is at capacity. Please try a different model.",
+  intervalMs = 20,
+) {
+  createSseResponse(
+    res,
+    [
+      'event: error\n',
+      `data: ${JSON.stringify({ error: { message, type: "server_error" } })}\n\n`,
+    ],
+    intervalMs,
+  );
+}
+
 function startFakeUpstream(port) {
   const failBeforeResponseCounts = new Map();
   const capacityBeforeSuccessCounts = new Map();
@@ -143,6 +158,13 @@ function startFakeUpstream(port) {
           const capacityCount = (capacityBeforeSuccessCounts.get(capacityKey) || 0) + 1;
           capacityBeforeSuccessCounts.set(capacityKey, capacityCount);
           if (capacityCount <= parsed.test_capacity_before_success_times) {
+            if (parsed.stream) {
+              createCapacityErrorSseResponse(
+                res,
+                parsed.test_capacity_message || "Selected model is at capacity. Please try a different model.",
+              );
+              return;
+            }
             createJsonResponse(
               res,
               parsed.test_capacity_status ?? 503,
@@ -157,6 +179,13 @@ function startFakeUpstream(port) {
             return;
           }
         }
+        if (parsed.stream && parsed.test_capacity_error) {
+          createCapacityErrorSseResponse(
+            res,
+            parsed.test_capacity_message || "Selected model is at capacity. Please try a different model.",
+          );
+          return;
+        }
         if (parsed.stream) {
           createSseResponse(res, [
             'data: {"type":"response.output_text.delta","delta":"hello"}\n\n',
@@ -434,6 +463,22 @@ async function run() {
       "stream+capacity error 返回体未标记 retry trigger",
     );
 
+    const streamCapacityRecoveredResponse = await fetch(`http://127.0.0.1:${gatewayPort}/responses`, {
+      method: "POST",
+      headers: { "content-type": "application/json" },
+      body: JSON.stringify({ stream: true, test_capacity_before_success_times: 2, test_reasoning_tokens: 128 }),
+    });
+    const streamCapacityRecoveredText = await streamCapacityRecoveredResponse.text();
+    assert(streamCapacityRecoveredResponse.status === 200, `stream capacity 抖动后未自动恢复: ${streamCapacityRecoveredResponse.status}`);
+    assert(streamCapacityRecoveredText.includes("hello"), "stream capacity 恢复后未拿到正常 SSE 内容");
+
+    const requestsAfterStreamCapacityRecoveryResponse = await fetch(`http://127.0.0.1:${gatewayPort}/__codex_retry_gateway/api/requests?limit=20`);
+    const requestsAfterStreamCapacityRecovery = await requestsAfterStreamCapacityRecoveryResponse.json();
+    const streamCapacityRecoveredEntry = requestsAfterStreamCapacityRecovery?.entries?.find(
+      (entry) => entry.path === "/responses" && entry.status_code === 200 && entry.response_stream && entry.upstream_attempt_count >= 3,
+    );
+    assert(streamCapacityRecoveredEntry, "stream capacity 抖动恢复后的请求记录未保留重试次数");
+
     for (const streamPath of [
       "/responses",
       "/v1/responses",