mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-21 10:15:47 +00:00
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:
@@ -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)]},
|
||||
|
||||
@@ -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"}
|
||||
|
||||
|
||||
|
||||
@@ -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, {}) == {}
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user