fix(backend): remove transient image context after model calls (#4267)

* fix(backend): discard transient image context

* fix(backend): protect client image context ids

* docs(backend): clarify image checkpoint lifecycle

README.md

@@ -764,6 +764,8 @@ DeerFlow doesn't just *talk* about doing things. It has its own computer.
 
 Each task gets its own execution environment with a full filesystem view — skills, workspace, uploads, outputs. The agent reads, writes, and edits files. It can view images and, when configured safely, execute shell commands.
 
+Image bytes loaded for a vision-model call are transient: DeerFlow removes the hidden base64 message after the model consumes it so later checkpoints do not keep duplicating that payload.
+
 After each run, DeerFlow records a workspace change summary for the run-owned `workspace` and `outputs` directories. The Web UI shows a compact "files changed" badge on the assistant turn; opening it reveals created, modified, and deleted files with text diffs when safe to display. Uploads are excluded because they are user inputs, not agent-generated changes. Large, binary, or sensitive-looking files are shown as metadata only.
 
 With `AioSandboxProvider`, shell execution runs inside isolated containers. With `LocalSandboxProvider`, file tools still map to per-thread directories on the host, but host `bash` is disabled by default because it is not a secure isolation boundary. Re-enable host bash only for fully trusted local workflows. Host bash commands have a wall-clock timeout, and long-lived processes should be started in the background with output redirected to a workspace log.

backend/AGENTS.md

@@ -260,7 +260,7 @@ Before changing a later authorization phase, read the [authorization RFC](../doc
 20. **TokenUsageMiddleware** - *(optional, if `token_usage.enabled`)* Records token usage metrics; subagent usage is merged back into the dispatching AIMessage by message position
 21. **TitleMiddleware** - Auto-generates the thread title after the first complete exchange and normalizes structured message content before prompting the title model. If a first-turn run is interrupted before this middleware can write a title, `runtime/runs/worker.py` keeps the run in a finalizing state, persists a local fallback title from the latest checkpoint or original run input, and then syncs it to `threads_meta.display_name`. Replacement runs admitted by `multitask_strategy="interrupt"` / `"rollback"` wait for older same-thread finalization before entering the graph; the interrupted run only skips the fallback title write once a later run has started and may have advanced the checkpoint.
 22. **MemoryMiddleware** - Queues conversations for async memory update (filters to user + final AI responses)
-23. **ViewImageMiddleware** - *(optional, if the model supports vision)* Injects base64 image data before the LLM call
+23. **ViewImageMiddleware** - *(optional, if the model supports vision)* Injects a hidden HumanMessage with base64 image data, identified by a reserved ID prefix plus a server-owned metadata marker, before the LLM call. Because `before_model`, `model`, and `after_model` are separate graph nodes, the `before_model` and `model` node checkpoints for that call still contain the payload; `after_model` / `aafter_model` then emits `RemoveMessage`, so subsequent checkpoints do not retain it
 24. **McpRoutingMiddleware** - *(optional, if `tool_search.enabled` and PR1 MCP routing metadata produce a routing index)* Auto-promotes matching deferred MCP tool schemas before the model call by writing a minimal `promoted` state update. It matches only the latest real `HumanMessage`, uses the global `tool_search.auto_promote_top_k` limit (default 3, clamped to 1..5), never executes tools, and must be installed before `DeferredToolFilterMiddleware`
 25. **DeferredToolFilterMiddleware** - *(optional, if `tool_search.enabled`)* Hides deferred (MCP) tool schemas from the bound model until `tool_search` or `McpRoutingMiddleware` promotes them (reads per-thread promotions from `ThreadState.promoted`, hash-scoped)
 26. **SystemMessageCoalescingMiddleware** - Merges every SystemMessage into a single leading SystemMessage per request; provider-agnostic fix for strict backends (vLLM/SGLang/Qwen/Anthropic) that reject non-leading system messages. Touches the per-request payload only (checkpoint state unchanged); on midnight crossings only the latest `dynamic_context_reminder` SystemMessage survives
@@ -983,7 +983,7 @@ See [docs/summarization.md](docs/summarization.md) for details.
 For models with `supports_vision: true`:
 - `ViewImageMiddleware` processes images in conversation
 - `view_image_tool` added to agent's toolset
-- Images automatically converted to base64 and injected into state
+- Images are converted to base64 and injected into a hidden message carrying both a reserved ID prefix and a server-owned metadata marker for the model call; Gateway strips that marker from untrusted input, and the middleware requires both identifiers before removing the message. The `before_model` and `model` node checkpoints for that call still contain the payload; after `after_model` cleanup, subsequent checkpoints retain only lightweight `viewed_images` metadata, while client-chosen IDs survive
 
 ## Code Style
 

backend/app/gateway/services.py

@@ -30,6 +30,7 @@ from app.gateway.internal_auth import (
 )
 from app.gateway.utils import sanitize_log_param
 from deerflow.agents.middlewares.dynamic_context_middleware import _DYNAMIC_CONTEXT_REMINDER_KEY, _REMINDER_DATE_KEY
+from deerflow.agents.middlewares.view_image_middleware import _IMAGE_CONTEXT_MESSAGE_MARKER_KEY
 from deerflow.config.app_config import get_app_config
 from deerflow.runtime import (
     END_SENTINEL,
@@ -67,10 +68,11 @@ _TERMINAL_RUN_STATUSES = {
     RunStatus.interrupted,
 }
 
-_SERVER_OWNED_DYNAMIC_CONTEXT_KEYS = frozenset(
+_SERVER_OWNED_MESSAGE_METADATA_KEYS = frozenset(
     {
         _DYNAMIC_CONTEXT_REMINDER_KEY,
         _REMINDER_DATE_KEY,
+        _IMAGE_CONTEXT_MESSAGE_MARKER_KEY,
     }
 )
 
@@ -141,7 +143,7 @@ def _strip_external_message_metadata(message: Any) -> Any:
         return message
     additional_kwargs = dict(message.additional_kwargs)
     additional_kwargs.pop(ORIGINAL_USER_CONTENT_KEY, None)
-    for key in _SERVER_OWNED_DYNAMIC_CONTEXT_KEYS:
+    for key in _SERVER_OWNED_MESSAGE_METADATA_KEYS:
         additional_kwargs.pop(key, None)
     if additional_kwargs == message.additional_kwargs:
         return message
@@ -162,9 +164,10 @@ def normalize_input(raw_input: dict[str, Any] | None, *, trusted_internal: bool
     of bubbling up as a 500.  The gateway is a system boundary, so per-entry
     validation errors are the right shape for clients to retry against.
 
-    ``original_user_content`` and dynamic-context reminder markers are
-    server-owned. External callers cannot supply them; trusted internal channel
-    calls may preserve metadata they added before invoking this boundary.
+    ``original_user_content``, dynamic-context reminder markers, and the
+    transient view-image context marker are server-owned. External callers
+    cannot supply them; trusted internal channel calls may preserve metadata
+    they added before invoking this boundary.
     """
     if raw_input is None:
         return {}

backend/packages/harness/deerflow/agents/middlewares/view_image_middleware.py

@@ -5,9 +5,10 @@ import base64
 import logging
 from pathlib import Path
 from typing import override
+from uuid import uuid4
 
 from langchain.agents.middleware import AgentMiddleware
-from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage, ToolMessage
 from langgraph.runtime import Runtime
 
 from deerflow.agents.thread_state import ThreadState
@@ -18,6 +19,8 @@ logger = logging.getLogger(__name__)
 # 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
+_IMAGE_CONTEXT_MESSAGE_ID_PREFIX = "view-image-context:"
+_IMAGE_CONTEXT_MESSAGE_MARKER_KEY = "deerflow_view_image_context"
 
 
 class ViewImageMiddlewareState(ThreadState):
@@ -33,6 +36,7 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
     3. Verifies all tool calls in that message have been completed (have corresponding ToolMessages)
     4. If conditions are met, creates a human message with all viewed image details (including base64 data)
     5. Adds the message to state so the LLM can see and analyze the images
+    6. Removes the transient message after the LLM call so later checkpoints do not retain its base64 data
 
     This enables the LLM to automatically receive and analyze images that were loaded via view_image tool,
     without requiring explicit user prompts to describe the images.
@@ -40,6 +44,11 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
 
     state_schema = ViewImageMiddlewareState
 
+    @staticmethod
+    def _is_image_context_message(message: object) -> bool:
+        """Return whether a message is trusted transient image context."""
+        return isinstance(message, HumanMessage) and bool(message.id) and message.id.startswith(_IMAGE_CONTEXT_MESSAGE_ID_PREFIX) and message.additional_kwargs.get(_IMAGE_CONTEXT_MESSAGE_MARKER_KEY) is True
+
     def _get_last_assistant_message(self, messages: list) -> AIMessage | None:
         """Get the last assistant message from the message list.
 
@@ -204,6 +213,8 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
         assistant_idx = messages.index(last_assistant_msg)
         for msg in messages[assistant_idx + 1 :]:
             if isinstance(msg, HumanMessage):
+                if self._is_image_context_message(msg):
+                    return False
                 content_str = str(msg.content)
                 if "Here are the images you've viewed" in content_str or "Here are the details of the images you've viewed" in content_str:
                     # Already added, don't add again
@@ -211,6 +222,26 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
 
         return True
 
+    @staticmethod
+    def _create_image_context_message(content: list[str | dict]) -> HumanMessage:
+        """Create an identifiable, model-only image context message."""
+        return HumanMessage(
+            id=f"{_IMAGE_CONTEXT_MESSAGE_ID_PREFIX}{uuid4().hex}",
+            content=content,
+            additional_kwargs={
+                "hide_from_ui": True,
+                _IMAGE_CONTEXT_MESSAGE_MARKER_KEY: True,
+            },
+        )
+
+    @staticmethod
+    def _remove_image_context_messages(state: ViewImageMiddlewareState) -> dict | None:
+        """Remove transient image context messages after the model consumed them."""
+        removals = [RemoveMessage(id=msg.id) for msg in state.get("messages", []) if ViewImageMiddleware._is_image_context_message(msg)]
+        if not removals:
+            return None
+        return {"messages": removals}
+
     def _inject_image_message(self, state: ViewImageMiddlewareState) -> dict | None:
         """Internal helper to inject image details message.
 
@@ -229,7 +260,7 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
         # Create a new human message with mixed content (text + images). This is
         # internal context for the model only, so hide it from the chat UI and IM
         # channels (matches the other middleware-injected context messages).
-        human_msg = HumanMessage(content=image_content, additional_kwargs={"hide_from_ui": True})
+        human_msg = self._create_image_context_message(image_content)
 
         logger.debug("Injecting image details message with images before LLM call")
 
@@ -273,6 +304,16 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
         # 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})
+        human_msg = self._create_image_context_message(image_content)
         logger.debug("Injecting image details message with images before LLM call")
         return {"messages": [human_msg]}
+
+    @override
+    def after_model(self, state: ViewImageMiddlewareState, runtime: Runtime) -> dict | None:
+        """Remove model-only image data before subsequent checkpoints (sync version)."""
+        return self._remove_image_context_messages(state)
+
+    @override
+    async def aafter_model(self, state: ViewImageMiddlewareState, runtime: Runtime) -> dict | None:
+        """Remove model-only image data before subsequent checkpoints (async version)."""
+        return self._remove_image_context_messages(state)

backend/tests/test_gateway_services.py

@@ -186,6 +186,31 @@ def test_normalize_input_strips_external_dynamic_context_metadata():
     assert result["messages"][0].additional_kwargs == {"hide_from_ui": True, "custom": "keep-me"}
 
 
+def test_normalize_input_strips_external_view_image_context_marker():
+    from app.gateway.services import normalize_input
+    from deerflow.agents.middlewares.view_image_middleware import _IMAGE_CONTEXT_MESSAGE_MARKER_KEY
+
+    result = normalize_input(
+        {
+            "messages": [
+                {
+                    "role": "user",
+                    "id": "view-image-context:client-supplied",
+                    "content": "client-authored message",
+                    "additional_kwargs": {
+                        _IMAGE_CONTEXT_MESSAGE_MARKER_KEY: True,
+                        "custom": "keep-me",
+                    },
+                }
+            ]
+        }
+    )
+
+    message = result["messages"][0]
+    assert message.id == "view-image-context:client-supplied"
+    assert message.additional_kwargs == {"custom": "keep-me"}
+
+
 def test_normalize_input_preserves_trusted_internal_original_user_content():
     from app.gateway.services import normalize_input
     from deerflow.agents.middlewares.dynamic_context_middleware import _DYNAMIC_CONTEXT_REMINDER_KEY, _REMINDER_DATE_KEY

backend/tests/test_view_image_middleware.py

@@ -16,6 +16,8 @@ Covered behavior:
 - `_inject_image_message` returns a state update with a HumanMessage, or None
   when injection is not warranted.
 - `before_model` and `abefore_model` expose the same behavior sync/async.
+- `after_model` and `aafter_model` remove only the transient image message so
+  later checkpoints do not retain its base64 payload.
 """
 
 from pathlib import Path
@@ -23,9 +25,16 @@ from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
-from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
+from langchain.agents import create_agent
+from langchain_core.callbacks import BaseCallbackHandler
+from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
+from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage, SystemMessage, ToolMessage
+from langgraph.graph.message import add_messages
 
-from deerflow.agents.middlewares.view_image_middleware import ViewImageMiddleware
+from deerflow.agents.middlewares.view_image_middleware import (
+    _IMAGE_CONTEXT_MESSAGE_MARKER_KEY,
+    ViewImageMiddleware,
+)
 
 
 def _view_image_call(call_id: str = "call_1", path: str = "/mnt/user-data/uploads/img.png") -> dict:
@@ -42,6 +51,14 @@ def _runtime() -> MagicMock:
     return MagicMock()
 
 
+class _CaptureChatMessages(BaseCallbackHandler):
+    def __init__(self):
+        self.messages = []
+
+    def on_chat_model_start(self, serialized, messages, **kwargs):
+        self.messages = messages[0]
+
+
 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
@@ -420,6 +437,9 @@ class TestInjectImageMessage:
         # Internal injection: must be hidden from the chat UI (and IM channels),
         # like the other middleware-injected context messages.
         assert injected.additional_kwargs.get("hide_from_ui") is True
+        assert injected.additional_kwargs.get(_IMAGE_CONTEXT_MESSAGE_MARKER_KEY) is True
+        assert injected.id is not None
+        assert injected.id.startswith("view-image-context:")
 
 
 class TestBeforeModel:
@@ -458,3 +478,140 @@ class TestBeforeModel:
         mw = ViewImageMiddleware()
         state = {"messages": []}
         assert await mw.abefore_model(state, _runtime()) is None
+
+
+class TestAfterModel:
+    def test_graph_exposes_image_context_only_during_model_call(self, tmp_path):
+        capture = _CaptureChatMessages()
+        model = FakeMessagesListChatModel(
+            responses=[AIMessage(content="I can see the image.")],
+            callbacks=[capture],
+        )
+        graph = create_agent(model=model, tools=[], middleware=[ViewImageMiddleware()])
+        assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")])
+        img_meta = _make_viewed_image(tmp_path)
+
+        result = graph.invoke(
+            {
+                "messages": [assistant, ToolMessage(content="ok", tool_call_id="c1")],
+                "viewed_images": {"/img.png": img_meta},
+            }
+        )
+
+        model_image_messages = [message for message in capture.messages if isinstance(message, HumanMessage) and message.id and message.id.startswith("view-image-context:")]
+        assert len(model_image_messages) == 1
+        assert any(block.get("type") == "image_url" for block in model_image_messages[0].content)
+        assert all(not (isinstance(message, HumanMessage) and message.id and message.id.startswith("view-image-context:")) for message in result["messages"])
+
+    def test_removes_transient_image_message_from_later_state(self, tmp_path):
+        mw = ViewImageMiddleware()
+        assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")])
+        tool_result = ToolMessage(content="ok", tool_call_id="c1")
+        img_meta = _make_viewed_image(tmp_path)
+        before_state = {
+            "messages": [assistant, tool_result],
+            "viewed_images": {"/img.png": img_meta},
+        }
+        injected = mw.before_model(before_state, _runtime())["messages"][0]
+        model_response = AIMessage(content="I can see the image.")
+        model_state = {
+            **before_state,
+            "messages": [assistant, tool_result, injected, model_response],
+        }
+
+        result = mw.after_model(model_state, _runtime())
+
+        assert result is not None
+        assert len(result["messages"]) == 1
+        removal = result["messages"][0]
+        assert isinstance(removal, RemoveMessage)
+        assert removal.id == injected.id
+
+        checkpoint_messages = add_messages(model_state["messages"], result["messages"])
+        assert injected.id not in {message.id for message in checkpoint_messages}
+        assert model_response in checkpoint_messages
+
+    def test_does_not_remove_unmarked_human_messages(self):
+        mw = ViewImageMiddleware()
+        state = {
+            "messages": [
+                HumanMessage(
+                    id="view-image-context:client-supplied",
+                    content="Here are the images you've viewed: user-authored text",
+                    additional_kwargs={"hide_from_ui": True},
+                ),
+                AIMessage(content="response"),
+            ]
+        }
+
+        assert mw.after_model(state, _runtime()) is None
+
+    def test_graph_preserves_normalized_client_message_with_reserved_prefix(self, tmp_path):
+        from app.gateway.services import normalize_input
+
+        client_id = "view-image-context:client-supplied"
+        normalized = normalize_input(
+            {
+                "messages": [
+                    {
+                        "role": "user",
+                        "id": client_id,
+                        "content": "client-authored message",
+                        "additional_kwargs": {
+                            _IMAGE_CONTEXT_MESSAGE_MARKER_KEY: True,
+                            "custom": "keep-me",
+                        },
+                    }
+                ]
+            }
+        )
+        client_message = normalized["messages"][0]
+        assert _IMAGE_CONTEXT_MESSAGE_MARKER_KEY not in client_message.additional_kwargs
+
+        capture = _CaptureChatMessages()
+        model = FakeMessagesListChatModel(
+            responses=[AIMessage(content="I can see the image.")],
+            callbacks=[capture],
+        )
+        graph = create_agent(model=model, tools=[], middleware=[ViewImageMiddleware()])
+        assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")])
+
+        result = graph.invoke(
+            {
+                "messages": [
+                    client_message,
+                    assistant,
+                    ToolMessage(content="ok", tool_call_id="c1"),
+                ],
+                "viewed_images": {"/img.png": _make_viewed_image(tmp_path)},
+            }
+        )
+
+        assert any(message.id == client_id for message in capture.messages)
+        assert any(isinstance(message, HumanMessage) and message.id != client_id and message.additional_kwargs.get(_IMAGE_CONTEXT_MESSAGE_MARKER_KEY) is True for message in capture.messages)
+        persisted_client = next(message for message in result["messages"] if message.id == client_id)
+        assert persisted_client.content == "client-authored message"
+        assert persisted_client.additional_kwargs == {"custom": "keep-me"}
+        assert all(message.additional_kwargs.get(_IMAGE_CONTEXT_MESSAGE_MARKER_KEY) is not True for message in result["messages"] if isinstance(message, HumanMessage))
+
+    @pytest.mark.anyio
+    async def test_aafter_model_matches_sync_cleanup(self, tmp_path):
+        mw = ViewImageMiddleware()
+        assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")])
+        tool_result = ToolMessage(content="ok", tool_call_id="c1")
+        img_meta = _make_viewed_image(tmp_path)
+        before_state = {
+            "messages": [assistant, tool_result],
+            "viewed_images": {"/img.png": img_meta},
+        }
+        injected = (await mw.abefore_model(before_state, _runtime()))["messages"][0]
+        model_state = {
+            **before_state,
+            "messages": [assistant, tool_result, injected, AIMessage(content="response")],
+        }
+
+        result = await mw.aafter_model(model_state, _runtime())
+
+        assert result is not None
+        assert isinstance(result["messages"][0], RemoveMessage)
+        assert result["messages"][0].id == injected.id