fix(agent): persist image-shrink recovery

agent/conversation_compression.py

@@ -2537,10 +2537,75 @@ def _complete_compaction_lifecycle() -> None:
     return messages, existing_prompt
 
 
+def _image_source_to_data_url(source: Any) -> Optional[str]:
+    if not isinstance(source, dict) or source.get("type") != "base64":
+        return None
+    data = source.get("data")
+    if not isinstance(data, str) or not data:
+        return None
+    media_type = str(source.get("media_type") or "image/jpeg").strip()
+    if not media_type.startswith("image/"):
+        media_type = "image/jpeg"
+    return f"data:{media_type};base64,{data}"
+
+
+def _write_data_url_to_image_source(source: dict, data_url: str) -> None:
+    header, _, data = data_url.partition(",")
+    media_type = "image/jpeg"
+    if header.startswith("data:"):
+        candidate = header[len("data:"):].split(";", 1)[0].strip()
+        if candidate.startswith("image/"):
+            media_type = candidate
+    source["type"] = "base64"
+    source["media_type"] = media_type
+    source["data"] = data
+
+
+def apply_image_url_replacements_in_messages(
+    messages: list,
+    replacements: dict[str, str],
+) -> int:
+    """Mirror repaired API image payloads into canonical session messages."""
+    if not messages or not replacements:
+        return 0
+
+    changed = 0
+    for msg in messages:
+        if not isinstance(msg, dict) or not isinstance(msg.get("content"), list):
+            continue
+        for part in msg["content"]:
+            if not isinstance(part, dict):
+                continue
+            if part.get("type") == "image":
+                source = part.get("source")
+                old_url = _image_source_to_data_url(source)
+                new_url = replacements.get(old_url or "")
+                if new_url and isinstance(source, dict):
+                    _write_data_url_to_image_source(source, new_url)
+                    changed += 1
+                continue
+            if part.get("type") not in {"image_url", "input_image"}:
+                continue
+            image_value = part.get("image_url")
+            if isinstance(image_value, dict):
+                old_url = image_value.get("url")
+                new_url = replacements.get(old_url) if isinstance(old_url, str) else None
+                if new_url:
+                    image_value["url"] = new_url
+                    changed += 1
+            elif isinstance(image_value, str):
+                new_url = replacements.get(image_value)
+                if new_url:
+                    part["image_url"] = new_url
+                    changed += 1
+    return changed
+
+
 def try_shrink_image_parts_in_messages(
     api_messages: list,
     *,
     max_dimension: int = 8000,
+    replacements: Optional[dict[str, str]] = None,
 ) -> bool:
     """Re-encode all native image parts at a smaller size to recover from
     image-too-large errors (Anthropic 5 MB, unknown other providers).
@@ -2556,7 +2621,9 @@ def try_shrink_image_parts_in_messages(
     under Anthropic's 5 MB ceiling with header overhead) or whose longest side
     exceeds ``max_dimension``, write the base64 to a tempfile, call
     ``vision_tools._resize_image_for_vision`` to produce a smaller data
-    URL, and substitute it in place.
+    URL, and substitute it in place. When ``replacements`` is supplied, it is
+    populated with each original-to-repaired data URL pair so the caller can
+    mirror the exact repair into canonical session history without re-encoding.
 
     Non-data-URL images (http/https URLs) are not touched — the provider
     fetches those itself and the size limit is different.
@@ -2710,17 +2777,6 @@ def _shrink_data_url(url: str) -> tuple:
             logger.warning("image-shrink recovery: re-encode failed — %s", exc)
             return None, triggered_by is not None
 
-    def _source_to_data_url(source: Any) -> Optional[str]:
-        if not isinstance(source, dict) or source.get("type") != "base64":
-            return None
-        data = source.get("data")
-        if not isinstance(data, str) or not data:
-            return None
-        media_type = str(source.get("media_type") or "image/jpeg").strip()
-        if not media_type.startswith("image/"):
-            media_type = "image/jpeg"
-        return f"data:{media_type};base64,{data}"
-
     def _write_data_url_to_source(source: dict, data_url: str) -> dict:
         """Return a NEW source dict carrying the re-encoded payload.
 
@@ -2762,9 +2818,11 @@ def _write_data_url_to_source(source: dict, data_url: str) -> dict:
             ptype = part.get("type")
             if ptype == "image":
                 source = part.get("source")
-                url = _source_to_data_url(source)
+                url = _image_source_to_data_url(source)
                 resized, unshrinkable = _shrink_data_url(url or "")
                 if resized and isinstance(source, dict):
+                    if replacements is not None and url:
+                        replacements[url] = resized
                     if new_content is None:
                         new_content = list(content)
                     new_content[part_idx] = {
@@ -2784,6 +2842,8 @@ def _write_data_url_to_source(source: dict, data_url: str) -> dict:
                 url = image_value.get("url", "")
                 resized, unshrinkable = _shrink_data_url(url)
                 if resized:
+                    if replacements is not None:
+                        replacements[url] = resized
                     if new_content is None:
                         new_content = list(content)
                     new_content[part_idx] = {
@@ -2796,6 +2856,8 @@ def _write_data_url_to_source(source: dict, data_url: str) -> dict:
             elif isinstance(image_value, str):
                 resized, unshrinkable = _shrink_data_url(image_value)
                 if resized:
+                    if replacements is not None:
+                        replacements[image_value] = resized
                     if new_content is None:
                         new_content = list(content)
                     new_content[part_idx] = {**part, "image_url": resized}
@@ -2831,5 +2893,6 @@ def _write_data_url_to_source(source: dict, data_url: str) -> dict:
     "check_compression_model_feasibility",
     "replay_compression_warning",
     "compress_context",
+    "apply_image_url_replacements_in_messages",
     "try_shrink_image_parts_in_messages",
 ]

agent/conversation_loop.py

@@ -3728,10 +3728,22 @@ def _perform_api_call(next_api_kwargs):
                 ):
                     _retry.image_shrink_retry_attempted = True
                     image_max_dimension = _image_error_max_dimension(api_error) or 8000
+                    image_replacements: dict[str, str] = {}
                     if agent._try_shrink_image_parts_in_messages(
                         api_messages,
                         max_dimension=image_max_dimension,
+                        replacements=image_replacements,
                     ):
