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:
Axmr1
2026-08-19 16:46:13 -07:00
committed by Teknium
parent d0132b5821
commit 4511ba49dd
2 changed files with 285 additions and 42 deletions

View File

@@ -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,
},
)

View File

@@ -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")