diff --git a/plugins/image_gen/openai-codex/__init__.py b/plugins/image_gen/openai-codex/__init__.py index 2cee336051..3104aec471 100644 --- a/plugins/image_gen/openai-codex/__init__.py +++ b/plugins/image_gen/openai-codex/__init__.py @@ -291,6 +291,20 @@ def _normalize_input_images( return [_to_input_image_part(value) for value in values] +# Progressive preview frames (partial_image_b64) are intermediate renders. +# Saving them as finals produced the long-running "smear" failure mode on the +# Codex Responses path. Defense in depth: +# 1) request layer prefers no progressive frames when the backend honors it +# 2) extractor never lets a partial overwrite a final result +# 3) generate() only delivers source=final; partial-only / empty are not success +# Live streams sometimes still emit a partial event even with 0; that is fine as +# long as only a final ``result`` can be saved. +_PARTIAL_IMAGES_REQUESTED = 0 +# Content-agnostic retries when the stream does not yield a final result +# (empty stream or progressive-only). No prompt-class branching. +_NONFINAL_RETRIES = 1 + + def _build_responses_payload( *, prompt: str, @@ -318,7 +332,12 @@ def _build_responses_payload( "quality": quality, "output_format": "png", "background": "opaque", - "partial_images": 1, + # Prefer 0 progressive preview frames. Preview frames can arrive + # without a later final ``result`` and look like smeared / + # unfinished images if saved as the deliverable. Even when the + # backend still emits a partial event, generate() refuses to + # deliver anything except source=final. + "partial_images": _PARTIAL_IMAGES_REQUESTED, }], # No ``tool_choice`` is sent: the chatgpt.com/backend-api/codex backend # rejects every shape we have for forcing the hosted ``image_generation`` @@ -333,27 +352,57 @@ def _build_responses_payload( } +def _extract_image_candidates(value: Any) -> Tuple[Optional[str], Optional[str]]: + """Return ``(final_result_b64, latest_partial_b64)`` from a payload tree. + + Final ``image_generation_call.result`` and progressive ``partial_image_b64`` + are tracked separately so a partial can never overwrite a genuine final, + including when both coexist in the same event payload. + """ + result_b64: Optional[str] = None + partial_b64: Optional[str] = None + + def walk(node: Any) -> None: + nonlocal result_b64, partial_b64 + if isinstance(node, dict): + if node.get("type") == "image_generation_call": + result = node.get("result") + if isinstance(result, str) and result: + result_b64 = result + partial = node.get("partial_image_b64") + if isinstance(partial, str) and partial: + partial_b64 = partial + for child in node.values(): + walk(child) + elif isinstance(node, list): + for child in node: + walk(child) + + walk(value) + return result_b64, partial_b64 + + def _extract_image_b64(value: Any) -> Optional[str]: - """Return the newest image b64 embedded in a Responses event payload.""" - found: Optional[str] = None - if isinstance(value, dict): - if value.get("type") == "image_generation_call": - result = value.get("result") - if isinstance(result, str) and result: - found = result - partial = value.get("partial_image_b64") - if isinstance(partial, str) and partial: - found = partial - for child in value.values(): - nested = _extract_image_b64(child) - if nested: - found = nested - elif isinstance(value, list): - for child in value: - nested = _extract_image_b64(child) - if nested: - found = nested - return found + """Return image b64 from a payload, preferring final result over partial. + + Progressive ``partial_image_b64`` is only used when no final + ``image_generation_call.result`` is present in the same payload tree. + """ + result_b64, partial_b64 = _extract_image_candidates(value) + return result_b64 or partial_b64 + + +def _png_pixel_size(raw: bytes) -> Optional[str]: + """Return ``\"{w}x{h}\"`` for a PNG payload, or None if not a PNG IHDR.""" + import struct + + if len(raw) < 24 or raw[:8] != b"\x89PNG\r\n\x1a\n": + return None + # IHDR: length(4) + type(4) + width(4) + height(4) + if raw[12:16] != b"IHDR": + return None + width, height = struct.unpack(">II", raw[16:24]) + return f"{width}x{height}" def _iter_sse_json(response: Any): @@ -410,8 +459,15 @@ def _collect_image_b64( size: str, quality: str, input_images: Optional[List[Dict[str, str]]] = None, -) -> Optional[str]: - """Stream a Codex Responses image_generation call and return the b64 image.""" +) -> Optional[Dict[str, str]]: + """Stream a Codex Responses image_generation call. + + Returns ``{\"b64\": ..., \"source\": \"final\"|\"partial\"}`` or ``None``. + + Final ``result`` frames are preferred across the whole stream. A progressive + ``partial_image_b64`` is retained only when no final result ever arrives; + callers must not treat partial-only as an unconditional success. + """ import httpx from agent.auxiliary_client import _codex_cloudflare_headers @@ -429,7 +485,8 @@ def _collect_image_b64( ) timeout = httpx.Timeout(300.0, connect=30.0, read=300.0, write=30.0, pool=30.0) - image_b64: Optional[str] = None + final_b64: Optional[str] = None + partial_b64: Optional[str] = None with httpx.Client(timeout=timeout, headers=headers) as http: with http.stream("POST", f"{_CODEX_BASE_URL}/responses", json=payload) as response: try: @@ -441,11 +498,17 @@ def _collect_image_b64( f"{_summarize_error_body(exc.response.text)}" ) from exc for event in _iter_sse_json(response): - found = _extract_image_b64(event) - if found: - image_b64 = found + result_b64, event_partial = _extract_image_candidates(event) + if result_b64: + final_b64 = result_b64 + if event_partial: + partial_b64 = event_partial - return image_b64 + if final_b64: + return {"b64": final_b64, "source": "final"} + if partial_b64: + return {"b64": partial_b64, "source": "partial"} + return None # --------------------------------------------------------------------------- @@ -578,13 +641,32 @@ class OpenAICodexImageGenProvider(ImageGenProvider): ) try: - b64 = _collect_image_b64( - token, - prompt=prompt, - size=size, - quality=meta["quality"], - input_images=input_images or None, - ) + collected: Optional[Dict[str, str]] = None + for attempt in range(_NONFINAL_RETRIES + 1): + collected = _collect_image_b64( + token, + prompt=prompt, + size=size, + quality=meta["quality"], + input_images=input_images or None, + ) + if collected and collected.get("source") == "final" and collected.get("b64"): + break + if attempt < _NONFINAL_RETRIES: + kind = ( + "progressive-only partial frame" + if collected and collected.get("source") == "partial" + else "no image_generation_call result" + ) + logger.warning( + "Codex image stream ended with %s (attempt %s/%s); " + "retrying once before failing closed.", + kind, + attempt + 1, + _NONFINAL_RETRIES + 1, + ) + continue + break except Exception as exc: logger.debug("Codex image generation failed", exc_info=True) return error_response( @@ -596,9 +678,12 @@ class OpenAICodexImageGenProvider(ImageGenProvider): aspect_ratio=aspect, ) - if not b64: + if not collected or not collected.get("b64"): return error_response( - error="Codex response contained no image_generation_call result", + error=( + "Codex response contained no image_generation_call result " + f"after {_NONFINAL_RETRIES + 1} attempt(s)" + ), error_type="empty_response", provider="openai-codex", model=tier_id, @@ -606,7 +691,46 @@ class OpenAICodexImageGenProvider(ImageGenProvider): aspect_ratio=aspect, ) + image_source = collected.get("source") or "unknown" + b64 = collected["b64"] + + # Defense in depth: never deliver a progressive-only frame as success. + # Partials are intermediate previews and have presented as smeared / + # unfinished images when saved as finals. + if image_source != "final": + pixel_hint = None + try: + import base64 as _b64mod + + pixel_hint = _png_pixel_size(_b64mod.b64decode(b64, validate=False)) + except Exception: + pixel_hint = None + detail = ( + "Codex returned only a progressive partial image frame after " + f"{_NONFINAL_RETRIES + 1} attempt(s); refusing to save it " + "as a final deliverable." + ) + if pixel_hint: + detail = f"{detail} partial_pixel_size={pixel_hint}." + err = error_response( + error=detail, + error_type="incomplete_image", + provider="openai-codex", + model=tier_id, + prompt=prompt, + aspect_ratio=aspect, + ) + err["image_source"] = image_source + err["requested_size"] = size + err["partial_pixel_size"] = pixel_hint + err["nonfinal_retries"] = _NONFINAL_RETRIES + return err + try: + import base64 as _b64mod + + raw_bytes = _b64mod.b64decode(b64) + pixel_size = _png_pixel_size(raw_bytes) saved_path = save_b64_image(b64, prefix=f"openai_codex_{tier_id}") except Exception as exc: return error_response( @@ -625,7 +749,14 @@ class OpenAICodexImageGenProvider(ImageGenProvider): aspect_ratio=aspect, provider="openai-codex", modality="image" if input_images else "text", - extra={"size": size, "quality": meta["quality"], "input_image_count": len(input_images)}, + extra={ + "size": size, + "quality": meta["quality"], + "input_image_count": len(input_images), + "image_source": image_source, + "requested_size": size, + "pixel_size": pixel_size, + }, ) diff --git a/tests/plugins/image_gen/test_openai_codex_provider.py b/tests/plugins/image_gen/test_openai_codex_provider.py index 150b92d3d3..244b720bf1 100644 --- a/tests/plugins/image_gen/test_openai_codex_provider.py +++ b/tests/plugins/image_gen/test_openai_codex_provider.py @@ -104,7 +104,7 @@ class TestGenerate: def test_generate_uses_codex_stream_path(self, provider, monkeypatch, tmp_path): monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token") - monkeypatch.setattr(codex_plugin, "_collect_image_b64", lambda *a, **kw: _b64_png()) + monkeypatch.setattr(codex_plugin, "_collect_image_b64", lambda *a, **kw: {"b64": _b64_png(), "source": "final"}) result = provider.generate("a cat", aspect_ratio="landscape") @@ -112,6 +112,8 @@ class TestGenerate: assert result["model"] == "gpt-image-2-medium" assert result["provider"] == "openai-codex" assert result["quality"] == "medium" + assert result.get("image_source") == "final" + assert result.get("pixel_size") == "1x1" saved = Path(result["image"]) assert saved.exists() @@ -132,7 +134,7 @@ class TestGenerate: quality=quality, input_images=input_images, )) - return _b64_png() + return {"b64": _b64_png(), "source": "final"} monkeypatch.setattr(codex_plugin, "_collect_image_b64", _collect) @@ -156,7 +158,9 @@ class TestGenerate: assert tool["size"] == "1024x1536" assert tool["output_format"] == "png" assert tool["background"] == "opaque" - assert tool["partial_images"] == 1 + # Progressive previews disabled: partial frames were being saved as + # finals and presented as smeared/unfinished images. + assert tool["partial_images"] == 0 def test_capabilities_advertise_image_inputs(self, provider): caps = provider.capabilities() @@ -177,12 +181,52 @@ class TestGenerate: def test_partial_image_event_used_when_done_missing(self): - """If output_item.done is missing, partial_image_b64 is accepted.""" + """Extractor may surface partial b64 when no final exists (fallback only).""" payload = { "type": "response.image_generation_call.partial_image", "partial_image_b64": _b64_png(), } assert codex_plugin._extract_image_b64(payload) == _b64_png() + result, partial = codex_plugin._extract_image_candidates(payload) + assert result is None + assert partial == _b64_png() + + def test_final_result_wins_over_coexisting_partial_in_same_payload(self): + """Blind spot that shipped the smear bug: both fields in one payload. + + partial_image_b64 must never overwrite image_generation_call.result + when they coexist in the same event tree. + """ + final = _b64_png() + # Distinct non-empty stand-in so equality proves which field won. + partial = "cGFydGlhbC1vbmx5LW5vdC1hLXJlYWwtZmluYWw=" + payload = { + "type": "response.output_item.done", + "item": { + "type": "image_generation_call", + "status": "completed", + "result": final, + "partial_image_b64": partial, + }, + } + assert codex_plugin._extract_image_b64(payload) == final + result, got_partial = codex_plugin._extract_image_candidates(payload) + assert result == final + assert got_partial == partial + + def test_nested_final_wins_over_sibling_partial(self): + payload = { + "type": "response.completed", + "response": { + "output": [{ + "type": "image_generation_call", + "status": "completed", + "result": _b64_png(), + }], + }, + "partial_image_b64": "cGFydGlhbC1zaWJsaW5n", + } + assert codex_plugin._extract_image_b64(payload) == _b64_png() def test_sse_parser_handles_event_and_data_lines(self): class _Response: @@ -214,8 +258,76 @@ class TestGenerate: } assert codex_plugin._extract_image_b64(payload) == _b64_png() + def test_partial_only_stream_fails_closed_after_retry(self, provider, monkeypatch): + """Partial-only streams must not return success:true with a smear frame.""" + monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token") + calls = {"n": 0} + + def _partial_only(*args, **kwargs): + calls["n"] += 1 + return {"b64": _b64_png(), "source": "partial"} + + monkeypatch.setattr(codex_plugin, "_collect_image_b64", _partial_only) + + result = provider.generate("a cat") + assert result["success"] is False + assert result["error_type"] == "incomplete_image" + assert "partial" in result["error"].lower() + # One initial attempt + one content-agnostic retry. + assert calls["n"] == codex_plugin._NONFINAL_RETRIES + 1 + + def test_empty_stream_retries_then_fails(self, provider, monkeypatch): + monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token") + calls = {"n": 0} + + def _empty(*args, **kwargs): + calls["n"] += 1 + return None + + monkeypatch.setattr(codex_plugin, "_collect_image_b64", _empty) + + result = provider.generate("a cat") + assert result["success"] is False + assert result["error_type"] == "empty_response" + assert calls["n"] == codex_plugin._NONFINAL_RETRIES + 1 + + def test_partial_then_final_on_retry_succeeds(self, provider, monkeypatch): + monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token") + calls = {"n": 0} + + def _then_final(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + return {"b64": _b64_png(), "source": "partial"} + return {"b64": _b64_png(), "source": "final"} + + monkeypatch.setattr(codex_plugin, "_collect_image_b64", _then_final) + + result = provider.generate("a cat") + assert result["success"] is True + assert result.get("image_source") == "final" + assert calls["n"] == 2 + + def test_empty_then_final_on_retry_succeeds(self, provider, monkeypatch): + monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token") + calls = {"n": 0} + + def _then_final(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + return None + return {"b64": _b64_png(), "source": "final"} + + monkeypatch.setattr(codex_plugin, "_collect_image_b64", _then_final) + + result = provider.generate("a cat") + assert result["success"] is True + assert result.get("image_source") == "final" + assert calls["n"] == 2 + def test_empty_response_returns_error(self, provider, monkeypatch): monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token") + monkeypatch.setattr(codex_plugin, "_NONFINAL_RETRIES", 0) monkeypatch.setattr(codex_plugin, "_collect_image_b64", lambda *a, **kw: None) result = provider.generate("a cat")