fix(agents): stop persisting base64 image data in checkpoint state (#4140)

* fix(agents): stop persisting base64 image data in checkpoint state (#4138)

The viewed_images state field stored full base64-encoded image data,
which was duplicated across every subsequent checkpoint (O(n * steps)
growth). A single 1MB image viewed early in a conversation would be
re-stored in every checkpoint for the rest of the session.

Changes:
- ViewedImageData: replace base64 field with lightweight metadata
  (mime_type, size, actual_path)
- view_image_tool: store only metadata in state, no base64 encoding
- ViewImageMiddleware: read image files from disk on-demand in
  before_model and encode base64 temporarily for the model call
- Update all tests to use the new metadata-only format

This is the first step of #4138. The base64 data is no longer in
persistent state, but the injected HumanMessage (with base64 content)
still appears in the checkpoint for the step where it was injected.
Checkpoint retention policies and large tool result dedup are separate
follow-up items.

* fix(agents): address review feedback on #4140

- view_image_tool: remove stale 'convert to base64' comment, replace with
  'validate contents'; drop redundant image_size reassignment and add a
  TOCTOU guard that rejects files changed between stat() and read().
- view_image_middleware: extract _read_image_as_data_url helper that
  re-checks size against the recorded value AND the absolute cap
  (_MAX_IMAGE_BYTES). Document the trust assumption for actual_path
  (server-set, not client-settable) in the helper docstring.
- view_image_middleware: abefore_model now runs the blocking read+encode
  via asyncio.to_thread to avoid stalling the event loop on up to 20MB
  images.
- tests: add coverage for OSError during read, file-changed-since-view
  (TOCTOU), and size-exceeds-cap branches.
This commit is contained in:
qin-chenghan
2026-07-14 23:02:26 +08:00
committed by GitHub
parent 51bb19fa1b
commit 713ee544b7
7 changed files with 202 additions and 68 deletions
@@ -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]}
@@ -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]
@@ -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)]},
+2 -2
View File
@@ -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"}
+3 -3
View File
@@ -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, {}) == {}
+103 -44
View File
@@ -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
+2 -1
View File
@@ -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: