fix: resolve image generation 429 cascade returning 502 to clients

Concurrent /v1/images/edits and /v1/images/generations requests were returning
502 Bad Gateway to clients even though images were successfully generated on
the user's ChatGPT account (quota consumed). Three layered bugs caused every
image request to fail:

Bug A (root cause) - stream parser dropped file_ids
update_conversation_state gated file_id extraction on is_image_tool_event
which requires the full {author.role: tool, metadata.async_task_type:
image_gen} structure. ChatGPT delivers asset_pointer via subsequent patch
events that lack those fields, so file_ids were extracted by regex but
silently discarded. Every image request was forced into the polling
fallback path, which then hammered /backend-api/conversation/{id} and
tripped upstream rate limits.

Fix: open the gate when state.tool_invoked is True (set by an earlier
server_ste_metadata event) OR when the event is itself a patch event whose
payload references asset_pointer / file-service://. The patch-event-only
gate prevents user-controlled message content from injecting fake file ids.
The regex now uses a negative lookahead to exclude the literal "file-service"
URI prefix.

Bug B - polling had zero fault tolerance
_poll_image_results called _get_conversation without try/except. A single
429 / 5xx / network blip killed the entire 180s loop. Production logs also
showed the actual root cause of the spurious 429s: every poll started ~218ms
after the SSE stream closed, but the conversation document is not committed
yet on the upstream and returns a misleading 429 "Too many requests"
instead of 425 Too Early.

Fix: add image_poll_initial_wait_secs (default 10s, +jitter) before the
first poll - image generation takes ~30s anyway. Add exponential backoff
with jitter for 429/5xx/network errors, honoring upstream Retry-After.
All sleeps respect the timeout budget; on exhaustion returns ([], []).

Bug C - upstream error messages were Python repr, not valid JSON
ensure_ok used f-string formatting which produced "body={'detail': 'x'}"
with single quotes - not JSON-parseable by clients.

Fix: introduce UpstreamHTTPError(RuntimeError) subclass exposing
structured fields (status_code, body, retry_after) so callers can branch
programmatically. Format dict/list bodies with json.dumps; truncate long
bodies in the str representation to keep logs readable while preserving
the full body on the exception instance.

Files
services/protocol/conversation.py - relax file_id extraction gate (Bug A)
services/openai_backend_api.py - initial wait + retry in _poll_image_results (Bug B)
utils/helper.py - UpstreamHTTPError + JSON body formatting (Bug C)
services/config.py - image_poll_interval_secs, image_poll_initial_wait_secs

services/config.py

@@ -187,6 +187,23 @@ class ConfigStore:
         except (TypeError, ValueError):
             return 120
 
+    @property
+    def image_poll_interval_secs(self) -> float:
+        try:
+            return max(0.5, float(self.data.get("image_poll_interval_secs", 10.0)))
+        except (TypeError, ValueError):
+            return 10.0
+
+    @property
+    def image_poll_initial_wait_secs(self) -> float:
+        """Image generation upstream takes ~30s; polling immediately wastes requests
+        and trips a transient 429. Default 10s gives the conversation document time
+        to commit before the first poll."""
+        try:
+            return max(0.0, float(self.data.get("image_poll_initial_wait_secs", 10.0)))
+        except (TypeError, ValueError):
+            return 10.0
+
     @property
     def image_account_concurrency(self) -> int:
         try:
@@ -277,6 +294,8 @@ class ConfigStore:
         data["refresh_account_interval_minute"] = self.refresh_account_interval_minute
         data["image_retention_days"] = self.image_retention_days
         data["image_poll_timeout_secs"] = self.image_poll_timeout_secs
+        data["image_poll_interval_secs"] = self.image_poll_interval_secs
+        data["image_poll_initial_wait_secs"] = self.image_poll_initial_wait_secs
         data["image_account_concurrency"] = self.image_account_concurrency
         data["auto_remove_invalid_accounts"] = self.auto_remove_invalid_accounts
         data["auto_remove_rate_limited_accounts"] = self.auto_remove_rate_limited_accounts

services/openai_backend_api.py

@@ -1,5 +1,6 @@
 import base64
 import os
+import random
 import re
 import time
 from concurrent.futures import ThreadPoolExecutor
@@ -14,7 +15,7 @@ from PIL import Image
 from services.account_service import account_service
 from services.config import config
 from services.proxy_service import proxy_settings
