diff --git a/agent/context_compressor.py b/agent/context_compressor.py index 881e7cfe45..151c0856cd 100644 --- a/agent/context_compressor.py +++ b/agent/context_compressor.py @@ -1578,6 +1578,35 @@ def _tool_content_has_images(content: Any) -> bool: return _content_has_images(content) +def _strip_images_from_tool_msg(msg: Dict[str, Any]) -> Optional[Dict[str, Any]]: + """Return a copy of a tool message with its image payloads replaced. + + Handles the two image-bearing tool-result shapes: + + * ``{_multimodal: True, ...}`` envelopes collapse to a short + ``"[screenshot removed] "`` string; + * OpenAI-style part lists have image parts swapped for text + placeholders via :func:`_strip_image_parts_from_parts`. + + Returns ``None`` when the message carries no strippable image (the + caller should leave it untouched). The returned copy has its stale + ``api_content`` sidecar dropped so replay cannot resend the + pre-rewrite bytes. The input message is never mutated. + """ + content = msg.get("content") + if isinstance(content, dict) and content.get("_multimodal"): + summary = content.get("text_summary") or "[screenshot removed to save context]" + new_msg = {**msg, "content": f"[screenshot removed] {str(summary)[:200]}"} + drop_stale_api_content(new_msg) + return new_msg + stripped = _strip_image_parts_from_parts(content) + if stripped is None: + return None + new_msg = {**msg, "content": stripped} + drop_stale_api_content(new_msg) + return new_msg + + def _retire_stale_tool_result_images( result: List[Dict[str, Any]], keep_newest: int = _MAX_KEEP_TOOL_IMAGES, @@ -1598,24 +1627,14 @@ def _retire_stale_tool_result_images( msg = result[i] if not isinstance(msg, dict) or msg.get("role") != "tool": continue - content = msg.get("content") - if not _tool_content_has_images(content): + if not _tool_content_has_images(msg.get("content")): continue seen += 1 if seen <= keep_newest: continue - if isinstance(content, dict) and content.get("_multimodal"): - summary = content.get("text_summary") or "[screenshot removed to save context]" - new_msg = {**msg, "content": f"[screenshot removed] {summary[:200]}"} - drop_stale_api_content(new_msg) - result[i] = new_msg - pruned += 1 + new_msg = _strip_images_from_tool_msg(msg) + if new_msg is None: continue - stripped = _strip_image_parts_from_parts(content) - if stripped is None: - continue - new_msg = {**msg, "content": stripped} - drop_stale_api_content(new_msg) result[i] = new_msg pruned += 1 return pruned @@ -3761,16 +3780,15 @@ class ContextCompressor(ContextEngine): if msg.get("role") != "tool": return False content = msg.get("content", "") - if isinstance(content, list): - stripped = _strip_image_parts_from_parts(content) - if stripped is not None: - result[idx] = {**msg, "content": stripped} - pruned += 1 - return True - return False - if isinstance(content, dict) and content.get("_multimodal"): - summary = content.get("text_summary") or "[screenshot removed to save context]" - result[idx] = {**msg, "content": f"[screenshot removed] {summary[:200]}"} + if isinstance(content, list) or ( + isinstance(content, dict) and content.get("_multimodal") + ): + # Image-bearing shapes share one strip policy with pass 3.5 + # (also drops the stale api_content sidecar on rewrite). + new_msg = _strip_images_from_tool_msg(msg) + if new_msg is None: + return False + result[idx] = new_msg pruned += 1 return True if not isinstance(content, str): diff --git a/tests/agent/test_compressor_stale_tool_images.py b/tests/agent/test_compressor_stale_tool_images.py index a6972b397a..4ef0ddd836 100644 --- a/tests/agent/test_compressor_stale_tool_images.py +++ b/tests/agent/test_compressor_stale_tool_images.py @@ -138,3 +138,74 @@ class TestRetireStaleToolImagesInProtectedTail: ) assert _content_has_images(out[0]["content"]) assert out[0]["content"][1]["image_url"]["url"].endswith("USERUPLOAD") + + +class TestSharedImageStripHelper: + """One strip policy for pass 3.5 and the demote pass (#92783 follow-up).""" + + def test_demote_pass_drops_stale_api_content_on_image_strip(self): + """Pass 2's image demotion must drop the api_content sidecar. + + Before the shared _strip_images_from_tool_msg helper, only pass 3.5 + dropped the sidecar; the demote branches left it behind, letting + replay restore pre-strip bytes. + """ + from agent.context_compressor import _strip_images_from_tool_msg + + msg = { + "role": "tool", + "tool_call_id": "c1", + "api_content": "stale exact-wire copy with image bytes", + "content": [ + {"type": "text", "text": "shot"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64," + "A" * 400}, + }, + ], + } + new_msg = _strip_images_from_tool_msg(msg) + assert new_msg is not None + assert "api_content" not in new_msg + # Input untouched (copy-on-write). + assert "api_content" in msg + assert _content_has_images(msg["content"]) + assert not _content_has_images(new_msg["content"]) + + def test_envelope_collapses_to_summary_string(self): + from agent.context_compressor import _strip_images_from_tool_msg + + msg = { + "role": "tool", + "tool_call_id": "c2", + "api_content": "stale", + "content": { + "_multimodal": True, + "content": [ + {"type": "text", "text": "s"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,XYZ"}, + }, + ], + "text_summary": "native shot", + }, + } + new_msg = _strip_images_from_tool_msg(msg) + assert new_msg is not None + assert isinstance(new_msg["content"], str) + assert "screenshot removed" in new_msg["content"] + assert "native shot" in new_msg["content"] + assert "api_content" not in new_msg + + def test_imageless_content_returns_none(self): + from agent.context_compressor import _strip_images_from_tool_msg + + msg = {"role": "tool", "tool_call_id": "c3", "content": "plain text"} + assert _strip_images_from_tool_msg(msg) is None + msg2 = { + "role": "tool", + "tool_call_id": "c4", + "content": [{"type": "text", "text": "no images here"}], + } + assert _strip_images_from_tool_msg(msg2) is None