diff --git a/backend/packages/harness/deerflow/agents/middlewares/view_image_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/view_image_middleware.py index 97e95e46c2..060eeadc7c 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/view_image_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/view_image_middleware.py @@ -50,6 +50,29 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]): state_schema = ViewImageMiddlewareState + def __init__( + self, + *, + max_injected_images: int = 10, + max_injected_total_chars: int = 32_000_000, + ) -> None: + """Bound the per-request image context injected after view_image rounds. + + ``viewed_images`` only grows (#6825): every model call that follows a + ``view_image`` round re-encodes and re-sends every entry ever viewed, + so a long PDF-inspection thread pushes the request body past provider + limits (e.g. Anthropic's 32 MB) even though each single image is + capped. These knobs cap how many images and how much base64 payload + (in data-url characters) may be injected per model call; admission + runs newest-first so the most recently viewed images survive, and + over-budget entries keep their text bullet with a placeholder, so + nothing silently disappears. The 32M-character default keeps a single + max-size image injectable (20 MiB source -> ~27 MB base64) while + bounding accumulation below typical provider request limits. + """ + self._max_injected_images = max(1, int(max_injected_images)) + self._max_injected_total_chars = max(1, int(max_injected_total_chars)) + @staticmethod def _is_image_context_message(message: object) -> bool: """Return whether a message is trusted transient image context.""" @@ -312,7 +335,7 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]): *, host_path_allowed: Callable[[str, str], bool] | None = None, ) -> list[str | dict]: - """Create a formatted message with all viewed image details. + """Create a formatted message with the viewed image details that fit the injection budget. Reads image files on-demand from the active sandbox when available and encodes them as base64 for the model. The base64 data is NOT persisted in @@ -320,6 +343,11 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]): source sandbox id when applicable) is stored in ``viewed_images``, avoiding large duplicate payloads across every checkpoint (see #4138). + Content blocks are emitted oldest-first, but budget admission itself + runs newest-first (reversed insertion order) so the most recently + viewed images survive the caps (see #6825); an entry over budget keeps + its text bullet plus a placeholder instead of its base64 payload. + Args: state: Current state containing viewed_images @@ -331,6 +359,25 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]): # Return a properly formatted text block, not a plain string array return [{"type": "text", "text": "No images have been viewed."}] + # Budget admission runs newest-first so the most recently viewed + # images survive the caps, and the oldest entries over budget keep + # their text bullet with a placeholder in place (#6825). Admission is + # decided from the recorded byte size before any file is read: base64 + # expands 4/3 and the data-url prefix adds ~100 characters, so the + # estimate tracks the encoded payload closely. Omitted entries are + # never read from the sandbox, which also bounds the work this + # middleware performs per model call. + admitted: set[str] = set() + admitted_chars = 0 + admitted_images = 0 + for image_path, image_data in reversed(viewed_images.items()): + estimated_chars = int(image_data.get("size", 0) or 0) * 4 // 3 + 128 + if admitted_images >= self._max_injected_images or admitted_chars + estimated_chars > self._max_injected_total_chars: + continue + admitted_images += 1 + admitted_chars += estimated_chars + admitted.add(image_path) + # Build the message with image information content_blocks: list[str | dict] = [{"type": "text", "text": "Here are the images you've viewed:"}] @@ -345,6 +392,10 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]): # Add text description content_blocks.append({"type": "text", "text": f"\n- **{image_path}** ({mime_type})"}) + if image_path not in admitted: + content_blocks.append({"type": "text", "text": " (omitted: image injection budget reached; call view_image again to re-view this image)"}) + continue + # Read the image file on-demand and encode as base64 for the model data_url = self._read_image_as_data_url( state, diff --git a/backend/packages/harness/deerflow/agents/thread_state.py b/backend/packages/harness/deerflow/agents/thread_state.py index edf813d0d1..883723ec10 100644 --- a/backend/packages/harness/deerflow/agents/thread_state.py +++ b/backend/packages/harness/deerflow/agents/thread_state.py @@ -115,6 +115,12 @@ def merge_viewed_images(existing: dict[str, ViewedImageData] | None, new: dict[s Special case: If new is an empty dict {}, it clears the existing images. This allows middlewares to clear the viewed_images state after processing. + + Dict order tracks recency: a key updated by ``new`` is re-inserted at the + end, so a re-viewed image becomes the newest entry. The image injection + middleware admits entries newest-first and tells the model to re-view an + omitted image, so the recovery only works if the re-viewed entry actually + moves to the newest position (#6825, review on #6829). """ if existing is None: return new or {} @@ -123,8 +129,11 @@ def merge_viewed_images(existing: dict[str, ViewedImageData] | None, new: dict[s # Special case: empty dict means clear all viewed images if len(new) == 0: return {} - # Merge dictionaries, new values override existing ones for same keys - return {**existing, **new} + # Merge dictionaries; new values override existing ones for same keys, and + # updated keys move to the end so insertion order reflects recency. + merged = {key: value for key, value in existing.items() if key not in new} + merged.update(new) + return merged def merge_todos(existing: list | None, new: list | None) -> list | None: diff --git a/backend/tests/test_thread_state_reducers.py b/backend/tests/test_thread_state_reducers.py index 227309fbff..5699fadd08 100644 --- a/backend/tests/test_thread_state_reducers.py +++ b/backend/tests/test_thread_state_reducers.py @@ -129,6 +129,23 @@ class TestMergeViewedImages: merged = merge_viewed_images(existing, new) assert set(merged.keys()) == {"k1", "k2"} + def test_updated_key_moves_to_end(self): + """Re-viewing an existing path must make it the newest entry (#6829 review).""" + existing = { + "k1": {"mime_type": "image/png", "size": 1, "actual_path": "/a"}, + "k2": {"mime_type": "image/png", "size": 2, "actual_path": "/b"}, + } + new = {"k1": {"mime_type": "image/png", "size": 3, "actual_path": "/a"}} + merged = merge_viewed_images(existing, new) + assert list(merged.keys()) == ["k2", "k1"] + assert merged["k1"]["size"] == 3 + + def test_new_keys_append_after_untouched(self): + existing = {"k1": {"mime_type": "image/png", "size": 1, "actual_path": "/a"}} + new = {"k2": {"mime_type": "image/png", "size": 2, "actual_path": "/b"}} + merged = merge_viewed_images(existing, new) + assert list(merged.keys()) == ["k1", "k2"] + def test_empty_dict_clears(self): existing = {"k1": {"mime_type": "image/png", "size": 1, "actual_path": "/a"}} assert merge_viewed_images(existing, {}) == {} diff --git a/backend/tests/test_view_image_middleware.py b/backend/tests/test_view_image_middleware.py index e6e5c0657c..aa87aeb3c5 100644 --- a/backend/tests/test_view_image_middleware.py +++ b/backend/tests/test_view_image_middleware.py @@ -39,6 +39,7 @@ from deerflow.agents.middlewares.view_image_middleware import ( _IMAGE_CONTEXT_MESSAGE_MARKER_KEY, ViewImageMiddleware, ) +from deerflow.agents.thread_state import merge_viewed_images from deerflow.config.paths import Paths from deerflow.storage import BlobRef @@ -994,3 +995,132 @@ def test_custom_base_host_copy_rejects_unscoped_thread_data_root(tmp_path, monke } assert _image_urls(ViewImageMiddleware()._inject(request)) == [] + + +def _text_blocks(blocks: list) -> list[str]: + return [block["text"] for block in blocks if isinstance(block, dict) and block.get("type") == "text"] + + +def test_injection_budget_keeps_newest_images_and_omits_oldest(tmp_path): + """The image cap admits newest-first; omitted entries keep a placeholder bullet (see #6825).""" + viewed_images = {} + for i in range(4): + viewed_images[f"/mnt/user-data/outputs/img_{i}.png"] = _make_viewed_image(tmp_path, f"img_{i}.png", data=bytes([i]) + b"\x89PNG\r\n\x1a\n" + b"\x00" * 16) + state = {"viewed_images": viewed_images} + + mw = ViewImageMiddleware(max_injected_images=2) + blocks = mw._create_image_details_message(state, host_path_allowed=lambda _, actual: Path(actual).is_relative_to(tmp_path)) + + image_urls = [b["image_url"]["url"] for b in blocks if isinstance(b, dict) and b.get("type") == "image_url"] + assert len(image_urls) == 2 + + # The two newest entries (img_2, img_3) are injected; the two oldest keep a placeholder. + # Walk the raw block list so each bullet's successor block is checked in place. + for i, block in enumerate(blocks): + if not (isinstance(block, dict) and block.get("type") == "text"): + continue + text = block["text"] + for omitted in ("img_0.png", "img_1.png"): + if omitted in text: + assert "omitted: image injection budget reached" in blocks[i + 1]["text"] + for injected in ("img_2.png", "img_3.png"): + if injected in text: + assert blocks[i + 1].get("type") == "image_url" + + # Distinct payloads prove which entries made it. + assert base64.b64decode(image_urls[0].split(",", 1)[1]).startswith(bytes([2])) + assert base64.b64decode(image_urls[1].split(",", 1)[1]).startswith(bytes([3])) + + +def test_injection_byte_budget_drops_oldest_when_total_exceeded(tmp_path): + """The byte cap evicts the oldest oversized entry while a newer small one still fits.""" + big = b"\x89PNG\r\n\x1a\n" + b"\x00" * (900_000 - 8) + small = b"\x89PNG\r\n\x1a\n" + b"\x00" * (100_000 - 8) + state = { + "viewed_images": { + "/mnt/user-data/outputs/big_0.png": _make_viewed_image(tmp_path, "big_0.png", data=big), + "/mnt/user-data/outputs/big_1.png": _make_viewed_image(tmp_path, "big_1.png", data=big), + "/mnt/user-data/outputs/small_2.png": _make_viewed_image(tmp_path, "small_2.png", data=small), + } + } + + mw = ViewImageMiddleware(max_injected_total_chars=2_000_000) + blocks = mw._create_image_details_message(state, host_path_allowed=lambda _, actual: Path(actual).is_relative_to(tmp_path)) + + image_urls = [b["image_url"]["url"] for b in blocks if isinstance(b, dict) and b.get("type") == "image_url"] + assert len(image_urls) == 2 + for i, block in enumerate(blocks): + if not (isinstance(block, dict) and block.get("type") == "text"): + continue + if "big_0.png" in block["text"]: + assert "omitted: image injection budget reached" in blocks[i + 1]["text"] + if "small_2.png" in block["text"]: + assert blocks[i + 1].get("type") == "image_url" + + +def test_injection_budget_defaults_keep_small_threads_unbounded(tmp_path): + """Defaults must not change threads whose image count and sizes stay small.""" + viewed_images = {} + for i in range(3): + viewed_images[f"/mnt/user-data/outputs/img_{i}.png"] = _make_viewed_image(tmp_path, f"img_{i}.png") + state = {"viewed_images": viewed_images} + + blocks = ViewImageMiddleware()._create_image_details_message(state, host_path_allowed=lambda _, actual: Path(actual).is_relative_to(tmp_path)) + + assert len([b for b in blocks if isinstance(b, dict) and b.get("type") == "image_url"]) == 3 + assert not any("omitted" in t for t in _text_blocks(blocks)) + + +def test_injection_budget_skips_reads_for_omitted_entries(tmp_path): + """Entries over budget are never read from disk or sandbox.""" + viewed_images = {} + for i in range(3): + viewed_images[f"/mnt/user-data/outputs/img_{i}.png"] = _make_viewed_image(tmp_path, f"img_{i}.png", data=bytes([i]) + b"\x89PNG\r\n\x1a\n" + b"\x00" * 16) + state = {"viewed_images": viewed_images} + + mw = ViewImageMiddleware(max_injected_images=1) + calls = [] + + def _fake_read(*args, **kwargs): + calls.append(args) + return "data:image/png;base64,QUJD" + + mw._read_image_as_data_url = _fake_read + blocks = mw._create_image_details_message(state) + + assert len([b for b in blocks if isinstance(b, dict) and b.get("type") == "image_url"]) == 1 + assert len(calls) == 1 + + +def test_reviewing_omitted_image_rescues_it_from_budget(tmp_path): + """Re-viewing an omitted image must move it to the newest slot so the + placeholder's recovery advice actually works (see #6829 review).""" + viewed_images = {} + for i in range(4): + viewed_images[f"/mnt/user-data/outputs/img_{i}.png"] = _make_viewed_image(tmp_path, f"img_{i}.png", data=bytes([i]) + b"\x89PNG\r\n\x1a\n" + b"\x00" * 16) + + mw = ViewImageMiddleware(max_injected_images=2) + + def host_ok(_, actual): + return Path(actual).is_relative_to(tmp_path) + + blocks = mw._create_image_details_message({"viewed_images": viewed_images}, host_path_allowed=host_ok) + injected = [b for b in blocks if isinstance(b, dict) and b.get("type") == "image_url"] + assert len(injected) == 2 # img_2, img_3 admitted; img_0/img_1 omitted + + # The model follows the placeholder advice and re-views img_0; the reducer + # re-inserts the updated key at the end, making it the newest entry. + updated = merge_viewed_images( + viewed_images, + {"/mnt/user-data/outputs/img_0.png": _make_viewed_image(tmp_path, "img_0.png", data=bytes([0]) + b"\x89PNG\r\n\x1a\n" + b"\x00" * 16)}, + ) + assert list(updated.keys())[-1] == "/mnt/user-data/outputs/img_0.png" + + blocks = mw._create_image_details_message({"viewed_images": updated}, host_path_allowed=host_ok) + urls = [b["image_url"]["url"] for b in blocks if isinstance(b, dict) and b.get("type") == "image_url"] + assert len(urls) == 2 + # Render order is oldest-first, so the admitted slots are img_3 then the + # re-viewed img_0; img_1/img_2 stay omitted. + payloads = [base64.b64decode(url.split(",", 1)[1]) for url in urls] + assert payloads[0].startswith(bytes([3])) + assert payloads[1].startswith(bytes([0]))