feat: track thread ids and retry stream disconnects

scripts/test-gateway-e2e.mjs

@@ -232,7 +232,7 @@ function startFakeUpstream(port) {
         }
         if (parsed.stream) {
           createSseResponse(res, [
-            'data: {"type":"response.output_text.delta","delta":"hello"}\n\n',
+            `data: ${JSON.stringify({ type: "response.output_text.delta", delta: "hello", response_id: "resp_stream", thread_id: parsed.thread_id || "thread_stream" })}\n\n`,
             `data: {"response":{"usage":{"output_tokens_details":{"reasoning_tokens":${reasoning}}}}}\n\n`,
             "data: [DONE]\n\n",
           ], parsed.test_stream_chunk_delay_ms ?? 20);
@@ -243,6 +243,7 @@ function startFakeUpstream(port) {
           200,
           {
             id: "resp_test",
+            thread_id: parsed.thread_id || "thread_test",
             retry_attempt: parsed.test_fail_before_response_once
               ? failBeforeResponseCounts.get(`${req.url}:fail-before-response-once`) || 0
               : 0,
@@ -383,7 +384,10 @@ async function run() {
     endpoints: ["/responses", "/chat/completions", "/v1/responses", "/v1/chat/completions"],
     reasoning_equals: [516],
     retryable_status_codes: [429, 503],
-    retryable_error_messages: ["Selected model is at capacity. Please try a different model."],
+    retryable_error_messages: [
+      "Selected model is at capacity. Please try a different model.",
+      "stream disconnected before completion: Concurrency limit exceeded for account, please retry later",
+    ],
     upstream_fetch_retry_attempts: 5,
     upstream_fetch_retry_backoff_ms: 25,
     non_stream_status_code: 502,
@@ -461,6 +465,22 @@ async function run() {
     );
     assert(recoveredEntry?.request_id, "请求记录未生成 request_id");
 
+    const threadTrackedResponse = await fetch(`http://127.0.0.1:${gatewayPort}/responses`, {
+      method: "POST",
+      headers: { "content-type": "application/json" },
+      body: JSON.stringify({ test_reasoning_tokens: 128, thread_id: "thread_nonstream" }),
+    });
+    const threadTrackedBody = await threadTrackedResponse.json();
+    assert(threadTrackedResponse.status === 200, `thread non-stream 请求失败: ${threadTrackedResponse.status}`);
+    assert(threadTrackedBody?.id === "resp_test", "thread non-stream 返回体缺少 response id");
+    assert(threadTrackedBody?.thread_id === "thread_nonstream", "thread non-stream 返回体缺少 thread_id");
+
+    const threadRequestsResponse = await fetch(`http://127.0.0.1:${gatewayPort}/__codex_retry_gateway/api/requests?query=${encodeURIComponent("thread_nonstream")}`);
+    const threadRequestsPayload = await threadRequestsResponse.json();
+    const threadEntry = (threadRequestsPayload?.entries || []).find((entry) => entry.thread_id === "thread_nonstream");
+    assert(threadEntry?.response_id === "resp_test", "non-stream 请求记录未保留 response_id");
+    assert(threadEntry?.thread_id === "thread_nonstream", "non-stream 请求记录未保留 thread_id");
+
     const sameRequestPayload = JSON.stringify({ test_reasoning_tokens: 128, test_request_id_marker: "same" });
     const sameRequestFirstResponse = await fetch(`http://127.0.0.1:${gatewayPort}/responses`, {
       method: "POST",
@@ -669,6 +689,31 @@ async function run() {
     );
     assert(streamCapacityResponseFailedRecoveredEntry, "stream response.failed capacity 恢复后的请求记录未保留重试次数");
 
+    const streamThreadResponse = await readSseUntilClose(
+      `http://127.0.0.1:${gatewayPort}/responses`,
+      { stream: true, test_reasoning_tokens: 128, thread_id: "thread_stream_ok" },
+    );
+    assert(streamThreadResponse.status === 200, `stream thread 请求失败: ${streamThreadResponse.status}`);
+    const streamThreadRequestsResponse = await fetch(`http://127.0.0.1:${gatewayPort}/__codex_retry_gateway/api/requests?query=${encodeURIComponent("thread_stream_ok")}`);
+    const streamThreadRequestsPayload = await streamThreadRequestsResponse.json();
+    const streamThreadEntry = (streamThreadRequestsPayload?.entries || []).find((entry) => entry.thread_id === "thread_stream_ok");
+    assert(streamThreadEntry?.response_id === "resp_stream", "stream 请求记录未保留 response_id");
+    assert(streamThreadEntry?.thread_id === "thread_stream_ok", "stream 请求记录未保留 thread_id");
+
+    const streamDisconnectedRetryResponse = 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_capacity_message: "stream disconnected before completion: Concurrency limit exceeded for account, please retry later",
+        test_reasoning_tokens: 128,
+      }),
+    });
+    const streamDisconnectedRetryText = await streamDisconnectedRetryResponse.text();
+    assert(streamDisconnectedRetryResponse.status === 200, `stream disconnected capacity 抖动后未自动恢复: ${streamDisconnectedRetryResponse.status}`);
+    assert(streamDisconnectedRetryText.includes("hello"), "stream disconnected capacity 恢复后未拿到正常 SSE 内容");
+
     for (const streamPath of [
       "/responses",
       "/v1/responses",