From 6f09e2481d570f4b9a70f16e27b28e3837d1063d Mon Sep 17 00:00:00 2001 From: embwl0x Date: Fri, 10 Jul 2026 06:21:06 -0400 Subject: [PATCH 1/2] fix(agent): persist image-shrink recovery --- agent/conversation_compression.py | 102 +++++++++++++----- agent/conversation_loop.py | 12 +++ run_agent.py | 2 + tests/run_agent/test_image_shrink_recovery.py | 62 +++++++++++ 4 files changed, 153 insertions(+), 25 deletions(-) diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index 843960cc2818..fa1e4da03359 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -1102,10 +1102,75 @@ def _compress_context_via_codex_app_server( 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). @@ -1121,7 +1186,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. @@ -1275,28 +1342,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) -> 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 - for msg in api_messages: if not isinstance(msg, dict): continue @@ -1309,10 +1354,12 @@ def _write_data_url_to_source(source: dict, data_url: str) -> None: 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): - _write_data_url_to_source(source, resized) + if replacements is not None and url: + replacements[url] = resized + _write_data_url_to_image_source(source, resized) changed_count += 1 elif unshrinkable: unshrinkable_oversized += 1 @@ -1326,6 +1373,8 @@ def _write_data_url_to_source(source: dict, data_url: str) -> None: url = image_value.get("url", "") resized, unshrinkable = _shrink_data_url(url) if resized: + if replacements is not None: + replacements[url] = resized image_value["url"] = resized changed_count += 1 elif unshrinkable: @@ -1333,6 +1382,8 @@ def _write_data_url_to_source(source: dict, data_url: str) -> None: elif isinstance(image_value, str): resized, unshrinkable = _shrink_data_url(image_value) if resized: + if replacements is not None: + replacements[image_value] = resized part["image_url"] = resized changed_count += 1 elif unshrinkable: @@ -1363,5 +1414,6 @@ def _write_data_url_to_source(source: dict, data_url: str) -> None: "check_compression_model_feasibility", "replay_compression_warning", "compress_context", + "apply_image_url_replacements_in_messages", "try_shrink_image_parts_in_messages", ] diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index f3b2bc0d39e3..704cd43bff00 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -2650,10 +2650,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...", diff --git a/run_agent.py b/run_agent.py index fe378f396ae9..da7bb338ade6 100644 --- a/run_agent.py +++ b/run_agent.py @@ -5142,12 +5142,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( diff --git a/tests/run_agent/test_image_shrink_recovery.py b/tests/run_agent/test_image_shrink_recovery.py index bdbb905d66e0..998b1700a8db 100644 --- a/tests/run_agent/test_image_shrink_recovery.py +++ b/tests/run_agent/test_image_shrink_recovery.py @@ -23,6 +23,7 @@ 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 @@ -227,6 +228,67 @@ 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_many_image_dimension_limit_rewritten(self, monkeypatch): """A 2000px many-image rejection must shrink images below the cap.""" agent = _make_agent() From 31ae7be9bd20720d09a7acd2ee81c730ea4bf80b Mon Sep 17 00:00:00 2001 From: embwl0x Date: Sat, 11 Jul 2026 07:19:35 -0500 Subject: [PATCH 2/2] test(agent): cover image-shrink retry persistence --- tests/run_agent/test_image_shrink_recovery.py | 113 +++++++++++++++++- 1 file changed, 109 insertions(+), 4 deletions(-) diff --git a/tests/run_agent/test_image_shrink_recovery.py b/tests/run_agent/test_image_shrink_recovery.py index 998b1700a8db..af72719a405b 100644 --- a/tests/run_agent/test_image_shrink_recovery.py +++ b/tests/run_agent/test_image_shrink_recovery.py @@ -9,17 +9,17 @@ 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 @@ -164,6 +164,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()