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 7aa1e6a0f..7ad567a2a 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/view_image_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/view_image_middleware.py @@ -1,6 +1,9 @@ """Middleware for injecting image details into conversation before LLM call.""" +import asyncio +import base64 import logging +from pathlib import Path from typing import override from langchain.agents.middleware import AgentMiddleware @@ -11,6 +14,11 @@ from deerflow.agents.thread_state import ThreadState logger = logging.getLogger(__name__) +# Mirror the tool-side size cap as a defense-in-depth check. The tool +# enforces this at write time; the middleware re-checks at read time in +# case the file grew on disk between view and injection. +_MAX_IMAGE_BYTES = 20 * 1024 * 1024 + class ViewImageMiddlewareState(ThreadState): """Reuse the thread state so reducer-backed keys keep their annotations.""" @@ -91,9 +99,43 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]): # Check if all tool calls have been completed return tool_call_ids.issubset(completed_tool_ids) + @staticmethod + def _read_image_as_data_url(actual_path: str, mime_type: str, expected_size: int) -> str | None: + """Read image file and return a `data:` URL, or None on failure. + + Trust assumption: ``actual_path`` is set by ``view_image_tool`` + (server-side, validated against the allowed virtual roots at write + time) and held in LangGraph-controlled state. Client input cannot + reach this field, so the read scope is trusted. We still re-check + size at read time to defend against TOCTOU growth and skip files + exceeding ``_MAX_IMAGE_BYTES``. + """ + try: + file_path = Path(actual_path) + if not file_path.exists() or not file_path.is_file(): + return None + current_size = file_path.stat().st_size + if current_size != expected_size: + # File changed between view and inject - skip. + return None + if current_size > _MAX_IMAGE_BYTES: + return None + with open(file_path, "rb") as f: + image_bytes = f.read() + base64_data = base64.b64encode(image_bytes).decode("utf-8") + return f"data:{mime_type};base64,{base64_data}" + except OSError: + return None + def _create_image_details_message(self, state: ViewImageMiddlewareState) -> list[str | dict]: """Create a formatted message with all viewed image details. + Reads image files from disk on-demand and encodes them as base64 + for the model. The base64 data is NOT persisted in state -- only + lightweight metadata (path, mime_type, size) is stored in + ``viewed_images``, avoiding large duplicate payloads across every + checkpoint (see #4138). + Args: state: Current state containing viewed_images @@ -110,19 +152,24 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]): for image_path, image_data in viewed_images.items(): mime_type = image_data.get("mime_type", "unknown") - base64_data = image_data.get("base64", "") + actual_path = image_data.get("actual_path", "") + expected_size = image_data.get("size", 0) # Add text description content_blocks.append({"type": "text", "text": f"\n- **{image_path}** ({mime_type})"}) - # Add the actual image data so LLM can "see" it - if base64_data: - content_blocks.append( - { - "type": "image_url", - "image_url": {"url": f"data:{mime_type};base64,{base64_data}"}, - } - ) + # Read the image file on-demand and encode as base64 for the model + if actual_path: + data_url = self._read_image_as_data_url(actual_path, mime_type, expected_size) + if data_url: + content_blocks.append( + { + "type": "image_url", + "image_url": {"url": data_url}, + } + ) + else: + content_blocks.append({"type": "text", "text": f" (file unavailable or changed on disk: {actual_path})"}) return content_blocks @@ -221,4 +268,11 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]): Returns: State update with additional human message, or None if no update needed """ - return self._inject_image_message(state) + if not self._should_inject_image_message(state): + return None + # Image reads + base64 encoding can be slow (up to 20MB), so offload + # the blocking work to a thread rather than stalling the event loop. + image_content = await asyncio.to_thread(self._create_image_details_message, state) + human_msg = HumanMessage(content=image_content, additional_kwargs={"hide_from_ui": True}) + logger.debug("Injecting image details message with images before LLM call") + return {"messages": [human_msg]} diff --git a/backend/packages/harness/deerflow/agents/thread_state.py b/backend/packages/harness/deerflow/agents/thread_state.py index 908906cba..0b10cd5d3 100644 --- a/backend/packages/harness/deerflow/agents/thread_state.py +++ b/backend/packages/harness/deerflow/agents/thread_state.py @@ -18,8 +18,17 @@ class ThreadDataState(TypedDict): class ViewedImageData(TypedDict): - base64: str + """Metadata for a viewed image file. + + Only lightweight metadata is persisted in checkpoint state; the actual + image bytes are read on-demand from disk when the model needs them. + This avoids duplicating large base64 payloads across every checkpoint + (see #4138). + """ + mime_type: str + size: int + actual_path: str def merge_sandbox(existing: SandboxState | None, new: SandboxState | None) -> SandboxState | None: @@ -235,7 +244,7 @@ class ThreadState(AgentState): todos: Annotated[list | None, merge_todos] goal: Annotated[GoalState | None, merge_goal] uploaded_files: NotRequired[list[dict] | None] - viewed_images: Annotated[dict[str, ViewedImageData], merge_viewed_images] # image_path -> {base64, mime_type} + viewed_images: Annotated[dict[str, ViewedImageData], merge_viewed_images] # image_path -> metadata (no base64) promoted: Annotated[PromotedTools | None, merge_promoted] delegations: Annotated[list[DelegationEntry], merge_delegations] skill_context: Annotated[list[SkillEntry], merge_skill_context] diff --git a/backend/packages/harness/deerflow/tools/builtins/view_image_tool.py b/backend/packages/harness/deerflow/tools/builtins/view_image_tool.py index 3895cfd12..10e22f4f7 100644 --- a/backend/packages/harness/deerflow/tools/builtins/view_image_tool.py +++ b/backend/packages/harness/deerflow/tools/builtins/view_image_tool.py @@ -1,4 +1,3 @@ -import base64 import mimetypes from pathlib import Path from typing import Annotated @@ -132,7 +131,7 @@ def view_image_tool( update={"messages": [ToolMessage(f"Error: Image file is too large: {image_size} bytes. Maximum supported size is {_MAX_IMAGE_BYTES} bytes", tool_call_id=tool_call_id)]}, ) - # Read image file and convert to base64 + # Read image file to validate contents (magic bytes + size) try: with open(actual_path, "rb") as f: image_data = f.read() @@ -141,6 +140,12 @@ def view_image_tool( update={"messages": [ToolMessage(f"Error reading image file: {_sanitize_image_error(e, thread_data)}", tool_call_id=tool_call_id)]}, ) + if len(image_data) != image_size: + # File changed between stat() and read() - reject for safety. + return Command( + update={"messages": [ToolMessage("Error: Image file changed during read", tool_call_id=tool_call_id)]}, + ) + detected_mime_type = _detect_image_mime(image_data) if detected_mime_type is None: return Command( @@ -151,11 +156,17 @@ def view_image_tool( update={"messages": [ToolMessage(f"Error: Image contents are {detected_mime_type}, but file extension indicates {expected_mime_type}", tool_call_id=tool_call_id)]}, ) mime_type = detected_mime_type - image_base64 = base64.b64encode(image_data).decode("utf-8") - # Update viewed_images in state - # The merge_viewed_images reducer will handle merging with existing images - new_viewed_images = {image_path: {"base64": image_base64, "mime_type": mime_type}} + # Store only lightweight metadata in state (not base64) to avoid + # duplicating large payloads across every checkpoint (see #4138). + # The middleware reads the file on-demand when the model needs it. + new_viewed_images = { + image_path: { + "mime_type": mime_type, + "size": image_size, + "actual_path": str(actual_path), + } + } return Command( update={"viewed_images": new_viewed_images, "messages": [ToolMessage("Successfully read image", tool_call_id=tool_call_id)]}, diff --git a/backend/tests/test_sandbox_middleware.py b/backend/tests/test_sandbox_middleware.py index b34449ed0..53f61a080 100644 --- a/backend/tests/test_sandbox_middleware.py +++ b/backend/tests/test_sandbox_middleware.py @@ -336,7 +336,7 @@ def test_wrap_tool_call_merges_with_existing_command_update() -> None: return Command( update={ "messages": [tool_msg], - "viewed_images": {"a.png": {"base64": "x", "mime_type": "image/png"}}, + "viewed_images": {"a.png": {"mime_type": "image/png", "size": 1, "actual_path": "/tmp/a.png"}}, }, goto="next-node", ) @@ -347,7 +347,7 @@ def test_wrap_tool_call_merges_with_existing_command_update() -> None: assert result.goto == "next-node" assert isinstance(result.update, dict) assert result.update["messages"] == [tool_msg] - assert result.update["viewed_images"] == {"a.png": {"base64": "x", "mime_type": "image/png"}} + assert result.update["viewed_images"] == {"a.png": {"mime_type": "image/png", "size": 1, "actual_path": "/tmp/a.png"}} assert result.update["sandbox"] == {"sandbox_id": "new-sandbox"} diff --git a/backend/tests/test_thread_state_reducers.py b/backend/tests/test_thread_state_reducers.py index f89d66c93..864fb690a 100644 --- a/backend/tests/test_thread_state_reducers.py +++ b/backend/tests/test_thread_state_reducers.py @@ -122,13 +122,13 @@ class TestMergeViewedImages: """Sanity check for the existing viewed_images reducer.""" def test_merges_dicts(self): - existing = {"k1": {"base64": "x", "mime_type": "image/png"}} - new = {"k2": {"base64": "y", "mime_type": "image/jpeg"}} + existing = {"k1": {"mime_type": "image/png", "size": 1, "actual_path": "/a"}} + new = {"k2": {"mime_type": "image/jpeg", "size": 2, "actual_path": "/b"}} merged = merge_viewed_images(existing, new) assert set(merged.keys()) == {"k1", "k2"} def test_empty_dict_clears(self): - existing = {"k1": {"base64": "x", "mime_type": "image/png"}} + 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 f899dd5ec..1a8503e32 100644 --- a/backend/tests/test_view_image_middleware.py +++ b/backend/tests/test_view_image_middleware.py @@ -9,7 +9,8 @@ Covered behavior: - `_get_last_assistant_message` returns the most recent AIMessage (or None). - `_has_view_image_tool` only matches assistant messages with `view_image` tool calls. - `_all_tools_completed` verifies every tool call id has a matching ToolMessage. -- `_create_image_details_message` produces correctly structured content blocks. +- `_create_image_details_message` produces correctly structured content blocks, + reading image files on-demand from disk (no base64 stored in state). - `_should_inject_image_message` gates injection on all preconditions, including deduplication when an image-details message was already added. - `_inject_image_message` returns a state update with a HumanMessage, or None @@ -17,6 +18,7 @@ Covered behavior: - `before_model` and `abefore_model` expose the same behavior sync/async. """ +from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock @@ -40,6 +42,17 @@ def _runtime() -> MagicMock: return MagicMock() +def _make_viewed_image(tmp_path, filename="img.png", mime_type="image/png", data=b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR"): + """Create a real image file and return viewed_images metadata dict.""" + img_path = tmp_path / filename + img_path.write_bytes(data) + return { + "mime_type": mime_type, + "size": len(data), + "actual_path": str(img_path), + } + + class TestGetLastAssistantMessage: def test_returns_none_on_empty_list(self): mw = ViewImageMiddleware() @@ -162,11 +175,12 @@ class TestCreateImageDetailsMessage: blocks = mw._create_image_details_message({}) assert blocks == [{"type": "text", "text": "No images have been viewed."}] - def test_builds_blocks_for_single_image(self): + def test_builds_blocks_for_single_image(self, tmp_path): mw = ViewImageMiddleware() + img_meta = _make_viewed_image(tmp_path, "cat.png") state = { "viewed_images": { - "/path/to/cat.png": {"base64": "BASE64DATA", "mime_type": "image/png"}, + "/path/to/cat.png": img_meta, } } blocks = mw._create_image_details_message(state) @@ -177,17 +191,17 @@ class TestCreateImageDetailsMessage: assert blocks[1]["type"] == "text" assert "/path/to/cat.png" in blocks[1]["text"] assert "image/png" in blocks[1]["text"] - assert blocks[2] == { - "type": "image_url", - "image_url": {"url": "data:image/png;base64,BASE64DATA"}, - } + assert blocks[2]["type"] == "image_url" + assert blocks[2]["image_url"]["url"].startswith("data:image/png;base64,") - def test_builds_blocks_for_multiple_images(self): + def test_builds_blocks_for_multiple_images(self, tmp_path): mw = ViewImageMiddleware() + img1 = _make_viewed_image(tmp_path, "a.png", data=b"\x89PNG\r\n\x1a\nfake-png") + img2 = _make_viewed_image(tmp_path, "b.jpg", mime_type="image/jpeg", data=b"\xff\xd8\xff\xe0fake-jpeg") state = { "viewed_images": { - "/a.png": {"base64": "AAA", "mime_type": "image/png"}, - "/b.jpg": {"base64": "BBB", "mime_type": "image/jpeg"}, + "/a.png": img1, + "/b.jpg": img2, } } blocks = mw._create_image_details_message(state) @@ -197,26 +211,31 @@ class TestCreateImageDetailsMessage: image_url_blocks = [b for b in blocks if isinstance(b, dict) and b.get("type") == "image_url"] assert len(image_url_blocks) == 2 urls = {b["image_url"]["url"] for b in image_url_blocks} - assert "data:image/png;base64,AAA" in urls - assert "data:image/jpeg;base64,BBB" in urls + assert any(u.startswith("data:image/png;base64,") for u in urls) + assert any(u.startswith("data:image/jpeg;base64,") for u in urls) - def test_omits_image_url_block_when_base64_missing(self): + def test_omits_image_url_block_when_file_missing(self, tmp_path): mw = ViewImageMiddleware() state = { "viewed_images": { - "/broken.png": {"base64": "", "mime_type": "image/png"}, + "/broken.png": { + "mime_type": "image/png", + "size": 0, + "actual_path": str(tmp_path / "nonexistent.png"), + }, } } blocks = mw._create_image_details_message(state) - # header + description only (no image_url since base64 is empty) - assert len(blocks) == 2 + # header + description + error text (file no longer available) + assert len(blocks) == 3 assert all(not (isinstance(b, dict) and b.get("type") == "image_url") for b in blocks) - def test_uses_unknown_mime_type_when_missing(self): + def test_uses_unknown_mime_type_when_missing(self, tmp_path): mw = ViewImageMiddleware() + img_meta = _make_viewed_image(tmp_path, "mystery.bin", mime_type="unknown") state = { "viewed_images": { - "/mystery.bin": {"base64": "XYZ"}, # no mime_type key + "/mystery.bin": img_meta, } } blocks = mw._create_image_details_message(state) @@ -225,6 +244,51 @@ class TestCreateImageDetailsMessage: assert len(description_blocks) == 1 assert "unknown" in description_blocks[0]["text"] + def test_omits_image_url_when_read_raises_oserror(self, tmp_path, monkeypatch): + """A failure during on-demand read must not crash the middleware.""" + img_meta = _make_viewed_image(tmp_path, "ok.png") + state = { + "viewed_images": { + "/ok.png": img_meta, + } + } + + def _raise(*args, **kwargs): + raise OSError("disk error") + + monkeypatch.setattr("builtins.open", _raise) + + mw = ViewImageMiddleware() + blocks = mw._create_image_details_message(state) + # header + description + 'unavailable' text, no image_url block + assert all(not (isinstance(b, dict) and b.get("type") == "image_url") for b in blocks) + unavailable = [b for b in blocks if isinstance(b, dict) and b.get("type") == "text" and "unavailable" in b.get("text", "")] + assert len(unavailable) == 1 + + def test_omits_image_url_when_size_changes_between_view_and_inject(self, tmp_path): + """Defense against TOCTOU growth: skip if current size differs from recorded size.""" + img_meta = _make_viewed_image(tmp_path, "shrinking.png", data=b"original-larger-content") + # Grow the file after the metadata was written + img_meta_path = Path(img_meta["actual_path"]) + img_meta_path.write_bytes(b"much-much-much-larger-content-bytes") + + state = {"viewed_images": {"/shrinking.png": img_meta}} + mw = ViewImageMiddleware() + blocks = mw._create_image_details_message(state) + assert all(not (isinstance(b, dict) and b.get("type") == "image_url") for b in blocks) + + def test_omits_image_url_when_size_exceeds_cap(self, tmp_path): + """Records a small size but the actual file is large - the cap kicks in regardless.""" + img_meta = _make_viewed_image(tmp_path, "huge.png", data=b"x" * 100) + img_meta_path = Path(img_meta["actual_path"]) + # Grow past the cap (20 MB) + img_meta_path.write_bytes(b"y" * (21 * 1024 * 1024)) + + state = {"viewed_images": {"/huge.png": img_meta}} + mw = ViewImageMiddleware() + blocks = mw._create_image_details_message(state) + assert all(not (isinstance(b, dict) and b.get("type") == "image_url") for b in blocks) + class TestShouldInjectImageMessage: def test_false_when_no_messages(self): @@ -254,36 +318,34 @@ class TestShouldInjectImageMessage: state = {"messages": [assistant]} # no ToolMessage yet assert mw._should_inject_image_message(state) is False - def test_true_when_all_preconditions_met(self): + def test_true_when_all_preconditions_met(self, tmp_path): mw = ViewImageMiddleware() assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")]) + img_meta = _make_viewed_image(tmp_path) state = { "messages": [assistant, ToolMessage(content="ok", tool_call_id="c1")], - "viewed_images": { - "/img.png": {"base64": "AAA", "mime_type": "image/png"}, - }, + "viewed_images": {"/img.png": img_meta}, } assert mw._should_inject_image_message(state) is True - def test_false_when_already_injected(self): + def test_false_when_already_injected(self, tmp_path): """If a HumanMessage with the recognized header is already present after the assistant turn, we must not inject a duplicate.""" mw = ViewImageMiddleware() assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")]) already_injected = HumanMessage(content="Here are the images you've viewed: /img.png") + img_meta = _make_viewed_image(tmp_path) state = { "messages": [ assistant, ToolMessage(content="ok", tool_call_id="c1"), already_injected, ], - "viewed_images": { - "/img.png": {"base64": "AAA", "mime_type": "image/png"}, - }, + "viewed_images": {"/img.png": img_meta}, } assert mw._should_inject_image_message(state) is False - def test_false_when_already_injected_with_list_content(self): + def test_false_when_already_injected_with_list_content(self, tmp_path): """Deduplication must recognize the real injected payload shape. The middleware's own `_inject_image_message` creates a HumanMessage @@ -294,7 +356,8 @@ class TestShouldInjectImageMessage: """ mw = ViewImageMiddleware() assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")]) - viewed_images = {"/img.png": {"base64": "AAA", "mime_type": "image/png"}} + img_meta = _make_viewed_image(tmp_path) + viewed_images = {"/img.png": img_meta} # Build content the same way the middleware would. real_injected_content = mw._create_image_details_message({"viewed_images": viewed_images}) # Sanity: this is a list of blocks, not a plain string. @@ -311,21 +374,20 @@ class TestShouldInjectImageMessage: } assert mw._should_inject_image_message(state) is False - def test_false_when_legacy_details_marker_present(self): + def test_false_when_legacy_details_marker_present(self, tmp_path): """The middleware also recognizes the legacy 'Here are the details of the images you've viewed' marker as an already-injected signal.""" mw = ViewImageMiddleware() assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")]) legacy = HumanMessage(content="Here are the details of the images you've viewed: ...") + img_meta = _make_viewed_image(tmp_path) state = { "messages": [ assistant, ToolMessage(content="ok", tool_call_id="c1"), legacy, ], - "viewed_images": { - "/img.png": {"base64": "AAA", "mime_type": "image/png"}, - }, + "viewed_images": {"/img.png": img_meta}, } assert mw._should_inject_image_message(state) is False @@ -336,14 +398,13 @@ class TestInjectImageMessage: state = {"messages": []} assert mw._inject_image_message(state) is None - def test_returns_state_update_with_human_message(self): + def test_returns_state_update_with_human_message(self, tmp_path): mw = ViewImageMiddleware() assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")]) + img_meta = _make_viewed_image(tmp_path) state = { "messages": [assistant, ToolMessage(content="ok", tool_call_id="c1")], - "viewed_images": { - "/img.png": {"base64": "AAA", "mime_type": "image/png"}, - }, + "viewed_images": {"/img.png": img_meta}, } result = mw._inject_image_message(state) @@ -367,28 +428,26 @@ class TestBeforeModel: state = {"messages": [HumanMessage(content="hi")]} assert mw.before_model(state, _runtime()) is None - def test_before_model_returns_injection_when_ready(self): + def test_before_model_returns_injection_when_ready(self, tmp_path): mw = ViewImageMiddleware() assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")]) + img_meta = _make_viewed_image(tmp_path) state = { "messages": [assistant, ToolMessage(content="ok", tool_call_id="c1")], - "viewed_images": { - "/img.png": {"base64": "AAA", "mime_type": "image/png"}, - }, + "viewed_images": {"/img.png": img_meta}, } result = mw.before_model(state, _runtime()) assert result is not None assert isinstance(result["messages"][0], HumanMessage) @pytest.mark.anyio - async def test_abefore_model_matches_sync_behavior(self): + async def test_abefore_model_matches_sync_behavior(self, tmp_path): mw = ViewImageMiddleware() assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")]) + img_meta = _make_viewed_image(tmp_path) state = { "messages": [assistant, ToolMessage(content="ok", tool_call_id="c1")], - "viewed_images": { - "/img.png": {"base64": "AAA", "mime_type": "image/png"}, - }, + "viewed_images": {"/img.png": img_meta}, } result = await mw.abefore_model(state, _runtime()) assert result is not None diff --git a/backend/tests/test_view_image_tool.py b/backend/tests/test_view_image_tool.py index eb7db890c..a136b58d8 100644 --- a/backend/tests/test_view_image_tool.py +++ b/backend/tests/test_view_image_tool.py @@ -68,8 +68,9 @@ def test_view_image_reads_virtual_uploads_path(tmp_path: Path) -> None: assert _message_content(result) == "Successfully read image" viewed_image = result.update["viewed_images"]["/mnt/user-data/uploads/sample.png"] - assert viewed_image["base64"] == base64.b64encode(PNG_BYTES).decode("utf-8") assert viewed_image["mime_type"] == "image/png" + assert viewed_image["size"] == len(PNG_BYTES) + assert viewed_image["actual_path"] == str(image_path) def test_view_image_rejects_spoofed_extension(tmp_path: Path) -> None: