diff --git a/backend/AGENTS.md b/backend/AGENTS.md
index 7b823a704..2a247457f 100644
--- a/backend/AGENTS.md
+++ b/backend/AGENTS.md
@@ -289,7 +289,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 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
+- Images are converted to base64 and appended to the model request as a hidden message carrying both a reserved ID prefix and a server-owned metadata marker; Gateway strips that marker from untrusted input, and the middleware requires both identifiers to recognize its own message. The middleware injects inside `wrap_model_call`, so the payload never enters graph state: checkpoints retain only lightweight `viewed_images` metadata, while client-chosen IDs survive. It also sweeps its own message out of every request before rebuilding it, so a payload stranded in an older checkpoint by an interrupted run stops being resent
 
 ## Code Style
 
diff --git a/backend/docs/middleware-execution-flow.md b/backend/docs/middleware-execution-flow.md
index 99d638938..6fd5ecec2 100644
--- a/backend/docs/middleware-execution-flow.md
+++ b/backend/docs/middleware-execution-flow.md
@@ -16,7 +16,7 @@
 | 7 | TodoMiddleware | | ✓ | ✓ | | ✓ | | ✓ | ✗ | `plan_mode` 参数 |
 | 8 | TitleMiddleware | | | ✓ | | | | ✓ | ✗ | `auto_title` |
 | 9 | MemoryMiddleware | | | | ✓ | | | ✓ | ✗ | `memory` |
-| 10 | ViewImageMiddleware | | ✓ | | | | | ✓ | ✗ | `vision` |
+| 10 | ViewImageMiddleware | | | | | ✓ | | ✓ | ✗ | `vision` |
 | 11 | SubagentLimitMiddleware | | | ✓ | | | | ✓ | ✗ | `subagent` |
 | 12 | LoopDetectionMiddleware | ✓ | | ✓ | ✓ | ✓ | | ✓ | ✗ | 始终开启 |
 | 13 | ClarificationMiddleware | | | | | | ✓ | ✓ | ✗ | 始终最后 |
@@ -38,18 +38,12 @@ graph TB
         TD["[0] ThreadData<br/>创建线程目录"] --> UL["[1] Uploads<br/>扫描上传文件"] --> SB["[2] Sandbox<br/>获取沙箱"] --> LD_BA["[12] LoopDetection<br/>清理 stale warning"]
     end
 
-    subgraph BM ["<b>before_model</b> 正序 0→N"]
+    subgraph WM ["<b>wrap_model_call</b> 外→内 0→N"]
         direction TB
-        VI["[10] ViewImage<br/>注入图片 base64"]
+        DTC_WM["[3] DanglingToolCall<br/>补悬空 ToolMessage"] --> VI["[10] ViewImage<br/>注入图片 base64（仅请求内）"] --> LD_WM["[12] LoopDetection<br/>注入当前 run warning"]
     end
 
-    subgraph WM ["<b>wrap_model_call</b>"]
-        direction TB
-        DTC_WM["[3] DanglingToolCall<br/>补悬空 ToolMessage"] --> LD_WM["[12] LoopDetection<br/>注入当前 run warning"]
-    end
-
-    LD_BA --> VI
-    VI --> DTC_WM
+    LD_BA --> DTC_WM
     LD_WM --> M["<b>MODEL</b>"]
 
     subgraph AM ["<b>after_model</b> 反序 N→0"]
@@ -74,8 +68,8 @@ graph TB
     classDef afterAgentNode fill:#a0b5a8,stroke:#637a6b,color:#2d3239
     classDef terminalNode fill:#a8b5a0,stroke:#6b7a63,color:#2d3239
 
-    class TD,UL,SB,LD_BA,VI beforeNode
-    class DTC_WM,LD_WM wrapModelNode
+    class TD,UL,SB,LD_BA beforeNode
+    class DTC_WM,VI,LD_WM wrapModelNode
     class M modelNode
     class LD,SL,TI afterModelNode
     class LD_CLEAN,SBR,MEM afterAgentNode
@@ -91,8 +85,8 @@ sequenceDiagram
     participant UL as UploadsMiddleware
     participant SB as SandboxMiddleware
     participant LD as LoopDetectionMiddleware
-    participant VI as ViewImageMiddleware
     participant DTC as DanglingToolCallMiddleware
+    participant VI as ViewImageMiddleware
     participant M as MODEL
     participant SL as SubagentLimitMiddleware
     participant TI as TitleMiddleware
@@ -113,14 +107,14 @@ sequenceDiagram
     SB ->> LD: before_agent
     activate LD
     Note right of LD: before_agent 清理同 thread 旧 run 的 pending warning
-    LD ->> VI: before_model
-    activate VI
-    Note right of VI: before_model 注入图片 base64
-
-    VI ->> DTC: wrap_model_call
+    LD ->> DTC: wrap_model_call
     activate DTC
     Note right of DTC: wrap_model_call 补悬空 ToolMessage
-    DTC ->> LD: wrap_model_call
+
+    DTC ->> VI: wrap_model_call
+    activate VI
+    Note right of VI: wrap_model_call 把图片 base64 追加到请求（不写入 state）
+    VI ->> LD: wrap_model_call
     Note right of LD: wrap_model_call drain 当前 run warning 并追加到末尾
     LD ->> M: messages + tools
     activate M
@@ -138,14 +132,14 @@ sequenceDiagram
 
     activate TI
     Note right of TI: after_model 生成标题
-    TI -->> DTC: done
+    TI -->> VI: done
     deactivate TI
 
-    deactivate DTC
-
-    VI -->> SB: done
     deactivate VI
 
+    DTC -->> SB: done
+    deactivate DTC
+
     Note right of LD: after_agent 清理当前 run 未消费 warning
 
     Note right of MEM: after_agent 入队记忆
@@ -166,7 +160,8 @@ sequenceDiagram
 列表位置决定在洋葱中的层级 — 位置 0 最外层，位置 N 最内层：
 
 ```
-进入 before_*：   [0] → [1] → [2] → ... → [10] → MODEL
+进入 before_*：   [0] → [1] → [2] → ... → [7] → MODEL
+进入 wrap_model_call： [3] → [10] → [12] → MODEL（外→内，同样正序）
 退出 after_*：    MODEL → [13] → [11] → ... → [6] → [3] → [2] → [0]
                           ↑ 最内层最先执行
 ```
@@ -233,8 +228,8 @@ sequenceDiagram
     participant UL as Uploads
     participant SB as Sandbox
     participant LD as LoopDetection
-    participant VI as ViewImage
     participant DTC as DanglingToolCall
+    participant VI as ViewImage
     participant M as MODEL
     participant SL as SubagentLimit
     participant TI as Title
@@ -250,11 +245,11 @@ sequenceDiagram
     Note right of LD: before_agent 清理 stale pending warning
 
     loop 每轮对话（tool call 循环）
-        SB ->> VI: .
-        Note right of VI: before_model 注入图片
-        VI ->> DTC: .
+        SB ->> DTC: .
         Note right of DTC: wrap_model_call 补悬空工具结果
-        DTC ->> LD: .
+        DTC ->> VI: .
+        Note right of VI: wrap_model_call 把图片追加到请求
+        VI ->> LD: .
         Note right of LD: wrap_model_call 注入当前 run warning
         LD ->> M: messages + tools
         M -->> LD: AI response
diff --git a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md
index 79d8b26ca..6870deef3 100644
--- a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md
+++ b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md
@@ -84,7 +84,7 @@ Before changing a later authorization phase, read the [authorization RFC](../../
 20. **TokenUsageMiddleware** - *(optional, if `token_usage.enabled`)* Records token usage metrics; subagent usage is read from terminal `ToolMessage.additional_kwargs` in the current run and merged back into the dispatching AIMessage by message position. The same state update marks the ToolMessage with `subagent_token_usage_attributed=true`, so checkpoint replay or middleware re-entry cannot add the cumulative snapshot twice; missing/malformed usage or a result with no matching dispatch remains unmarked and retryable.
 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); captures the runtime-resolved user so standalone LangGraph Server reads and writes stay in the same bucket
-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
+23. **ViewImageMiddleware** - *(optional, if the model supports vision)* Appends a hidden HumanMessage with base64 image data, identified by a reserved ID prefix plus a server-owned metadata marker, to `ModelRequest.messages` in `wrap_model_call` / `awrap_model_call`. The payload lives only in that request and is never returned as a state update, so no checkpoint carries it and an interrupted run cannot strand it in history; state keeps only the lightweight `viewed_images` metadata. It owns that context and rebuilds it per call: its own message is swept out of the request first — a thread checkpointed by the earlier `before_model`/`after_model` pair (which wrote the payload into state and took it back out with `RemoveMessage`) can carry one that reached state but was never removed, and leaving it in would resend that base64 in every later request for the life of the thread — then a freshly built one is appended when warranted. The sweep requires both the reserved ID prefix and the server-owned marker, so a client cannot get its own message dropped; unmarked leftovers predating the marker are left in place and merely not duplicated
 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. The subagent builder places its date-only context middleware immediately before this coalescer, so the built-in subagent prompt and hidden date reminder still reach providers as one leading system block
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 d1f1a4225..3f148cd54 100644
--- a/backend/packages/harness/deerflow/agents/middlewares/view_image_middleware.py
+++ b/backend/packages/harness/deerflow/agents/middlewares/view_image_middleware.py
@@ -1,16 +1,17 @@
-"""Middleware for injecting image details into conversation before LLM call."""
+"""Middleware for injecting image details into the model request."""
 
 import asyncio
 import base64
 import logging
+from collections.abc import Awaitable, Callable
 from pathlib import Path
 from typing import override
 from uuid import uuid4
 
 from deerflow_extension_api import ContentKind, provenance_kwargs
 from langchain.agents.middleware import AgentMiddleware
-from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage, ToolMessage
-from langgraph.runtime import Runtime
+from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse
+from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage
 
 from deerflow.agents.thread_state import ThreadState
 
@@ -29,18 +30,24 @@ class ViewImageMiddlewareState(ThreadState):
 
 
 class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
-    """Injects image details as a human message before LLM calls when view_image tools have completed.
+    """Injects image details into the model request when view_image tool calls have completed.
 
     This middleware:
-    1. Runs before each LLM call
+    1. Wraps each LLM call
     2. Checks if the last assistant message contains view_image tool calls
     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
+    4. If conditions are met, appends a human message with all viewed image details (including base64 data)
+    5. Hands the augmented request to the model so it can see and analyze the images
 
     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.
+
+    Injection happens in ``wrap_model_call`` on purpose: the message exists only
+    in ``ModelRequest.messages`` and is never returned as a state update, so no
+    checkpoint carries the base64 payload and an interrupted run cannot strand it
+    in history. Do not move this back to a ``before_model``/``after_model`` pair
+    -- that writes the payload into state and can only take it out again
+    afterwards (see #4267).
     """
 
     state_schema = ViewImageMiddlewareState
@@ -183,16 +190,15 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
 
         return content_blocks
 
-    def _should_inject_image_message(self, state: ViewImageMiddlewareState) -> bool:
-        """Determine if we should inject an image details message.
+    def _should_inject_image_message(self, messages: list[AnyMessage]) -> bool:
+        """Determine if we should append an image details message.
 
         Args:
-            state: Current state
+            messages: Messages about to be sent to the model
 
         Returns:
-            True if we should inject the message
+            True if we should append the message
         """
-        messages = state.get("messages", [])
         if not messages:
             return False
 
@@ -209,13 +215,14 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
         if not self._all_tools_completed(messages, last_assistant_msg):
             return False
 
-        # Check if we've already added an image details message
-        # Look for a human message after the last assistant message that contains image details
+        # Skip when image details are already present. ``_inject`` has stripped
+        # this middleware's own messages by now, so what remains are unmarked
+        # ones from checkpoints written before the marker existed -- those cannot
+        # be told apart from user-authored text with certainty, so they are left
+        # in place and simply not duplicated.
         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
@@ -236,86 +243,55 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
             },
         )
 
-    @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.
+    def _inject(self, request: ModelRequest) -> ModelRequest:
+        """Rebuild the request's image context from ``viewed_images``.
 
         Args:
-            state: Current state
+            request: The pending model request
 
         Returns:
-            State update with additional human message, or None if no update needed
+            A request whose messages carry exactly the image context this call
+            warrants -- one freshly built message, or none -- leaving the
+            original request untouched when there is nothing to change
         """
-        if not self._should_inject_image_message(state):
-            return None
+        # This middleware owns the image context and rebuilds it from
+        # ``viewed_images`` on every call, so drop any copy already in the list.
+        # A thread checkpointed by the earlier before_model/after_model pair can
+        # carry one that reached state but was never removed (the run died during
+        # the model call); left in place it would ride along in every later
+        # request for the life of the thread. Matching requires both the reserved
+        # ID prefix and the server-owned marker, and Gateway strips that marker
+        # from client input, so this can never drop a user-authored message.
+        messages = [message for message in request.messages if not self._is_image_context_message(message)]
+        dropped_stranded = len(messages) != len(request.messages)
+        if dropped_stranded:
+            logger.debug("Dropping %d stranded image context message(s) from the model request", len(request.messages) - len(messages))
 
-        # Create the image details message with text and image content
-        image_content = self._create_image_details_message(state)
+        if not self._should_inject_image_message(messages):
+            return request.override(messages=messages) if dropped_stranded else request
 
-        # 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 = self._create_image_context_message(image_content)
+        # Mixed content (text + images) for the model only, so hide it from the
+        # chat UI and IM channels (matches the other middleware-injected context
+        # messages) even though it never leaves this request.
+        image_content = self._create_image_details_message(request.state or {})
+        logger.debug("Injecting image details message with images into the model request")
 
-        logger.debug("Injecting image details message with images before LLM call")
-
-        # Return state update with the new message
-        return {"messages": [human_msg]}
+        return request.override(messages=[*messages, self._create_image_context_message(image_content)])
 
     @override
-    def before_model(self, state: ViewImageMiddlewareState, runtime: Runtime) -> dict | None:
-        """Inject image details message before LLM call if view_image tools have completed (sync version).
-
-        This runs before each LLM call, checking if the previous turn included view_image
-        tool calls that have all completed. If so, it injects a human message with the image
-        details so the LLM can see and analyze the images.
-
-        Args:
-            state: Current state
-            runtime: Runtime context (unused but required by interface)
-
-        Returns:
-            State update with additional human message, or None if no update needed
-        """
-        return self._inject_image_message(state)
+    def wrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], ModelResponse],
+    ) -> ModelCallResult:
+        return handler(self._inject(request))
 
     @override
-    async def abefore_model(self, state: ViewImageMiddlewareState, runtime: Runtime) -> dict | None:
-        """Inject image details message before LLM call if view_image tools have completed (async version).
-
-        This runs before each LLM call, checking if the previous turn included view_image
-        tool calls that have all completed. If so, it injects a human message with the image
-        details so the LLM can see and analyze the images.
-
-        Args:
-            state: Current state
-            runtime: Runtime context (unused but required by interface)
-
-        Returns:
-            State update with additional human message, or None if no update needed
-        """
-        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 = 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)
+    async def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelCallResult:
+        # Image reads + base64 encoding can be slow (up to 20MB), so offload the
+        # blocking work to a thread rather than stalling the event loop.
+        return await handler(await asyncio.to_thread(self._inject, request))
diff --git a/backend/packages/harness/deerflow/runtime/serialization.py b/backend/packages/harness/deerflow/runtime/serialization.py
index a9b5cea2f..6d0c84587 100644
--- a/backend/packages/harness/deerflow/runtime/serialization.py
+++ b/backend/packages/harness/deerflow/runtime/serialization.py
@@ -75,9 +75,11 @@ def strip_data_url_image_blocks(messages: list[dict[str, Any]]) -> list[dict[str
     """Remove ``data:``-scheme ``image_url`` blocks from *hide_from_ui* messages.
 
     The history and run-wait endpoints return checkpoint-persisted messages to
-    the frontend.  ``ViewImageMiddleware`` stores full base64 image payloads in
-    ``hide_from_ui`` human messages — these are internal model context and must
-    not be sent over the wire (huge response bodies, no UI value).
+    the frontend.  ``ViewImageMiddleware`` now keeps its base64 image payloads
+    inside the model request, but threads checkpointed by earlier versions still
+    hold them in ``hide_from_ui`` human messages — these are internal model
+    context and must not be sent over the wire (huge response bodies, no UI
+    value).
 
     Only content blocks of type ``image_url`` whose URL starts with ``data:``
     are stripped.  Text blocks, ``https://`` image URLs, and non-hidden
diff --git a/backend/tests/test_view_image_middleware.py b/backend/tests/test_view_image_middleware.py
index cf58ab46f..4018441d0 100644
--- a/backend/tests/test_view_image_middleware.py
+++ b/backend/tests/test_view_image_middleware.py
@@ -1,9 +1,9 @@
 """Unit tests for ViewImageMiddleware.
 
 Tests cover the middleware's ability to inject image details (including base64
-payloads) as a HumanMessage before the next LLM call, triggered only when the
-previous assistant turn contained `view_image` tool calls that have all been
-completed with corresponding ToolMessages.
+payloads) into the model request, triggered only when the previous assistant
+turn contained `view_image` tool calls that have all been completed with
+corresponding ToolMessages.
 
 Covered behavior:
 - `_get_last_assistant_message` returns the most recent AIMessage (or None).
@@ -12,12 +12,13 @@ Covered behavior:
 - `_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
-  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.
+  deduplication when an image-details message is already in the request.
+- `_inject` rebuilds the request's image context: it sweeps out any copy left
+  in the message list by an interrupted run before deciding whether to append a
+  freshly built one.
+- `wrap_model_call` and `awrap_model_call` expose the same behavior sync/async,
+  handing the payload to the model without ever writing it to state — so no
+  checkpoint retains it, even if the run is interrupted mid-call.
 """
 
 from pathlib import Path
@@ -26,10 +27,10 @@ from unittest.mock import MagicMock
 
 import pytest
 from langchain.agents import create_agent
+from langchain.agents.middleware.types import ModelRequest
 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 langchain_core.messages import AIMessage, AnyMessage, HumanMessage, SystemMessage, ToolMessage
 
 from deerflow.agents.middlewares.view_image_middleware import (
     _IMAGE_CONTEXT_MESSAGE_MARKER_KEY,
@@ -45,10 +46,19 @@ def _other_tool_call(call_id: str = "call_other", name: str = "bash") -> dict:
     return {"name": name, "id": call_id, "args": {"command": "ls"}}
 
 
-def _runtime() -> MagicMock:
-    """Minimal Runtime stub. The middleware doesn't use it today, but the
-    interface requires it."""
-    return MagicMock()
+def _model_request(messages: list[AnyMessage], viewed_images: dict | None = None) -> ModelRequest:
+    """Build a real ModelRequest so `.override()` behaves as it does in the graph."""
+    return ModelRequest(
+        model=FakeMessagesListChatModel(responses=[AIMessage(content="ok")]),
+        messages=list(messages),
+        system_message=None,
+        tool_choice=None,
+        tools=[],
+        response_format=None,
+        state={"messages": list(messages), "viewed_images": viewed_images or {}},
+        runtime=MagicMock(),
+        model_settings={},
+    )
 
 
 class _CaptureChatMessages(BaseCallbackHandler):
@@ -59,6 +69,10 @@ class _CaptureChatMessages(BaseCallbackHandler):
         self.messages = messages[0]
 
 
+def _image_context_messages(messages: list[AnyMessage]) -> list[HumanMessage]:
+    return [message for message in messages if isinstance(message, HumanMessage) and message.id and message.id.startswith("view-image-context:")]
+
+
 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
@@ -310,126 +324,98 @@ class TestCreateImageDetailsMessage:
 class TestShouldInjectImageMessage:
     def test_false_when_no_messages(self):
         mw = ViewImageMiddleware()
-        assert mw._should_inject_image_message({"messages": []}) is False
-
-    def test_false_when_messages_key_missing(self):
-        mw = ViewImageMiddleware()
-        assert mw._should_inject_image_message({}) is False
+        assert mw._should_inject_image_message([]) is False
 
     def test_false_when_no_assistant_message(self):
         mw = ViewImageMiddleware()
-        state = {"messages": [HumanMessage(content="hello")]}
-        assert mw._should_inject_image_message(state) is False
+        assert mw._should_inject_image_message([HumanMessage(content="hello")]) is False
 
     def test_false_when_no_view_image_tool_call(self):
         mw = ViewImageMiddleware()
         assistant = AIMessage(content="", tool_calls=[_other_tool_call()])
-        state = {
-            "messages": [assistant, ToolMessage(content="ok", tool_call_id="call_other")],
-        }
-        assert mw._should_inject_image_message(state) is False
+        messages = [assistant, ToolMessage(content="ok", tool_call_id="call_other")]
+        assert mw._should_inject_image_message(messages) is False
 
     def test_false_when_tool_not_completed(self):
         mw = ViewImageMiddleware()
         assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")])
-        state = {"messages": [assistant]}  # no ToolMessage yet
-        assert mw._should_inject_image_message(state) is False
+        assert mw._should_inject_image_message([assistant]) is False  # no ToolMessage yet
 
-    def test_true_when_all_preconditions_met(self, tmp_path):
+    def test_true_when_all_preconditions_met(self):
         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": img_meta},
-        }
-        assert mw._should_inject_image_message(state) is True
+        messages = [assistant, ToolMessage(content="ok", tool_call_id="c1")]
+        assert mw._should_inject_image_message(messages) is True
 
-    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."""
+    def test_false_when_already_injected(self):
+        """A checkpoint written by an older version (or by a run that died before
+        its `RemoveMessage` cleanup landed) can still carry image details. Do not
+        add a duplicate on top of one."""
         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": img_meta},
-        }
-        assert mw._should_inject_image_message(state) is False
+        messages = [
+            assistant,
+            ToolMessage(content="ok", tool_call_id="c1"),
+            already_injected,
+        ]
+        assert mw._should_inject_image_message(messages) is False
 
     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
-        whose `.content` is a *list* of dicts (text + image_url blocks), not a
-        plain string. This test reuses `_create_image_details_message` output
-        to reproduce the realistic shape and confirms `_should_inject_image_message`
-        still detects the marker via `str(msg.content)`.
+        An unmarked leftover carries `.content` as a *list* of dicts (text +
+        image_url blocks), not a plain string. This test reuses
+        `_create_image_details_message` output to reproduce the realistic shape
+        and confirms the marker is still detected via `str(msg.content)`.
         """
         mw = ViewImageMiddleware()
         assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")])
-        img_meta = _make_viewed_image(tmp_path)
-        viewed_images = {"/img.png": img_meta}
+        viewed_images = {"/img.png": _make_viewed_image(tmp_path)}
         # 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.
         assert isinstance(real_injected_content, list)
-        already_injected = HumanMessage(content=real_injected_content)
 
-        state = {
-            "messages": [
-                assistant,
-                ToolMessage(content="ok", tool_call_id="c1"),
-                already_injected,
-            ],
-            "viewed_images": viewed_images,
-        }
-        assert mw._should_inject_image_message(state) is False
+        messages = [
+            assistant,
+            ToolMessage(content="ok", tool_call_id="c1"),
+            HumanMessage(content=real_injected_content),
+        ]
+        assert mw._should_inject_image_message(messages) is False
 
-    def test_false_when_legacy_details_marker_present(self, tmp_path):
+    def test_false_when_legacy_details_marker_present(self):
         """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": img_meta},
-        }
-        assert mw._should_inject_image_message(state) is False
+        messages = [
+            assistant,
+            ToolMessage(content="ok", tool_call_id="c1"),
+            legacy,
+        ]
+        assert mw._should_inject_image_message(messages) is False
 
 
-class TestInjectImageMessage:
-    def test_returns_none_when_should_not_inject(self):
+class TestInject:
+    def test_returns_request_unchanged_when_should_not_inject(self):
         mw = ViewImageMiddleware()
-        state = {"messages": []}
-        assert mw._inject_image_message(state) is None
+        request = _model_request([HumanMessage(content="hi")])
+        assert mw._inject(request) is request
 
-    def test_returns_state_update_with_human_message(self, tmp_path):
+    def test_appends_image_context_message_to_request(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": img_meta},
-        }
+        original = [assistant, ToolMessage(content="ok", tool_call_id="c1")]
+        request = _model_request(original, {"/img.png": _make_viewed_image(tmp_path)})
 
-        result = mw._inject_image_message(state)
+        injected_request = mw._inject(request)
 
-        assert isinstance(result, dict)
-        assert "messages" in result
-        assert len(result["messages"]) == 1
-        injected = result["messages"][0]
+        assert injected_request is not request
+        # The payload is appended last, so it directly follows the tool results.
+        assert injected_request.messages[:-1] == original
+        injected = injected_request.messages[-1]
         assert isinstance(injected, HumanMessage)
         # Mixed-content payload: list of text + image_url blocks
         assert isinstance(injected.content, list)
@@ -441,110 +427,146 @@ class TestInjectImageMessage:
         assert injected.id is not None
         assert injected.id.startswith("view-image-context:")
 
-
-class TestBeforeModel:
-    def test_before_model_returns_none_when_preconditions_not_met(self):
-        mw = ViewImageMiddleware()
-        state = {"messages": [HumanMessage(content="hi")]}
-        assert mw.before_model(state, _runtime()) is None
-
-    def test_before_model_returns_injection_when_ready(self, tmp_path):
+    def test_replaces_a_stranded_payload_instead_of_stacking_a_second_one(self, tmp_path):
+        """A run that died during the model call can leave the old
+        before_model/after_model pair's message checkpointed. Rebuild it rather
+        than adding a second copy on top."""
         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": img_meta},
-        }
-        result = mw.before_model(state, _runtime())
-        assert result is not None
-        assert isinstance(result["messages"][0], HumanMessage)
+        stranded = ViewImageMiddleware._create_image_context_message([{"type": "text", "text": "stale"}])
+        request = _model_request(
+            [assistant, ToolMessage(content="ok", tool_call_id="c1"), stranded],
+            {"/img.png": _make_viewed_image(tmp_path)},
+        )
 
-    @pytest.mark.anyio
-    async def test_abefore_model_matches_sync_behavior(self, tmp_path):
+        injected = _image_context_messages(mw._inject(request).messages)
+
+        assert len(injected) == 1
+        assert injected[0].id != stranded.id
+        assert any(isinstance(b, dict) and b.get("type") == "image_url" for b in injected[0].content)
+
+    def test_drops_a_stranded_payload_even_when_no_injection_is_warranted(self):
+        """Otherwise the stale base64 would ride along in every later request for
+        the life of the thread -- the old `after_model` swept it, so dropping the
+        hook must not lose that."""
+        mw = ViewImageMiddleware()
+        stranded = ViewImageMiddleware._create_image_context_message([{"type": "text", "text": "stale"}])
+        request = _model_request([HumanMessage(content="hi"), stranded, AIMessage(content="done")])
+
+        prepared = mw._inject(request)
+
+        assert prepared is not request
+        assert _image_context_messages(prepared.messages) == []
+        assert [type(m) for m in prepared.messages] == [HumanMessage, AIMessage]
+
+    def test_never_drops_a_client_message_wearing_the_reserved_prefix(self):
+        """The prefix alone is not enough — Gateway strips the server-owned
+        marker from client input, and both are required to match."""
+        mw = ViewImageMiddleware()
+        client_message = HumanMessage(id="view-image-context:client-supplied", content="client-authored", additional_kwargs={"hide_from_ui": True})
+        request = _model_request([client_message, AIMessage(content="done")])
+
+        assert mw._inject(request) is request
+
+    def test_does_not_mutate_the_incoming_request(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": img_meta},
-        }
-        result = await mw.abefore_model(state, _runtime())
-        assert result is not None
-        assert isinstance(result["messages"][0], HumanMessage)
+        request = _model_request(
+            [assistant, ToolMessage(content="ok", tool_call_id="c1")],
+            {"/img.png": _make_viewed_image(tmp_path)},
+        )
+
+        mw._inject(request)
+
+        assert _image_context_messages(request.messages) == []
+
+
+class TestWrapModelCall:
+    def test_handler_receives_the_image_context_message(self, tmp_path):
+        mw = ViewImageMiddleware()
+        assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")])
+        request = _model_request(
+            [assistant, ToolMessage(content="ok", tool_call_id="c1")],
+            {"/img.png": _make_viewed_image(tmp_path)},
+        )
+        seen: list[ModelRequest] = []
+
+        def handler(prepared: ModelRequest) -> AIMessage:
+            seen.append(prepared)
+            return AIMessage(content="I can see the image.")
+
+        result = mw.wrap_model_call(request, handler)
+
+        assert result.content == "I can see the image."
+        assert len(_image_context_messages(seen[0].messages)) == 1
+
+    def test_handler_receives_request_unchanged_when_not_warranted(self):
+        mw = ViewImageMiddleware()
+        request = _model_request([HumanMessage(content="hi")])
+        seen: list[ModelRequest] = []
+
+        mw.wrap_model_call(request, lambda prepared: seen.append(prepared) or AIMessage(content="ok"))
+
+        assert seen[0] is request
 
     @pytest.mark.anyio
-    async def test_abefore_model_returns_none_when_no_injection(self):
+    async def test_awrap_model_call_matches_sync_behavior(self, tmp_path):
         mw = ViewImageMiddleware()
-        state = {"messages": []}
-        assert await mw.abefore_model(state, _runtime()) is None
+        assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")])
+        request = _model_request(
+            [assistant, ToolMessage(content="ok", tool_call_id="c1")],
+            {"/img.png": _make_viewed_image(tmp_path)},
+        )
+        seen: list[ModelRequest] = []
+
+        async def handler(prepared: ModelRequest) -> AIMessage:
+            seen.append(prepared)
+            return AIMessage(content="I can see the image.")
+
+        result = await mw.awrap_model_call(request, handler)
+
+        assert result.content == "I can see the image."
+        assert len(_image_context_messages(seen[0].messages)) == 1
 
 
-class TestAfterModel:
-    def test_graph_exposes_image_context_only_during_model_call(self, tmp_path):
+class TestGraphIntegration:
+    def _graph_and_capture(self):
         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)
+        return create_agent(model=model, tools=[], middleware=[ViewImageMiddleware()]), capture
 
-        result = graph.invoke(
-            {
-                "messages": [assistant, ToolMessage(content="ok", tool_call_id="c1")],
-                "viewed_images": {"/img.png": img_meta},
-            }
-        )
+    def _input(self, tmp_path):
+        return {
+            "messages": [
+                AIMessage(content="", tool_calls=[_view_image_call("c1")]),
+                ToolMessage(content="ok", tool_call_id="c1"),
+            ],
+            "viewed_images": {"/img.png": _make_viewed_image(tmp_path)},
+        }
 
-        model_image_messages = [message for message in capture.messages if isinstance(message, HumanMessage) and message.id and message.id.startswith("view-image-context:")]
+    def test_image_context_reaches_the_model_but_never_the_state(self, tmp_path):
+        graph, capture = self._graph_and_capture()
+
+        result = graph.invoke(self._input(tmp_path))
+
+        model_image_messages = _image_context_messages(capture.messages)
         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"])
+        # Nothing is written back, so the payload is absent from every checkpoint
+        # rather than being added and then removed again.
+        assert _image_context_messages(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],
-        }
+    @pytest.mark.anyio
+    async def test_async_graph_matches_sync_behavior(self, tmp_path):
+        graph, capture = self._graph_and_capture()
 
-        result = mw.after_model(model_state, _runtime())
+        result = await graph.ainvoke(self._input(tmp_path))
 
-        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
+        assert len(_image_context_messages(capture.messages)) == 1
+        assert _image_context_messages(result["messages"]) == []
 
     def test_graph_preserves_normalized_client_message_with_reserved_prefix(self, tmp_path):
         from app.gateway.services import normalize_input
@@ -568,50 +590,14 @@ class TestAfterModel:
         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")])
+        graph, capture = self._graph_and_capture()
+        graph_input = self._input(tmp_path)
 
-        result = graph.invoke(
-            {
-                "messages": [
-                    client_message,
-                    assistant,
-                    ToolMessage(content="ok", tool_call_id="c1"),
-                ],
-                "viewed_images": {"/img.png": _make_viewed_image(tmp_path)},
-            }
-        )
+        result = graph.invoke({**graph_input, "messages": [client_message, *graph_input["messages"]]})
 
         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)
+        assert any(message.id != client_id and message.additional_kwargs.get(_IMAGE_CONTEXT_MESSAGE_MARKER_KEY) is True for message in _image_context_messages(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