+                        if image_replacements:
+                            from agent.conversation_compression import (
+                                apply_image_url_replacements_in_messages,
+                            )
+
+                            apply_image_url_replacements_in_messages(
+                                messages,
+                                image_replacements,
+                            )
+                            agent._persist_session(messages, conversation_history)
                         agent._vprint(
                             f"{agent.log_prefix}📐 Image(s) exceeded provider size limit — "
                             f"shrank and retrying...",

run_agent.py

@@ -6233,12 +6233,14 @@ def _try_shrink_image_parts_in_messages(
         api_messages: list,
         *,
         max_dimension: int = 8000,
+        replacements: Optional[Dict[str, str]] = None,
     ) -> bool:
         """Forwarder — see ``agent.conversation_compression.try_shrink_image_parts_in_messages``."""
         from agent.conversation_compression import try_shrink_image_parts_in_messages
         return try_shrink_image_parts_in_messages(
             api_messages,
             max_dimension=max_dimension,
+            replacements=replacements,
         )
 
     def _try_strip_image_parts_from_tool_messages(

tests/run_agent/test_image_shrink_recovery.py

@@ -9,20 +9,21 @@
      payload in-place, re-encoding native data: URL image parts to fit
      under 4 MB using vision_tools._resize_image_for_vision.
 
-The end-to-end wiring in the retry loop is not unit-tested here — it's
-covered by the live E2E in the PR description. These tests lock in the
-two pieces that matter independently: the classifier signal and the
-payload rewriter.
+The retry-loop regression uses detached canonical/API payloads to prove the
+repair is persisted before retry and the repaired image is not encoded twice.
 """
 
 from __future__ import annotations
 
 import base64
+import copy
 import sys
 from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
 
 
 from agent.conversation_loop import _image_error_max_dimension
+from agent.conversation_compression import apply_image_url_replacements_in_messages
 from agent.error_classifier import FailoverReason, classify_api_error
 
 
@@ -123,6 +124,111 @@ def _make_agent():
     return agent
 
 
+def _make_conversation_agent():
+    """Build an isolated agent capable of exercising the real retry loop."""
+    from run_agent import AIAgent
+
+    tool_defs = [{
+        "type": "function",
+        "function": {
+            "name": "web_search",
+            "description": "search",
+            "parameters": {"type": "object", "properties": {}},
+        },
+    }]
+    with (
+        patch("run_agent.get_tool_definitions", return_value=tool_defs),
+        patch("run_agent.check_toolset_requirements", return_value={}),
+        patch("run_agent.OpenAI"),
+    ):
+        agent = AIAgent(
+            model="anthropic/claude-sonnet-4.6",
+            provider="openrouter",
+            api_key="unused",
+            base_url="https://openrouter.ai/api/v1",
+            quiet_mode=True,
+            skip_context_files=True,
+            skip_memory=True,
+        )
+    agent.client = MagicMock()
+    return agent
+
+
+def _response(content: str):
+    message = SimpleNamespace(content=content, tool_calls=None)
+    choice = SimpleNamespace(message=message, finish_reason="stop")
+    return SimpleNamespace(choices=[choice], model="test/model", usage=None)
+
+
+def _first_image_url(messages) -> str | None:
+    for message in messages:
+        content = message.get("content")
+        if not isinstance(content, list):
+            continue
+        for part in content:
+            if not isinstance(part, dict) or part.get("type") != "image_url":
+                continue
+            image_url = part.get("image_url")
+            if isinstance(image_url, dict):
+                return image_url.get("url")
+    return None
+
+
+def test_retry_loop_persists_repair_before_reusing_image(monkeypatch):
+    agent = _make_conversation_agent()
+    oversized_url = _big_png_data_url(5000)
+    shrunk_url = "data:image/jpeg;base64," + "S" * 1000
+    resize_calls = []
+    events = []
+
+    def _resize(*args, **kwargs):
+        resize_calls.append((args, kwargs))
+        return shrunk_url
+
+    monkeypatch.setattr(
+        "tools.vision_tools._resize_image_for_vision",
+        _resize,
+        raising=False,
+    )
+
+    error = _FakeApiError(
+        400,
+        "messages.0.content.1.image.source.base64: image exceeds 5 MB maximum",
+    )
+    responses = [error, _response("recovered")]
+
+    def _api_call(api_kwargs):
+        events.append(("api", _first_image_url(api_kwargs["messages"])))
+        response = responses.pop(0)
+        if isinstance(response, Exception):
+            raise response
+        return response
+
+    def _persist(messages, _conversation_history):
+        events.append(("persist", _first_image_url(copy.deepcopy(messages))))
+
+    agent._interruptible_api_call = _api_call
+    agent._persist_session = _persist
+    agent._save_trajectory = lambda *args, **kwargs: None
+    # Prompt caching deep-copies nested content, reproducing the detached API
+    # payload that exposed the production persistence bug.
+    agent._use_prompt_caching = True
+
+    result = agent.run_conversation([{
+        "type": "image_url",
+        "image_url": {"url": oversized_url},
+    }])
+
+    assert result["completed"] is True
+    assert result["final_response"] == "recovered"
+    assert len(resize_calls) == 1
+    assert ("api", oversized_url) in events
+    first_api = events.index(("api", oversized_url))
+    repaired_persist = events.index(("persist", shrunk_url), first_api + 1)
+    repaired_retry = events.index(("api", shrunk_url), repaired_persist + 1)
+    assert first_api < repaired_persist < repaired_retry
+
+
 class TestShrinkImagePartsHelper:
     def test_no_messages_returns_false(self):
         agent = _make_agent()
@@ -180,6 +286,66 @@ def _fake_resize(path, mime_type=None, max_base64_bytes=None, max_dimension=None
         assert changed is True
         assert msgs[0]["content"][1]["image_url"]["url"] == shrunk
 
+    def test_records_replacement_for_canonical_session_history(self, monkeypatch):
+        """A successful API repair must expose the exact canonical rewrite."""
+        agent = _make_agent()
+        oversized_url = _big_png_data_url(5000)
+        shrunk = "data:image/jpeg;base64," + "A" * 1000
+        monkeypatch.setattr(
+            "tools.vision_tools._resize_image_for_vision",
+            lambda *args, **kwargs: shrunk,
+            raising=False,
+        )
+        api_messages = [{
+            "role": "user",
+            "content": [{
+                "type": "image_url",
+                "image_url": {"url": oversized_url},
+            }],
+        }]
+        replacements: dict[str, str] = {}
+
+        assert agent._try_shrink_image_parts_in_messages(
+            api_messages,
+            replacements=replacements,
+        ) is True
+        assert replacements == {oversized_url: shrunk}
+
+    def test_applies_replacements_to_detached_session_history(self):
+        old_openai = "data:image/png;base64,OPENAI"
+        old_anthropic = "data:image/png;base64,ANTHROPIC"
+        new_openai = "data:image/jpeg;base64,SMALL-OPENAI"
+        new_anthropic = "data:image/webp;base64,SMALL-ANTHROPIC"
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "image_url", "image_url": {"url": old_openai}},
+                {
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": "image/png",
+                        "data": "ANTHROPIC",
+                    },
+                },
+            ],
+        }]
+
+        changed = apply_image_url_replacements_in_messages(
+            messages,
+            {
+                old_openai: new_openai,
+                old_anthropic: new_anthropic,
+            },
+        )
+
+        assert changed == 2
+        assert messages[0]["content"][0]["image_url"]["url"] == new_openai
+        assert messages[0]["content"][1]["source"] == {
+            "type": "base64",
+            "media_type": "image/webp",
+            "data": "SMALL-ANTHROPIC",
+        }
 
     def test_anthropic_base64_image_source_rewritten(self, monkeypatch):
         """Anthropic-native image blocks are shrinkable after adapter conversion."""
@@ -527,4 +693,3 @@ def test_shrink_does_not_mutate_aliased_history_parts(self, monkeypatch):
         # ...but the stored history still has the original bytes.
         assert history_msg["content"][0] is history_part
         assert history_part["image_url"]["url"] == oversized_url
-