-from utils.helper import ensure_ok, iter_sse_payloads, new_uuid
+from utils.helper import UpstreamHTTPError, ensure_ok, iter_sse_payloads, new_uuid
 from utils.log import logger
 from utils.pow import build_legacy_requirements_token, build_proof_token, parse_pow_resources
 from utils.turnstile import solve_turnstile_token
@@ -639,13 +640,76 @@ class OpenAIBackendAPI:
         return sorted(records, key=lambda item: item["create_time"])
 
     def _poll_image_results(self, conversation_id: str, timeout_secs: float = 120.0) -> tuple[list[str], list[str]]:
-        """轮询 conversation，直到拿到图片文件 id 或超时。"""
+        """Poll the conversation document until image file ids appear or budget runs out.
+
+        - Sleeps image_poll_initial_wait_secs first (default 10s, +jitter). ChatGPT
+          image generation takes ~30s; polling immediately wastes requests and trips
+          a transient 429 the upstream returns within ~200ms of the SSE stream
+          closing (the conversation document is not yet committed).
+        - Subsequent polls are image_poll_interval_secs apart (default 10s).
+        - On upstream 429 / 5xx or network errors, backs off exponentially
+          (capped at 16s, +jitter) honoring Retry-After when present.
+        - All sleeps stay within timeout_secs; on exhaustion returns ([], []).
+        """
         start = time.time()
         attempt = 0
-        logger.info({"event": "image_poll_start", "conversation_id": conversation_id, "timeout_secs": timeout_secs})
-        while time.time() - start < timeout_secs:
+        interval = float(config.image_poll_interval_secs)
+        initial_wait = float(config.image_poll_initial_wait_secs)
+        logger.info({
+            "event": "image_poll_start",
+            "conversation_id": conversation_id,
+            "timeout_secs": timeout_secs,
+            "initial_wait_secs": initial_wait,
+            "interval_secs": interval,
+        })
+
+        def _remaining() -> float:
+            return timeout_secs - (time.time() - start)
+
+        if initial_wait > 0:
+            jitter = random.uniform(0, min(2.0, initial_wait * 0.2))
+            sleep_for = min(initial_wait + jitter, max(0.0, _remaining()))
+            if sleep_for > 0:
+                time.sleep(sleep_for)
+
+        def _retry_sleep(reason: str, status_code: int | None, error: str | None, retry_after: int | None) -> bool:
+            # retry_after=0 means "retry immediately" — must not be coerced via falsy check.
+            base = retry_after if retry_after is not None else min(2 ** min(attempt, 4), 16)
+            backoff = base + random.uniform(0, 0.5)
+            remaining = _remaining()
+            if remaining <= 0:
+                return False
+            sleep_for = min(backoff, remaining)
+            log_payload: Dict[str, Any] = {
+                "event": "image_poll_retry",
+                "conversation_id": conversation_id,
+                "attempt": attempt,
+                "reason": reason,
+                "sleep_secs": round(sleep_for, 2),
+            }
+            if status_code is not None:
+                log_payload["status_code"] = status_code
+            if error is not None:
+                log_payload["error"] = error
+            logger.warning(log_payload)
+            time.sleep(sleep_for)
+            return True
+
+        while _remaining() > 0:
             attempt += 1
-            conversation = self._get_conversation(conversation_id)
+            try:
+                conversation = self._get_conversation(conversation_id)
+            except UpstreamHTTPError as exc:
+                if exc.status_code in (429, 500, 502, 503, 504):
+                    if _retry_sleep("upstream_status", exc.status_code, None, exc.retry_after):
+                        continue
+                    break
+                raise
+            except requests.exceptions.RequestException as exc:
+                if _retry_sleep("network", None, str(exc), None):
+                    continue
+                break
+
             file_ids, sediment_ids = [], []
             for record in self._extract_image_tool_records(conversation):
                 for file_id in record["file_ids"]:
@@ -666,8 +730,17 @@ class OpenAIBackendAPI:
                 return [], sediment_ids
             logger.debug({"event": "image_poll_wait", "conversation_id": conversation_id,
                           "elapsed_secs": round(time.time() - start, 1)})
-            time.sleep(4)
-        logger.info({"event": "image_poll_timeout", "conversation_id": conversation_id, "timeout_secs": timeout_secs})
+            wait = min(interval, max(0.0, _remaining()))
+            if wait > 0:
+                time.sleep(wait)
+        logger.info({
+            "event": "image_poll_timeout",
+            "conversation_id": conversation_id,
+            "timeout_secs": timeout_secs,
+            "attempts_made": attempt,
+            # attempts_made == 0 means the initial_wait consumed the entire budget — no HTTP attempted.
+            "initial_wait_exhausted_budget": attempt == 0,
+        })
         return [], []
 
     def _get_file_download_url(self, file_id: str) -> str:

