From e762a5a4737e12a99c29f640f2bc2c9a0a77f7ba Mon Sep 17 00:00:00 2001 From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com> Date: Tue, 28 Jul 2026 19:15:28 +0500 Subject: [PATCH] fix(compression): copy-on-write in image-shrink recovery so degraded images never reach stored history MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With the selective prompt-cache copy (#57046), un-marked messages on the decorated api_messages list share their nested content parts with the persistent conversation history — the per-message copy in conversation_loop is shallow and decoration now deep-copies only the marked messages. try_shrink_image_parts_in_messages previously wrote the re-encoded image INTO the aliased part/source dicts, so an image-too-large retry on an Anthropic route would silently replace the original image bytes in agent.messages (and persist the degraded copy). Replace the in-place writes with copy-on-write: rebuild the content list with fresh part/source/image_url dicts and reassign msg['content'] — a top-level write on the per-call copy that never reaches history. Adds two regression tests simulating the aliasing; both fail against the old in-place implementation (mutation-verified). --- agent/conversation_compression.py | 48 ++++++++++++--- tests/run_agent/test_image_shrink_recovery.py | 59 +++++++++++++++++++ 2 files changed, 99 insertions(+), 8 deletions(-) diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index 506181f2d390..8308b618da76 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -2721,16 +2721,28 @@ def try_shrink_image_parts_in_messages( media_type = "image/jpeg" return f"data:{media_type};base64,{data}" - def _write_data_url_to_source(source: dict, data_url: str) -> None: + def _write_data_url_to_source(source: dict, data_url: str) -> dict: + """Return a NEW source dict carrying the re-encoded payload. + + Copy-on-write: content parts on the per-call ``api_messages`` list may + be shared references into the persistent conversation history (the + per-message copy is shallow, and cache decoration only deep-copies the + marked messages). Mutating the existing dict would rewrite the stored + transcript with the degraded image — so the caller replaces the part, + never edits it in place. + """ 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 + return { + **source, + "type": "base64", + "media_type": media_type, + "data": data, + } for msg in api_messages: if not isinstance(msg, dict): @@ -2738,7 +2750,13 @@ def try_shrink_image_parts_in_messages( content = msg.get("content") if not isinstance(content, list): continue - for part in content: + # Copy-on-write per message: never mutate part/source dicts in place — + # they can alias the stored conversation history (see + # _write_data_url_to_source). Build a replacement content list on the + # first shrunken part and reassign msg["content"] (a top-level write on + # the per-call message copy, which never reaches history). + new_content: list | None = None + for part_idx, part in enumerate(content): if not isinstance(part, dict): continue ptype = part.get("type") @@ -2747,7 +2765,12 @@ def try_shrink_image_parts_in_messages( url = _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 new_content is None: + new_content = list(content) + new_content[part_idx] = { + **part, + "source": _write_data_url_to_source(source, resized), + } changed_count += 1 elif unshrinkable: unshrinkable_oversized += 1 @@ -2761,17 +2784,26 @@ def try_shrink_image_parts_in_messages( url = image_value.get("url", "") resized, unshrinkable = _shrink_data_url(url) if resized: - image_value["url"] = resized + if new_content is None: + new_content = list(content) + new_content[part_idx] = { + **part, + "image_url": {**image_value, "url": resized}, + } changed_count += 1 elif unshrinkable: unshrinkable_oversized += 1 elif isinstance(image_value, str): resized, unshrinkable = _shrink_data_url(image_value) if resized: - part["image_url"] = resized + if new_content is None: + new_content = list(content) + new_content[part_idx] = {**part, "image_url": resized} changed_count += 1 elif unshrinkable: unshrinkable_oversized += 1 + if new_content is not None: + msg["content"] = new_content if changed_count: logger.info( diff --git a/tests/run_agent/test_image_shrink_recovery.py b/tests/run_agent/test_image_shrink_recovery.py index bdbb905d66e0..522c79541ab0 100644 --- a/tests/run_agent/test_image_shrink_recovery.py +++ b/tests/run_agent/test_image_shrink_recovery.py @@ -661,3 +661,62 @@ class TestShrinkImagePartsHelper: # Default cap (8000) — no explicit max_dimension passed. assert agent._try_shrink_image_parts_in_messages(msgs) is True assert msgs[0]["content"][0]["image_url"]["url"] == shrunk + + +class TestShrinkCopyOnWriteHistoryIsolation: + """The shrink recovery must never rewrite the stored conversation history. + + With selective prompt-cache copying (#57046 salvage), un-marked messages + on the decorated per-call list share their nested content parts with + ``agent.messages``. The shrink helper therefore replaces parts + copy-on-write and reassigns ``msg["content"]`` instead of mutating the + aliased part/source dicts in place. + """ + + def test_shrink_does_not_mutate_aliased_history_parts(self, monkeypatch): + agent = _make_agent() + oversized_url = _big_png_data_url(5000) + shrunk = "data:image/jpeg;base64," + "C" * 1000 + + monkeypatch.setattr( + "tools.vision_tools._resize_image_for_vision", + lambda *a, **kw: shrunk, + raising=False, + ) + + # Simulate the persistent history and the per-call api_messages list: + # top-level dicts are shallow copies, nested content parts are ALIASED + # (exactly what conversation_loop's msg.copy() + selective cache + # decoration produce for un-marked messages). + history_part = {"type": "image_url", "image_url": {"url": oversized_url}} + history_msg = {"role": "user", "content": [history_part]} + api_msg = history_msg.copy() # shares the content list + part dicts + + assert agent._try_shrink_image_parts_in_messages([api_msg]) is True + # The outgoing copy carries the shrunken image... + assert api_msg["content"][0]["image_url"]["url"] == shrunk + # ...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 + + def test_shrink_anthropic_source_does_not_mutate_history(self, monkeypatch): + agent = _make_agent() + raw_b64 = _big_png_data_url(5000).split(",", 1)[1] + shrunk = "data:image/jpeg;base64," + "D" * 1000 + + monkeypatch.setattr( + "tools.vision_tools._resize_image_for_vision", + lambda *a, **kw: shrunk, + raising=False, + ) + + source = {"type": "base64", "media_type": "image/png", "data": raw_b64} + history_part = {"type": "image", "source": source} + history_msg = {"role": "user", "content": [history_part]} + api_msg = history_msg.copy() + + assert agent._try_shrink_image_parts_in_messages([api_msg]) is True + assert api_msg["content"][0]["source"]["data"] == "D" * 1000 + # Aliased history source untouched. + assert history_part["source"] is source + assert source["data"] == raw_b64