fix(image_gen/openai-codex): do not save progressive partial frames as finals
Codex Responses streams can emit partial_image_b64 previews without a final image_generation_call.result. The provider treated any b64 as success and could let a partial overwrite a coexisting final in the same payload, delivering smeared intermediates as finished GPT Image 2 outputs. Request partial_images=0, prefer final over partial in extraction, fail closed (with one content-agnostic retry) unless source=final, and surface image_source plus pixel_size for QA.
This commit is contained in:
@@ -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,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user