From 7ed99f04e626c6da3152313830be42dd878359da Mon Sep 17 00:00:00 2001 From: Tranquil-Flow <66773372+Tranquil-Flow@users.noreply.github.com> Date: Wed, 17 Jun 2026 22:32:52 +0200 Subject: [PATCH] fix(agent): accept dimension-correct image shrinks (#48013) --- agent/conversation_compression.py | 104 +++++++++++------ tests/run_agent/test_image_shrink_recovery.py | 105 +++++++++++++++++- 2 files changed, 175 insertions(+), 34 deletions(-) diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index 318e67d0faf2d..15d0a59a287ec 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -702,34 +702,58 @@ def try_shrink_image_parts_in_messages( # actually brought under the target. unshrinkable_oversized = 0 - def _shrink_data_url(url: str) -> Optional[str]: - """Return a smaller data URL, or None if shrink can't help.""" + def _shrink_data_url(url: str) -> tuple[Optional[str], bool]: + """Return ``(resized_url, unshrinkable)`` for a data URL. + + ``unshrinkable`` is True only when the image exceeded a constraint + (byte-size or dimensions) and resizing failed to satisfy that same + constraint. This prevents a doomed retry if one dimension-oversized + image remains over the provider cap while a different image was + successfully rewritten. + """ if not isinstance(url, str) or not url.startswith("data:"): - return None + return None, False - # Check both byte size AND pixel dimensions. - needs_shrink = len(url) > target_bytes # over byte budget - if not needs_shrink: - # Even if bytes are fine, check pixel dimensions against the - # provider's reported per-side cap. A screenshot can be tiny in - # bytes yet too large in pixels. + def _decode_pixels(data_url: str) -> Optional[tuple[int, int]]: + """Return (width, height) of a base64 data URL, or None on any failure. + + Soft-depends on Pillow; returns None (caller falls back to + bytes-only check) if Pillow is missing or the data is corrupt. + """ try: - import base64 as _b64_dim - header_d, _, data_d = url.partition(",") - if not data_d: + import base64 as _b64 + import io as _io + header, _, data = data_url.partition(",") + if not data or not data_url.startswith("data:"): return None - raw_d = _b64_dim.b64decode(data_d) from PIL import Image as _PILImage - import io as _io_dim - with _PILImage.open(_io_dim.BytesIO(raw_d)) as _img: - if max(_img.size) <= max_dimension: - return None # both bytes and pixels are fine - needs_shrink = True # pixels exceed limit, force shrink + with _PILImage.open(_io.BytesIO(_b64.b64decode(data))) as _img: + return _img.size except Exception: - # If we can't check dimensions (Pillow unavailable, corrupt - # image, etc.), fall back to byte-only check. return None + # Check both byte size AND pixel dimensions. Track which constraint + # triggered the shrink — the accept/reject gate must use the same + # constraint, not blindly compare bytes. A PNG screenshot that + # dimensionally shrinks from 4000x3000 to 2000x1500 can re-encode to + # MORE bytes (PNG compression is non-monotonic in image size for + # smooth regions), and rejecting those results permanently wedges + # sessions on the Anthropic many-image 2000px path (#48013). + needs_shrink = len(url) > target_bytes # over byte budget + triggered_by = "bytes" if needs_shrink else None + if not needs_shrink: + # Even if bytes are fine, check pixel dimensions against the + # provider's reported per-side cap. A screenshot can be tiny in + # bytes yet too large in pixels. + dims = _decode_pixels(url) + if dims is None: + # Pillow missing or corrupt data — fall back to byte-only. + return None, False + if max(dims) <= max_dimension: + return None, False # both bytes and pixels are fine + needs_shrink = True + triggered_by = "dimension" + try: header, _, data = url.partition(",") mime = "image/jpeg" @@ -760,13 +784,33 @@ def _shrink_data_url(url: str) -> Optional[str]: Path(tmp.name).unlink(missing_ok=True) except Exception: pass - if not resized or len(resized) >= len(url): - # Shrink didn't help (or made it bigger — corrupt input?). - return None - return resized + if not resized: + # Resize returned nothing — Pillow couldn't help. + return None, True + if triggered_by == "bytes": + # Byte budget was the binding constraint — bytes must shrink. + if len(resized) >= len(url): + return None, True # re-encode made it bigger + return resized, False + # triggered_by == "dimension": dimension cap is the binding + # constraint. The re-encode may have grown in bytes (PNG + # screenshots re-encode larger when smaller); accept the result + # if it's now within the per-side cap. If we can't check + # dimensions (Pillow missing or corrupt re-encode), fall back to + # the historical "bytes must shrink" gate to avoid regressions. + new_dims = _decode_pixels(resized) + if new_dims is not None: + if max(new_dims) <= max_dimension: + return resized, False + # Still too tall/wide — re-encode didn't help. + return None, True + # Couldn't verify dimension — only accept if bytes also shrank. + if len(resized) >= len(url): + return None, True + return resized, False except Exception as exc: logger.warning("image-shrink recovery: re-encode failed — %s", exc) - return None + return None, triggered_by is not None for msg in api_messages: if not isinstance(msg, dict): @@ -785,20 +829,18 @@ def _shrink_data_url(url: str) -> Optional[str]: # OpenAI Responses: {"image_url": "data:..."} if isinstance(image_value, dict): url = image_value.get("url", "") - resized = _shrink_data_url(url) + resized, unshrinkable = _shrink_data_url(url) if resized: image_value["url"] = resized changed_count += 1 - elif isinstance(url, str) and url.startswith("data:") \ - and len(url) > target_bytes: + elif unshrinkable: unshrinkable_oversized += 1 elif isinstance(image_value, str): - resized = _shrink_data_url(image_value) + resized, unshrinkable = _shrink_data_url(image_value) if resized: part["image_url"] = resized changed_count += 1 - elif image_value.startswith("data:") \ - and len(image_value) > target_bytes: + elif unshrinkable: unshrinkable_oversized += 1 if changed_count: diff --git a/tests/run_agent/test_image_shrink_recovery.py b/tests/run_agent/test_image_shrink_recovery.py index 240546ea14cad..754a9edaa079d 100644 --- a/tests/run_agent/test_image_shrink_recovery.py +++ b/tests/run_agent/test_image_shrink_recovery.py @@ -108,11 +108,21 @@ def _big_png_data_url(size_kb: int) -> str: return "data:image/png;base64," + base64.b64encode(raw).decode("ascii") -def _install_fake_pillow(monkeypatch, size: tuple[int, int]) -> None: +def _install_fake_pillow( + monkeypatch, + size: tuple[int, int], + *, + shrunk_size: tuple[int, int] | None = None, + sizes: list[tuple[int, int]] | None = None, +) -> None: """Install the tiny subset of Pillow used by the shrink preflight.""" + call_count = {"n": 0} + target_sizes = sizes or [size, shrunk_size if shrunk_size is not None else size] + class _FakeImage: def __init__(self): - self.size = size + self.size = target_sizes[min(call_count["n"], len(target_sizes) - 1)] + call_count["n"] += 1 def __enter__(self): return self @@ -205,7 +215,7 @@ def _fake_resize(path, mime_type=None, max_base64_bytes=None, max_dimension=None def test_many_image_dimension_limit_rewritten(self, monkeypatch): """A 2000px many-image rejection must shrink images below 8000px.""" agent = _make_agent() - _install_fake_pillow(monkeypatch, (2501, 100)) + _install_fake_pillow(monkeypatch, (2501, 100), shrunk_size=(1500, 60)) oversized_for_many = _big_png_data_url(100) shrunk = "data:image/jpeg;base64," + "M" * 1000 seen = {} @@ -392,3 +402,92 @@ def fake_resize(path, *a, **kw): assert msgs[0]["content"][0]["image_url"]["url"] == small # The unshrinkable one is left as-is (caller surfaces original error). assert msgs[0]["content"][1]["image_url"]["url"] == unshrinkable + + def test_dimension_shrink_with_byte_growth_accepted(self, monkeypatch): + """Dimension-driven shrink may accept a byte-larger PNG re-encode.""" + agent = _make_agent() + _install_fake_pillow(monkeypatch, (2501, 100), shrunk_size=(1500, 60)) + original_url = _big_png_data_url(100) + dimensionally_shrunk = "data:image/png;base64," + "G" * 200 * 1024 + seen = {} + + def _fake_resize(path, mime_type=None, max_base64_bytes=None, max_dimension=None): + seen["max_dimension"] = max_dimension + return dimensionally_shrunk + + monkeypatch.setattr( + "tools.vision_tools._resize_image_for_vision", + _fake_resize, + raising=False, + ) + + msgs = [{ + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": original_url}}, + ], + }] + assert agent._try_shrink_image_parts_in_messages( + msgs, max_dimension=2000, + ) is True + assert seen["max_dimension"] == 2000 + assert msgs[0]["content"][0]["image_url"]["url"] == dimensionally_shrunk + + def test_dimension_shrink_failure_still_blocks_retry(self, monkeypatch): + """A dimension-oversized image that remains oversized is unshrinkable.""" + agent = _make_agent() + _install_fake_pillow(monkeypatch, (2501, 100)) + original_url = _big_png_data_url(100) + still_oversized = "data:image/png;base64," + "H" * 120 * 1024 + + monkeypatch.setattr( + "tools.vision_tools._resize_image_for_vision", + lambda *a, **kw: still_oversized, + raising=False, + ) + + msgs = [{ + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": original_url}}, + ], + }] + assert agent._try_shrink_image_parts_in_messages( + msgs, max_dimension=2000, + ) is False + assert msgs[0]["content"][0]["image_url"]["url"] == original_url + + def test_mixed_dimension_failure_returns_false(self, monkeypatch): + """Partial dimension-path progress must not burn the one retry.""" + agent = _make_agent() + _install_fake_pillow( + monkeypatch, + (2501, 100), + sizes=[(2501, 100), (1500, 60), (2501, 100), (2501, 100)], + ) + first = _big_png_data_url(100) + second = _big_png_data_url(90) + calls = {"n": 0} + + def _fake_resize(path, mime_type=None, max_base64_bytes=None, max_dimension=None): + calls["n"] += 1 + if calls["n"] == 1: + return "data:image/png;base64," + "G" * 200 * 1024 + return "data:image/png;base64," + "H" * 120 * 1024 + + monkeypatch.setattr( + "tools.vision_tools._resize_image_for_vision", + _fake_resize, + raising=False, + ) + + msgs = [{ + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": first}}, + {"type": "image_url", "image_url": {"url": second}}, + ], + }] + assert agent._try_shrink_image_parts_in_messages( + msgs, max_dimension=2000, + ) is False