fix(agent): accept dimension-correct image shrinks (#48013)

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:

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