services/protocol/conversation.py

@@ -362,7 +362,8 @@ def add_unique(values: list[str], candidates: list[str]) -> None:
 def extract_conversation_ids(payload: str) -> tuple[str, list[str], list[str]]:
     conversation_match = re.search(r'"conversation_id"\s*:\s*"([^"]+)"', payload)
     conversation_id = conversation_match.group(1) if conversation_match else ""
-    file_ids = re.findall(r"(file[-_][A-Za-z0-9]+)", payload)
+    # Negative lookahead excludes "file-service" (URI prefix, not a real id).
+    file_ids = re.findall(r"(file[-_](?!service\b)[A-Za-z0-9]+)", payload)
     sediment_ids = re.findall(r"sediment://([A-Za-z0-9_-]+)", payload)
     return conversation_id, file_ids, sediment_ids
 
@@ -381,7 +382,19 @@ def update_conversation_state(state: ConversationState, payload: str, event: dic
     conversation_id, file_ids, sediment_ids = extract_conversation_ids(payload)
     if conversation_id and not state.conversation_id:
         state.conversation_id = conversation_id
-    if isinstance(event, dict) and is_image_tool_event(event):
+    # Accept file_id / sediment_id when any of:
+    #   1) event is a complete image_gen tool message
+    #   2) prior server_ste_metadata already flipped tool_invoked True (in an image_gen turn)
+    #   3) patch event whose payload references asset_pointer / file-service://
+    # User messages (type=conversation.message) never satisfy these, so attacker-controlled
+    # substrings in user input cannot inject file ids into state.
+    is_patch_event = isinstance(event, dict) and event.get("o") == "patch"
+    image_context = (
+        (isinstance(event, dict) and is_image_tool_event(event))
+        or state.tool_invoked is True
+        or (is_patch_event and ("asset_pointer" in payload or "file-service://" in payload))
+    )
+    if image_context:
         add_unique(state.file_ids, file_ids)
         add_unique(state.sediment_ids, sediment_ids)
     if not isinstance(event, dict):

utils/helper.py

@@ -27,6 +27,41 @@ def is_image_chat_request(body: dict[str, object]) -> bool:
     return isinstance(modalities, list) and "image" in {str(item or "").strip().lower() for item in modalities}
 
 
+_UPSTREAM_BODY_LOG_LIMIT = 500
+
+
+class UpstreamHTTPError(RuntimeError):
+    """Raised when an upstream HTTP call returns a non-2xx status.
+
+    Carries structured fields (status_code, body, retry_after) so callers can
+    branch on status code instead of string-matching on str(exc). The full
+    body is preserved on the instance; the formatted message truncates it
+    to keep log lines reasonable.
+    """
+
+    def __init__(
+        self,
+        context: str,
+        status_code: int,
+        body: Any,
+        retry_after: int | None = None,
+    ) -> None:
+        self.context = context
+        self.status_code = status_code
+        self.body = body
+        self.retry_after = retry_after
+        if isinstance(body, (dict, list)):
+            try:
+                body_str = json.dumps(body, ensure_ascii=False)
+            except (TypeError, ValueError):
+                body_str = repr(body)
+        else:
+            body_str = str(body)
+        if len(body_str) > _UPSTREAM_BODY_LOG_LIMIT:
+            body_str = body_str[:_UPSTREAM_BODY_LOG_LIMIT] + "…[truncated]"
+        super().__init__(f"{context} failed: status={status_code}, body={body_str}")
+
+
 def ensure_ok(response: requests.Response, context: str) -> None:
     if 200 <= response.status_code < 300:
         return
@@ -35,7 +70,13 @@ def ensure_ok(response: requests.Response, context: str) -> None:
         body = response.json()
     except Exception:
         pass
-    raise RuntimeError(f"{context} failed: status={response.status_code}, body={body}")
+    retry_after_header = response.headers.get("Retry-After") if hasattr(response, "headers") else None
+    retry_after: int | None = None
+    if retry_after_header is not None:
+        ra_str = str(retry_after_header).strip()
+        if ra_str.isdigit():
+            retry_after = int(ra_str)
+    raise UpstreamHTTPError(context, response.status_code, body, retry_after=retry_after)
 
 
 def sse_json_stream(items) -> Iterator[str]: