github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-09-09 21:49:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(middleware): stop checkpointing view_image base64 payloads (#5014)

Browse Source 
ViewImageMiddleware injected the viewed-image message from before_model and
removed it again from after_model. before_model, model, and after_model are
separate graph nodes, so every view_image turn cost two extra nodes and two
state writes, and up to 20MB of base64 sat in two checkpoints for the duration
of the model call. A run interrupted in that window (user cancel, restart)
stranded the payload in history for good.

Inject from wrap_model_call instead, so the message lives only in
ModelRequest.messages and is never returned as a state update:

- before_model/after_model (and the async pair) are replaced by
  wrap_model_call/awrap_model_call; _remove_image_context_messages and its
  RemoveMessage bookkeeping go with them. The async hook keeps the existing
  asyncio.to_thread offload for the file read and base64 encode.
- _should_inject_image_message gates on request.messages rather than state, so
  the decision is made against what the model will actually see.
- _inject sweeps this middleware's own message out of the request before
  rebuilding it. Dropping after_model also drops the cleanup it did on every
  call, so without the sweep a payload stranded by an older interrupted run
  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 a user message is never dropped.

Chain position is unchanged, and wrap_model_call nests first-registered
outermost, so TokenBudgetMiddleware still sees the image message and enforces
the input budget against it.

Checkpoint rows that already hold a stranded payload keep it on disk. It is
inert -- never sent to a provider, and strip_data_url_image_blocks keeps it off
the wire -- and reclaiming it would mean keeping the node this change removes.

tests/test_view_image_middleware.py is rewritten around the new hook (43
tests): sync/async at unit and graph level, the stranded sweep, and the
client-message protection. Docs: middleware chain entry 23, Vision Support, the
middleware-execution-flow hook matrix and diagrams, and the
strip_data_url_image_blocks docstring.
This commit is contained in:
Hyeonsang Cho 2026-08-27 10:12:37 +09:00 committed by GitHub
parent 788a890bd0
commit cb24bc2699
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 294 additions and 335 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
backend/AGENTS.md
View File

@ -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

53
backend/docs/middleware-execution-flow.md
View File

@ -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

2
backend/packages/harness/deerflow/agents/middlewares/AGENTS.md
View File

@ -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

152
backend/packages/harness/deerflow/agents/middlewares/view_image_middleware.py
View File

@ -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))

8
backend/packages/harness/deerflow/runtime/serialization.py
View File

@ -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

412
backend/tests/test_view_image_middleware.py
View File

@ -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