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

fix(view-image): read active sandbox images from sandbox (#5306)

Browse Source 
* fix(view-image): read remote sandbox images from sandbox

* docs(tools): clarify view_image sandbox behavior

* fix(view-image): address sandbox lifecycle review

* fix(view-image): preserve image provenance across sandbox replacement

* fix(view-image): address provider recovery review

* fix(view-image): drain cancelled tool reads
This commit is contained in:
Jun 2026-09-10 15:48:55 +08:00 committed by GitHub
parent 69f0f483eb
commit 3a6e681dee
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 1319 additions and 92 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1,7 +1,7 @@
"""Middleware for injecting image details into the model request."""
import asyncio
import base64
import hashlib
import logging
from collections.abc import Awaitable, Callable
from pathlib import Path
@ -14,12 +14,13 @@ from langchain.agents.middleware.types import ModelCallResult, ModelRequest, Mod
from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage
from deerflow.agents.thread_state import ThreadState
from deerflow.sandbox.lease import run_sync_lifecycle_operation
logger = logging.getLogger(__name__)
# Mirror the tool-side size cap as a defense-in-depth check. The tool
# enforces this at write time; the middleware re-checks at read time in
# case the file grew on disk between view and injection.
# case the file grew between view and injection.
_MAX_IMAGE_BYTES = 20 * 1024 * 1024
_IMAGE_CONTEXT_MESSAGE_ID_PREFIX = "view-image-context:"
_IMAGE_CONTEXT_MESSAGE_MARKER_KEY = "deerflow_view_image_context"
@ -117,41 +118,153 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
return tool_call_ids.issubset(completed_tool_ids)
@staticmethod
def _read_image_as_data_url(actual_path: str, mime_type: str, expected_size: int) -> str | None:
"""Read image file and return a `data:` URL, or None on failure.
def _encode_image_bytes(
image_bytes: bytes,
mime_type: str,
expected_size: int,
expected_sha256: str | None = None,
) -> str | None:
"""Validate image bytes against recorded metadata and return a data URL."""
current_size = len(image_bytes)
if current_size != expected_size or current_size > _MAX_IMAGE_BYTES:
return None
if expected_sha256 is not None and hashlib.sha256(image_bytes).hexdigest() != expected_sha256:
return None
base64_data = base64.b64encode(image_bytes).decode("utf-8")
return f"data:{mime_type};base64,{base64_data}"
Trust assumption: ``actual_path`` is set by ``view_image_tool``
(server-side, validated against the allowed virtual roots at write
time) and held in LangGraph-controlled state. Client input cannot
reach this field, so the read scope is trusted. We still re-check
size at read time to defend against TOCTOU growth and skip files
exceeding ``_MAX_IMAGE_BYTES``.
@classmethod
def _read_host_image_as_data_url(
cls,
actual_path: str,
mime_type: str,
expected_size: int,
expected_sha256: str | None = None,
) -> str | None:
"""Read a validated host mirror and return a data URL, or None on failure.
``actual_path`` is server-set by ``view_image_tool`` and held in
LangGraph-controlled state. The host path remains the compatibility path
for local execution and older checkpoints. Provenance-aware checkpoints
additionally verify the exact SHA-256 before a synchronized host copy can
stand in for bytes from an earlier sandbox generation.
"""
try:
file_path = Path(actual_path)
if not file_path.exists() or not file_path.is_file():
return None
current_size = file_path.stat().st_size
if current_size != expected_size:
# File changed between view and inject - skip.
return None
if current_size > _MAX_IMAGE_BYTES:
if current_size != expected_size or current_size > _MAX_IMAGE_BYTES:
return None
with open(file_path, "rb") as f:
image_bytes = f.read()
base64_data = base64.b64encode(image_bytes).decode("utf-8")
return f"data:{mime_type};base64,{base64_data}"
return cls._encode_image_bytes(
image_bytes,
mime_type,
expected_size,
expected_sha256,
)
except OSError:
return None
@classmethod
def _read_image_as_data_url(
cls,
state: ViewImageMiddlewareState,
image_path: str,
actual_path: str,
mime_type: str,
expected_size: int,
expected_sha256: str | None,
source_sandbox_id: str | None,
) -> str | None:
"""Read the exact image bytes represented by ``viewed_images`` metadata.
A live sandbox is authoritative only for metadata recorded from that same
sandbox generation. If the thread now points at a replacement sandbox,
the previous image can be reconstructed from the synchronized host mirror
only when its SHA-256 exactly matches the bytes that ``view_image`` saw.
Legacy metadata without a digest never authorizes this cross-generation
fallback. When no live sandbox exists, the historical host compatibility
path remains available (digest-checked when present).
"""
from deerflow.sandbox.overwrite import unwrap_sandbox
from deerflow.sandbox.sandbox_provider import get_sandbox_provider
sandbox_state, _ = unwrap_sandbox(state.get("sandbox"))
sandbox_id = sandbox_state.get("sandbox_id") if isinstance(sandbox_state, dict) else None
sandbox = get_sandbox_provider().get(sandbox_id) if sandbox_id else None
if sandbox is not None:
provenance_matches_live = source_sandbox_id == sandbox_id
provenance_identifies_other_source = expected_sha256 is not None and source_sandbox_id != sandbox_id
if provenance_identifies_other_source:
# The current client belongs to a different generation (or the
# image was originally read from the host). Reproduce the exact
# historical bytes rather than letting an unrelated same-path
# file in the replacement sandbox win.
if actual_path:
host_data_url = cls._read_host_image_as_data_url(
actual_path,
mime_type,
expected_size,
expected_sha256,
)
if host_data_url is not None:
return host_data_url
try:
image_bytes = sandbox.download_file(image_path)
except Exception:
logger.warning(
"Failed to recover viewed image %s from replacement sandbox %s",
image_path,
sandbox_id,
exc_info=True,
)
return None
return cls._encode_image_bytes(
image_bytes,
mime_type,
expected_size,
expected_sha256,
)
if not provenance_matches_live and expected_sha256 is None:
# A legacy checkpoint cannot prove which sandbox generation
# supplied these bytes. Do not silently reinterpret historical
# image context through a newly active remote filesystem.
return None
try:
image_bytes = sandbox.download_file(image_path)
except Exception:
logger.warning("Failed to read viewed image %s from sandbox %s", image_path, sandbox_id, exc_info=True)
return None
return cls._encode_image_bytes(
image_bytes,
mime_type,
expected_size,
expected_sha256,
)
if not actual_path:
return None
return cls._read_host_image_as_data_url(
actual_path,
mime_type,
expected_size,
expected_sha256,
)
def _create_image_details_message(self, state: ViewImageMiddlewareState) -> list[str | dict]:
"""Create a formatted message with all viewed image details.
Reads image files from disk on-demand and encodes them as base64
for the model. The base64 data is NOT persisted in state -- only
lightweight metadata (path, mime_type, size) is stored in
``viewed_images``, avoiding large duplicate payloads across every
checkpoint (see #4138).
Reads image files on-demand from the active sandbox when available and
encodes them as base64 for the model. The base64 data is NOT persisted in
state -- only lightweight metadata (path, mime_type, size, digest, and
source sandbox id when applicable) is stored in ``viewed_images``,
avoiding large duplicate payloads across every checkpoint (see #4138).
Args:
state: Current state containing viewed_images
@ -171,22 +284,31 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
mime_type = image_data.get("mime_type", "unknown")
actual_path = image_data.get("actual_path", "")
expected_size = image_data.get("size", 0)
expected_sha256 = image_data.get("sha256")
source_sandbox_id = image_data.get("source_sandbox_id")
# Add text description
content_blocks.append({"type": "text", "text": f"\n- **{image_path}** ({mime_type})"})
# Read the image file on-demand and encode as base64 for the model
if actual_path:
data_url = self._read_image_as_data_url(actual_path, mime_type, expected_size)
if data_url:
content_blocks.append(
{
"type": "image_url",
"image_url": {"url": data_url},
}
)
else:
content_blocks.append({"type": "text", "text": f" (file unavailable or changed on disk: {actual_path})"})
data_url = self._read_image_as_data_url(
state,
image_path,
actual_path,
mime_type,
expected_size,
expected_sha256 if isinstance(expected_sha256, str) else None,
source_sandbox_id if isinstance(source_sandbox_id, str) else None,
)
if data_url:
content_blocks.append(
{
"type": "image_url",
"image_url": {"url": data_url},
}
)
else:
content_blocks.append({"type": "text", "text": f" (file unavailable or changed: {image_path})"})
return content_blocks
@ -284,6 +406,9 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
request: ModelRequest,
handler: Callable[[ModelRequest], ModelResponse],
) -> ModelCallResult:
# Sync injection executes inline on this call stack. There is no detached
# worker to drain: an outer sandbox lease cannot reach its finally/release
# boundary until this blocking read returns or raises.
return handler(self._inject(request))
@override
@ -293,5 +418,8 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
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))
# blocking work without allowing cancellation to outlive a sandbox
# client operation. The outer run lease may release the client as soon as
# cancellation propagates.
injected_request = await run_sync_lifecycle_operation(self._inject, request)
return await handler(injected_request)

10
backend/packages/harness/deerflow/agents/thread_state.py
View File

@ -52,15 +52,17 @@ class BackgroundTaskState(TypedDict):
class ViewedImageData(TypedDict):
"""Metadata for a viewed image file.
Only lightweight metadata is persisted in checkpoint state; the actual
image bytes are read on-demand from disk when the model needs them.
This avoids duplicating large base64 payloads across every checkpoint
(see #4138).
Only lightweight metadata is persisted in checkpoint state. Image bytes are
read on-demand from the active sandbox or from a synchronized host copy whose
size and SHA-256 match the previously viewed bytes. This avoids duplicating
large base64 payloads across every checkpoint (see #4138).
"""
mime_type: str
size: int
actual_path: str
sha256: str
source_sandbox_id: NotRequired[str]
def merge_sandbox(existing: SandboxState | None, new: SandboxState | None) -> SandboxState | None:

4
backend/packages/harness/deerflow/tools/AGENTS.md
View File

@ -6,7 +6,7 @@
3. **Built-in tools**:
- `present_files` - Make output files visible to user (only `/mnt/user-data/outputs`); virtual paths use `resolve_runtime_user_id(runtime)` so validation resolves the same user-scoped outputs directory established by `ThreadDataMiddleware`
- `ask_clarification` - Request clarification (intercepted by ClarificationMiddleware, which preserves text fallback and adds `artifact.human_input` for Web UI Human Input Cards). Beyond free text and single choice, the request-side v2 protocol supports `fields` (structured form card collecting several values at once; field types: text/textarea/number/select/multi_select/checkbox/date, validated and normalized server-side in the middleware — invalid entries are dropped, unknown types degrade to `text`; a standalone multi-select question is a one-field form). Replies stay on the v1 response protocol (`text`/`option`): the form card submits a readable text summary
- `view_image` - Read image as base64 (added only if model supports vision)
- `view_image` - Read image bytes for vision-capable models; live sandbox bytes win for the same sandbox generation, replacement-sandbox recovery uses only SHA-256-verified synchronized host bytes, and async tool invocation drains blocking reads before cancellation may release the sandbox lease
- `setup_agent` - Bootstrap-only: persist a brand-new custom agent's `SOUL.md` and `config.yaml`. Bound only when `is_bootstrap=True`.
- `update_agent` - Custom-agent-only: persist self-updates to the current agent's `SOUL.md` / `config.yaml` from inside a normal chat (partial update + atomic write). Bound when `agent_name` is set and `is_bootstrap=False`.
4. **Subagent tool** (if enabled):
@ -42,4 +42,4 @@ E2B output sync records remote file versions and actual host file metadata in a
- MiniMax Code speaks ACP directly: configure `command: mcode` with `args: ["acp"]`. It receives DeerFlow's enabled MCP servers and uses the per-thread ACP workspace; the Gateway process must have an authenticated `mcode` executable on `PATH`
- ACP results collect only `agent_message_chunk` text. Thought chunks remain internal and must not be concatenated into the tool result
- Missing ACP executables now return an actionable error message instead of a raw `[Errno 2]`
- Each ACP agent uses a per-thread workspace at `{base_dir}/users/{user_id}/threads/{thread_id}/acp-workspace/`. The workspace is accessible to the lead agent via the virtual path `/mnt/acp-workspace/` (read-only). In docker sandbox mode, the directory is volume-mounted into the container at `/mnt/acp-workspace` (read-only); in local sandbox mode, path translation is handled by `tools.py`
- Each ACP agent uses a per-thread workspace at `{base_dir}/users/{user_id}/threads/{thread_id}/acp-workspace/`. The workspace is accessible to the lead agent via the virtual path `/mnt/acp-workspace/` (read-only). In docker sandbox mode, the directory is volume-mounted into the container at `/mnt/acp-workspace` (read-only); in local sandbox mode, path translation is handled by `tools.py`

233
backend/packages/harness/deerflow/tools/builtins/view_image_tool.py
View File

@ -1,9 +1,11 @@
import hashlib
import mimetypes
from pathlib import Path
from typing import Annotated
from langchain.tools import InjectedToolCallId, tool
from langchain.tools import InjectedToolCallId
from langchain_core.messages import ToolMessage
from langchain_core.tools import StructuredTool
from langgraph.types import Command
from deerflow.agents.thread_state import ThreadDataState
@ -48,8 +50,60 @@ def _sanitize_image_error(error: Exception, thread_data: ThreadDataState | None)
return mask_local_paths_in_output(f"{type(error).__name__}: {error}", thread_data)
@tool("view_image", parse_docstring=True)
def view_image_tool(
def _is_file_not_found_error(error: BaseException) -> bool:
"""Recognize an explicit missing-file signal through provider wrappers.
``Sandbox.download_file`` promises ``OSError`` for read failures, while
remote SDKs expose missing paths in different explicit forms: builtin or
provider-defined ``FileNotFoundError`` types, E2B's
``FileNotFoundException``, and HTTP-style exceptions carrying
``status_code == 404``. Walk only explicit ``raise ... from`` causes so an
unrelated exception being handled when a transport failure is raised cannot
accidentally authorize historical host recovery. Error-message strings are
deliberately never parsed.
"""
current: BaseException | None = error
seen: set[int] = set()
while current is not None and id(current) not in seen:
seen.add(id(current))
error_type = type(current)
if isinstance(current, FileNotFoundError) or error_type.__name__ == "FileNotFoundError":
return True
if error_type.__name__ == "FileNotFoundException" and error_type.__module__.split(".", 1)[0] == "e2b":
return True
if getattr(current, "status_code", None) == 404:
return True
current = current.__cause__
return False
def _read_verified_host_copy(
actual_path: str | Path,
*,
expected_size: int,
expected_sha256: str,
) -> bytes | None:
"""Read a synchronized host image only when it matches prior metadata."""
path = Path(actual_path)
try:
if not path.exists() or not path.is_file():
return None
size = path.stat().st_size
if size != expected_size or size > _MAX_IMAGE_BYTES:
return None
data = path.read_bytes()
except OSError:
return None
if len(data) != size:
return None
if hashlib.sha256(data).hexdigest() != expected_sha256:
return None
return data
def _view_image(
runtime: Runtime,
image_path: str,
tool_call_id: Annotated[str, InjectedToolCallId],
@ -69,6 +123,8 @@ def view_image_tool(
image_path: Absolute /mnt/user-data virtual path to the image file. Common formats supported: jpg, jpeg, png, webp, gif.
"""
from deerflow.sandbox.exceptions import SandboxRuntimeError
from deerflow.sandbox.overwrite import unwrap_sandbox
from deerflow.sandbox.sandbox_provider import get_sandbox_provider
from deerflow.sandbox.tools import (
get_thread_data,
resolve_and_validate_user_data_path,
@ -97,58 +153,108 @@ def view_image_tool(
update={"messages": [ToolMessage(f"Error: {str(e)}", tool_call_id=tool_call_id)]},
)
path = Path(actual_path)
# Validate that the file exists
if not path.exists():
return Command(
update={"messages": [ToolMessage(f"Error: Image file not found: {image_path}", tool_call_id=tool_call_id)]},
)
# Validate that it's a file (not a directory)
if not path.is_file():
return Command(
update={"messages": [ToolMessage(f"Error: Path is not a file: {image_path}", tool_call_id=tool_call_id)]},
)
# Validate image extension
expected_mime_type = _EXTENSION_TO_MIME.get(path.suffix.lower())
image_suffix = Path(image_path).suffix.lower()
expected_mime_type = _EXTENSION_TO_MIME.get(image_suffix)
if expected_mime_type is None:
return Command(
update={"messages": [ToolMessage(f"Error: Unsupported image format: {path.suffix}. Supported formats: {', '.join(_EXTENSION_TO_MIME)}", tool_call_id=tool_call_id)]},
update={"messages": [ToolMessage(f"Error: Unsupported image format: {image_suffix}. Supported formats: {', '.join(_EXTENSION_TO_MIME)}", tool_call_id=tool_call_id)]},
)
# Detect MIME type from file extension
mime_type, _ = mimetypes.guess_type(actual_path)
mime_type, _ = mimetypes.guess_type(image_path)
if mime_type is None:
mime_type = expected_mime_type
try:
image_size = path.stat().st_size
except OSError as e:
return Command(
update={"messages": [ToolMessage(f"Error reading image metadata: {_sanitize_image_error(e, thread_data)}", tool_call_id=tool_call_id)]},
)
state = runtime.state or {}
sandbox_state, _ = unwrap_sandbox(state.get("sandbox"))
sandbox_id = sandbox_state.get("sandbox_id") if isinstance(sandbox_state, dict) else None
sandbox = get_sandbox_provider().get(sandbox_id) if sandbox_id else None
viewed_images = state.get("viewed_images")
previous_view = viewed_images.get(image_path) if isinstance(viewed_images, dict) else None
previous_source_id = previous_view.get("source_sandbox_id") if isinstance(previous_view, dict) else None
read_source_sandbox_id: str | None = None
if sandbox is not None:
try:
image_data = sandbox.download_file(image_path)
read_source_sandbox_id = sandbox_id
except IsADirectoryError:
return Command(
update={"messages": [ToolMessage(f"Error: Path is not a file: {image_path}", tool_call_id=tool_call_id)]},
)
except Exception as e:
# A replacement sandbox may be live without containing files from
# the earlier generation. Recover only from an explicitly missing
# file and only when the synchronized host copy matches the exact
# metadata of the previously viewed image. Other live-client
# failures stay fail-closed so a stale mirror cannot mask them.
if _is_file_not_found_error(e) and isinstance(previous_view, dict) and previous_source_id != sandbox_id:
previous_size = previous_view.get("size")
previous_sha256 = previous_view.get("sha256")
if isinstance(previous_size, int) and isinstance(previous_sha256, str):
recovered = _read_verified_host_copy(
actual_path,
expected_size=previous_size,
expected_sha256=previous_sha256,
)
if recovered is not None:
image_data = recovered
else:
return Command(
update={"messages": [ToolMessage(f"Error: Image file not found: {image_path}", tool_call_id=tool_call_id)]},
)
else:
return Command(
update={"messages": [ToolMessage(f"Error: Image file not found: {image_path}", tool_call_id=tool_call_id)]},
)
elif _is_file_not_found_error(e):
return Command(
update={"messages": [ToolMessage(f"Error: Image file not found: {image_path}", tool_call_id=tool_call_id)]},
)
else:
return Command(
update={"messages": [ToolMessage(f"Error reading image file: {_sanitize_image_error(e, thread_data)}", tool_call_id=tool_call_id)]},
)
image_size = len(image_data)
else:
path = Path(actual_path)
if not path.exists():
return Command(
update={"messages": [ToolMessage(f"Error: Image file not found: {image_path}", tool_call_id=tool_call_id)]},
)
if not path.is_file():
return Command(
update={"messages": [ToolMessage(f"Error: Path is not a file: {image_path}", tool_call_id=tool_call_id)]},
)
try:
image_size = path.stat().st_size
except OSError as e:
return Command(
update={"messages": [ToolMessage(f"Error reading image metadata: {_sanitize_image_error(e, thread_data)}", tool_call_id=tool_call_id)]},
)
if image_size > _MAX_IMAGE_BYTES:
return Command(
update={"messages": [ToolMessage(f"Error: Image file is too large: {image_size} bytes. Maximum supported size is {_MAX_IMAGE_BYTES} bytes", tool_call_id=tool_call_id)]},
)
try:
with open(actual_path, "rb") as f:
image_data = f.read()
except Exception as e:
return Command(
update={"messages": [ToolMessage(f"Error reading image file: {_sanitize_image_error(e, thread_data)}", tool_call_id=tool_call_id)]},
)
if len(image_data) != image_size:
return Command(
update={"messages": [ToolMessage("Error: Image file changed during read", tool_call_id=tool_call_id)]},
)
if image_size > _MAX_IMAGE_BYTES:
return Command(
update={"messages": [ToolMessage(f"Error: Image file is too large: {image_size} bytes. Maximum supported size is {_MAX_IMAGE_BYTES} bytes", tool_call_id=tool_call_id)]},
)
# Read image file to validate contents (magic bytes + size)
try:
with open(actual_path, "rb") as f:
image_data = f.read()
except Exception as e:
return Command(
update={"messages": [ToolMessage(f"Error reading image file: {_sanitize_image_error(e, thread_data)}", tool_call_id=tool_call_id)]},
)
if len(image_data) != image_size:
# File changed between stat() and read() - reject for safety.
return Command(
update={"messages": [ToolMessage("Error: Image file changed during read", tool_call_id=tool_call_id)]},
)
detected_mime_type = _detect_image_mime(image_data)
if detected_mime_type is None:
return Command(
@ -160,17 +266,40 @@ def view_image_tool(
)
mime_type = detected_mime_type
# Store only lightweight metadata in state (not base64) to avoid
# duplicating large payloads across every checkpoint (see #4138).
# The middleware reads the file on-demand when the model needs it.
new_viewed_images = {
image_path: {
"mime_type": mime_type,
"size": image_size,
"actual_path": str(actual_path),
}
image_metadata = {
"mime_type": mime_type,
"size": image_size,
"actual_path": str(actual_path),
"sha256": hashlib.sha256(image_data).hexdigest(),
}
if read_source_sandbox_id is not None:
image_metadata["source_sandbox_id"] = read_source_sandbox_id
new_viewed_images = {image_path: image_metadata}
return Command(
update={"viewed_images": new_viewed_images, "messages": [ToolMessage("Successfully read image", tool_call_id=tool_call_id)]},
)
async def _aview_image(
runtime: Runtime,
image_path: str,
tool_call_id: Annotated[str, InjectedToolCallId],
) -> Command:
"""Run the blocking image read without letting cancellation outlive it."""
from deerflow.sandbox.lease import run_sync_lifecycle_operation
return await run_sync_lifecycle_operation(
_view_image,
runtime,
image_path,
tool_call_id,
)
view_image_tool = StructuredTool.from_function(
func=_view_image,
coroutine=_aview_image,
name="view_image",
parse_docstring=True,
)

130
backend/tests/test_view_image_provider_error_classification.py Normal file
View File

@ -0,0 +1,130 @@
import base64
import hashlib
from pathlib import Path
from types import SimpleNamespace
from deerflow.agents.thread_state import ViewedImageData
from deerflow.tools.builtins.view_image_tool import _is_file_not_found_error, view_image_tool
PNG_BYTES = base64.b64decode("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==")
class _HttpNotFoundError(Exception):
status_code = 404
class _ProviderFileNotFoundError(Exception):
pass
_ProviderFileNotFoundError.__name__ = "FileNotFoundError"
class _HttpMissingSandbox:
id = "remote-new"
def __init__(self) -> None:
self.downloads: list[str] = []
def download_file(self, path: str) -> bytes:
self.downloads.append(path)
try:
raise _HttpNotFoundError("not found")
except _HttpNotFoundError as error:
raise OSError(f"cannot read '{path}' from remote provider") from error
class _Provider:
def __init__(self, sandbox: _HttpMissingSandbox) -> None:
self.sandbox = sandbox
def get(self, sandbox_id: str):
return self.sandbox if sandbox_id == self.sandbox.id else None
def _thread_data(tmp_path: Path) -> dict[str, str]:
user_data = tmp_path / "threads" / "thread-1" / "user-data"
workspace = user_data / "workspace"
uploads = user_data / "uploads"
outputs = user_data / "outputs"
for directory in (workspace, uploads, outputs):
directory.mkdir(parents=True)
return {
"workspace_path": str(workspace),
"uploads_path": str(uploads),
"outputs_path": str(outputs),
}
def test_classifier_accepts_explicit_http_404_cause():
try:
try:
raise _HttpNotFoundError("not found")
except _HttpNotFoundError as cause:
raise OSError("wrapped remote read failure") from cause
except OSError as error:
assert _is_file_not_found_error(error)
def test_classifier_accepts_provider_defined_file_not_found_type():
try:
raise _ProviderFileNotFoundError("not found")
except _ProviderFileNotFoundError as error:
assert _is_file_not_found_error(error)
def test_classifier_ignores_implicit_exception_context():
try:
try:
raise FileNotFoundError("unrelated cleanup miss")
except FileNotFoundError:
raise OSError("transport timeout")
except OSError as error:
assert error.__cause__ is None
assert isinstance(error.__context__, FileNotFoundError)
assert not _is_file_not_found_error(error)
def test_http_404_replacement_sandbox_recovers_verified_host_copy(tmp_path, monkeypatch):
thread_data = _thread_data(tmp_path)
host_path = Path(thread_data["outputs_path"]) / "plot.png"
host_path.write_bytes(PNG_BYTES)
sandbox = _HttpMissingSandbox()
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: _Provider(sandbox),
)
runtime = SimpleNamespace(
state={
"thread_data": thread_data,
"sandbox": {"sandbox_id": sandbox.id},
"viewed_images": {
"/mnt/user-data/outputs/plot.png": {
"mime_type": "image/png",
"size": len(PNG_BYTES),
"actual_path": str(host_path),
"sha256": hashlib.sha256(PNG_BYTES).hexdigest(),
"source_sandbox_id": "remote-old",
}
},
},
context={"thread_id": "thread-1"},
config={},
)
result = view_image_tool.func(
runtime=runtime,
image_path="/mnt/user-data/outputs/plot.png",
tool_call_id="tc-http-not-found",
)
assert result.update["messages"][0].content == "Successfully read image"
viewed = result.update["viewed_images"]["/mnt/user-data/outputs/plot.png"]
assert viewed["sha256"] == hashlib.sha256(PNG_BYTES).hexdigest()
assert "source_sandbox_id" not in viewed
assert sandbox.downloads == ["/mnt/user-data/outputs/plot.png"]
def test_viewed_image_state_contract_includes_provenance_keys():
assert "sha256" in ViewedImageData.__required_keys__
assert "source_sandbox_id" in ViewedImageData.__optional_keys__

704
backend/tests/test_view_image_remote_sandbox.py Normal file
View File

@ -0,0 +1,704 @@
import asyncio
import base64
import hashlib
import threading
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
from langchain.agents.middleware.types import ModelRequest
from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
from langchain_core.messages import AIMessage, ToolMessage
from deerflow.agents.middlewares.view_image_middleware import ViewImageMiddleware
from deerflow.sandbox.lease import SandboxLeaseManager
from deerflow.tools.builtins.view_image_tool import view_image_tool
PNG_BYTES = base64.b64decode("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==")
STALE_SAME_SIZE_PNG_BYTES = PNG_BYTES[:-1] + bytes([PNG_BYTES[-1] ^ 1])
_E2BFileNotFound = type(
"FileNotFoundException",
(Exception,),
{"__module__": "e2b.filesystem.filesystem"},
)
class _RemoteSandbox:
def __init__(self, image_bytes: bytes, *, sandbox_id: str = "remote-1") -> None:
self.id = sandbox_id
self.image_bytes = image_bytes
self.downloads: list[str] = []
self.released_scopes: list[str] = []
def download_file(self, path: str) -> bytes:
self.downloads.append(path)
return self.image_bytes
def release_command_scope(self, scope_id: str) -> None:
self.released_scopes.append(scope_id)
class _FailingRemoteSandbox(_RemoteSandbox):
def download_file(self, path: str) -> bytes:
self.downloads.append(path)
raise OSError("remote download failed")
class _MissingRemoteSandbox(_RemoteSandbox):
def download_file(self, path: str) -> bytes:
self.downloads.append(path)
try:
raise _E2BFileNotFound("not found")
except _E2BFileNotFound as error:
raise OSError(f"Failed to download file '{path}' from sandbox: not found") from error
class _BlockingRemoteSandbox(_RemoteSandbox):
def __init__(self, image_bytes: bytes, *, sandbox_id: str = "remote-1") -> None:
super().__init__(image_bytes, sandbox_id=sandbox_id)
self.download_started = threading.Event()
self.allow_download = threading.Event()
def download_file(self, path: str) -> bytes:
self.downloads.append(path)
self.download_started.set()
assert self.allow_download.wait(timeout=5)
return self.image_bytes
class _Provider:
def __init__(self, sandbox: _RemoteSandbox | None) -> None:
self.sandbox = sandbox
self.acquire_calls: list[tuple[str | None, str | None]] = []
self.release_calls: list[str] = []
def acquire(self, thread_id=None, *, user_id=None):
self.acquire_calls.append((thread_id, user_id))
raise AssertionError("view_image must not acquire a replacement sandbox")
def get(self, sandbox_id: str):
if self.sandbox is None:
return None
return self.sandbox if sandbox_id == self.sandbox.id else None
def release(self, sandbox_id: str) -> None:
self.release_calls.append(sandbox_id)
def _make_thread_data(tmp_path: Path) -> dict[str, str]:
user_data = tmp_path / "threads" / "thread-1" / "user-data"
workspace = user_data / "workspace"
uploads = user_data / "uploads"
outputs = user_data / "outputs"
for directory in (workspace, uploads, outputs):
directory.mkdir(parents=True)
return {
"workspace_path": str(workspace),
"uploads_path": str(uploads),
"outputs_path": str(outputs),
}
def _image_metadata(
actual_path: Path,
image_bytes: bytes,
*,
source_sandbox_id: str | None = None,
include_digest: bool = True,
) -> dict[str, object]:
metadata: dict[str, object] = {
"mime_type": "image/png",
"size": len(image_bytes),
"actual_path": str(actual_path),
}
if include_digest:
metadata["sha256"] = hashlib.sha256(image_bytes).hexdigest()
if source_sandbox_id is not None:
metadata["source_sandbox_id"] = source_sandbox_id
return metadata
def _make_runtime(
thread_data: dict[str, str],
*,
sandbox_id: str = "remote-1",
viewed_images: dict[str, dict[str, object]] | None = None,
) -> SimpleNamespace:
state: dict[str, object] = {
"thread_data": thread_data,
"sandbox": {"sandbox_id": sandbox_id},
}
if viewed_images is not None:
state["viewed_images"] = viewed_images
return SimpleNamespace(
state=state,
context={"thread_id": "thread-1"},
config={},
)
def _make_model_request(state: dict) -> ModelRequest:
assistant = AIMessage(
content="",
tool_calls=[
{
"name": "view_image",
"id": "call-view-image",
"args": {"image_path": "/mnt/user-data/outputs/plot.png"},
}
],
)
messages = [
assistant,
ToolMessage(content="Successfully read image", tool_call_id="call-view-image"),
]
return ModelRequest(
model=FakeMessagesListChatModel(responses=[AIMessage(content="ok")]),
messages=messages,
system_message=None,
tool_choice=None,
tools=[],
response_format=None,
state={"messages": messages, **state},
runtime=MagicMock(),
model_settings={},
)
def _message_content(result) -> str:
return result.update["messages"][0].content
def _image_bytes_from_blocks(blocks: list[str | dict]) -> bytes:
image_blocks = [block for block in blocks if isinstance(block, dict) and block.get("type") == "image_url"]
assert len(image_blocks) == 1
data_url = image_blocks[0]["image_url"]["url"]
prefix = "data:image/png;base64,"
assert data_url.startswith(prefix)
return base64.b64decode(data_url[len(prefix) :])
def _image_block_count(blocks: list[str | dict]) -> int:
return sum(1 for block in blocks if isinstance(block, dict) and block.get("type") == "image_url")
def test_view_image_reads_active_sandbox_when_host_mirror_is_missing(tmp_path, monkeypatch):
thread_data = _make_thread_data(tmp_path)
sandbox = _RemoteSandbox(PNG_BYTES)
provider = _Provider(sandbox)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
host_path = Path(thread_data["outputs_path"]) / "plot.png"
assert not host_path.exists()
result = view_image_tool.func(
runtime=_make_runtime(thread_data),
image_path="/mnt/user-data/outputs/plot.png",
tool_call_id="tc-remote",
)
assert _message_content(result) == "Successfully read image"
viewed = result.update["viewed_images"]["/mnt/user-data/outputs/plot.png"]
assert viewed["size"] == len(PNG_BYTES)
assert viewed["sha256"] == hashlib.sha256(PNG_BYTES).hexdigest()
assert viewed["source_sandbox_id"] == sandbox.id
assert sandbox.downloads == ["/mnt/user-data/outputs/plot.png"]
assert provider.acquire_calls == []
def test_view_image_prefers_active_sandbox_over_stale_host_mirror(tmp_path, monkeypatch):
thread_data = _make_thread_data(tmp_path)
host_path = Path(thread_data["outputs_path"]) / "plot.png"
host_path.write_bytes(PNG_BYTES)
remote_bytes = PNG_BYTES + b"remote-version"
sandbox = _RemoteSandbox(remote_bytes)
provider = _Provider(sandbox)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
result = view_image_tool.func(
runtime=_make_runtime(thread_data),
image_path="/mnt/user-data/outputs/plot.png",
tool_call_id="tc-stale",
)
assert _message_content(result) == "Successfully read image"
viewed = result.update["viewed_images"]["/mnt/user-data/outputs/plot.png"]
assert viewed["size"] == len(remote_bytes)
assert viewed["sha256"] == hashlib.sha256(remote_bytes).hexdigest()
assert viewed["source_sandbox_id"] == sandbox.id
assert sandbox.downloads == ["/mnt/user-data/outputs/plot.png"]
assert provider.acquire_calls == []
def test_view_image_falls_back_to_host_when_saved_sandbox_has_no_live_client(tmp_path, monkeypatch):
thread_data = _make_thread_data(tmp_path)
host_path = Path(thread_data["outputs_path"]) / "plot.png"
host_path.write_bytes(PNG_BYTES)
provider = _Provider(None)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
result = view_image_tool.func(
runtime=_make_runtime(thread_data),
image_path="/mnt/user-data/outputs/plot.png",
tool_call_id="tc-host-fallback",
)
assert _message_content(result) == "Successfully read image"
viewed = result.update["viewed_images"]["/mnt/user-data/outputs/plot.png"]
assert viewed["size"] == len(PNG_BYTES)
assert viewed["sha256"] == hashlib.sha256(PNG_BYTES).hexdigest()
assert "source_sandbox_id" not in viewed
assert provider.acquire_calls == []
def test_view_image_does_not_fall_back_after_live_sandbox_download_failure(tmp_path, monkeypatch):
thread_data = _make_thread_data(tmp_path)
host_path = Path(thread_data["outputs_path"]) / "plot.png"
host_path.write_bytes(PNG_BYTES)
sandbox = _FailingRemoteSandbox(PNG_BYTES)
provider = _Provider(sandbox)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
previous = {
"/mnt/user-data/outputs/plot.png": _image_metadata(
host_path,
PNG_BYTES,
source_sandbox_id="remote-old",
)
}
result = view_image_tool.func(
runtime=_make_runtime(thread_data, viewed_images=previous),
image_path="/mnt/user-data/outputs/plot.png",
tool_call_id="tc-live-failure",
)
assert _message_content(result).startswith("Error reading image file:")
assert sandbox.downloads == ["/mnt/user-data/outputs/plot.png"]
assert provider.acquire_calls == []
def test_view_image_recovers_verified_host_copy_when_replacement_sandbox_lacks_file(tmp_path, monkeypatch):
thread_data = _make_thread_data(tmp_path)
host_path = Path(thread_data["outputs_path"]) / "plot.png"
host_path.write_bytes(PNG_BYTES)
sandbox = _MissingRemoteSandbox(PNG_BYTES, sandbox_id="remote-new")
provider = _Provider(sandbox)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
previous = {
"/mnt/user-data/outputs/plot.png": _image_metadata(
host_path,
PNG_BYTES,
source_sandbox_id="remote-old",
)
}
result = view_image_tool.func(
runtime=_make_runtime(
thread_data,
sandbox_id=sandbox.id,
viewed_images=previous,
),
image_path="/mnt/user-data/outputs/plot.png",
tool_call_id="tc-replacement-host-fallback",
)
assert _message_content(result) == "Successfully read image"
viewed = result.update["viewed_images"]["/mnt/user-data/outputs/plot.png"]
assert viewed["sha256"] == hashlib.sha256(PNG_BYTES).hexdigest()
assert "source_sandbox_id" not in viewed
assert sandbox.downloads == ["/mnt/user-data/outputs/plot.png"]
def test_view_image_rejects_same_size_stale_host_copy_after_replacement(tmp_path, monkeypatch):
thread_data = _make_thread_data(tmp_path)
host_path = Path(thread_data["outputs_path"]) / "plot.png"
host_path.write_bytes(STALE_SAME_SIZE_PNG_BYTES)
sandbox = _MissingRemoteSandbox(PNG_BYTES, sandbox_id="remote-new")
provider = _Provider(sandbox)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
previous = {
"/mnt/user-data/outputs/plot.png": _image_metadata(
host_path,
PNG_BYTES,
source_sandbox_id="remote-old",
)
}
result = view_image_tool.func(
runtime=_make_runtime(
thread_data,
sandbox_id=sandbox.id,
viewed_images=previous,
),
image_path="/mnt/user-data/outputs/plot.png",
tool_call_id="tc-stale-host-rejected",
)
assert _message_content(result) == "Error: Image file not found: /mnt/user-data/outputs/plot.png"
assert sandbox.downloads == ["/mnt/user-data/outputs/plot.png"]
def test_view_image_same_generation_missing_file_stays_fail_closed(tmp_path, monkeypatch):
thread_data = _make_thread_data(tmp_path)
host_path = Path(thread_data["outputs_path"]) / "plot.png"
host_path.write_bytes(PNG_BYTES)
sandbox = _MissingRemoteSandbox(PNG_BYTES)
provider = _Provider(sandbox)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
previous = {
"/mnt/user-data/outputs/plot.png": _image_metadata(
host_path,
PNG_BYTES,
source_sandbox_id=sandbox.id,
)
}
result = view_image_tool.func(
runtime=_make_runtime(thread_data, viewed_images=previous),
image_path="/mnt/user-data/outputs/plot.png",
tool_call_id="tc-same-generation-missing",
)
assert _message_content(result) == "Error: Image file not found: /mnt/user-data/outputs/plot.png"
assert sandbox.downloads == ["/mnt/user-data/outputs/plot.png"]
def test_view_image_legacy_metadata_does_not_authorize_cross_generation_fallback(tmp_path, monkeypatch):
thread_data = _make_thread_data(tmp_path)
host_path = Path(thread_data["outputs_path"]) / "plot.png"
host_path.write_bytes(PNG_BYTES)
sandbox = _MissingRemoteSandbox(PNG_BYTES, sandbox_id="remote-new")
provider = _Provider(sandbox)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
previous = {
"/mnt/user-data/outputs/plot.png": _image_metadata(
host_path,
PNG_BYTES,
source_sandbox_id="remote-old",
include_digest=False,
)
}
result = view_image_tool.func(
runtime=_make_runtime(
thread_data,
sandbox_id=sandbox.id,
viewed_images=previous,
),
image_path="/mnt/user-data/outputs/plot.png",
tool_call_id="tc-legacy-no-digest",
)
assert _message_content(result) == "Error: Image file not found: /mnt/user-data/outputs/plot.png"
def test_middleware_injects_image_from_active_sandbox_without_host_copy(tmp_path, monkeypatch):
sandbox = _RemoteSandbox(PNG_BYTES)
provider = _Provider(sandbox)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
state = {
"sandbox": {"sandbox_id": sandbox.id},
"viewed_images": {
"/mnt/user-data/outputs/plot.png": _image_metadata(
tmp_path / "missing-host-copy.png",
PNG_BYTES,
source_sandbox_id=sandbox.id,
)
},
}
blocks = ViewImageMiddleware()._create_image_details_message(state)
assert _image_bytes_from_blocks(blocks) == PNG_BYTES
assert sandbox.downloads == ["/mnt/user-data/outputs/plot.png"]
def test_middleware_prefers_active_sandbox_over_stale_host_mirror(tmp_path, monkeypatch):
host_path = tmp_path / "stale-host-copy.png"
host_path.write_bytes(PNG_BYTES)
remote_bytes = PNG_BYTES + b"remote-version"
sandbox = _RemoteSandbox(remote_bytes)
provider = _Provider(sandbox)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
state = {
"sandbox": {"sandbox_id": sandbox.id},
"viewed_images": {
"/mnt/user-data/outputs/plot.png": _image_metadata(
host_path,
remote_bytes,
source_sandbox_id=sandbox.id,
)
},
}
blocks = ViewImageMiddleware()._create_image_details_message(state)
assert _image_bytes_from_blocks(blocks) == remote_bytes
assert sandbox.downloads == ["/mnt/user-data/outputs/plot.png"]
def test_middleware_falls_back_to_host_when_saved_sandbox_has_no_live_client(tmp_path, monkeypatch):
host_path = tmp_path / "synced-host-copy.png"
host_path.write_bytes(PNG_BYTES)
provider = _Provider(None)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
state = {
"sandbox": {"sandbox_id": "remote-1"},
"viewed_images": {
"/mnt/user-data/outputs/plot.png": _image_metadata(
host_path,
PNG_BYTES,
source_sandbox_id="remote-old",
)
},
}
blocks = ViewImageMiddleware()._create_image_details_message(state)
assert _image_bytes_from_blocks(blocks) == PNG_BYTES
assert provider.acquire_calls == []
def test_middleware_uses_verified_host_copy_after_sandbox_replacement(tmp_path, monkeypatch):
host_path = tmp_path / "synced-old-image.png"
host_path.write_bytes(PNG_BYTES)
sandbox = _MissingRemoteSandbox(PNG_BYTES, sandbox_id="remote-new")
provider = _Provider(sandbox)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
state = {
"sandbox": {"sandbox_id": sandbox.id},
"viewed_images": {
"/mnt/user-data/outputs/plot.png": _image_metadata(
host_path,
PNG_BYTES,
source_sandbox_id="remote-old",
)
},
}
blocks = ViewImageMiddleware()._create_image_details_message(state)
assert _image_bytes_from_blocks(blocks) == PNG_BYTES
assert sandbox.downloads == []
def test_middleware_rejects_same_size_stale_host_after_sandbox_replacement(tmp_path, monkeypatch):
host_path = tmp_path / "stale-same-size.png"
host_path.write_bytes(STALE_SAME_SIZE_PNG_BYTES)
sandbox = _MissingRemoteSandbox(PNG_BYTES, sandbox_id="remote-new")
provider = _Provider(sandbox)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
state = {
"sandbox": {"sandbox_id": sandbox.id},
"viewed_images": {
"/mnt/user-data/outputs/plot.png": _image_metadata(
host_path,
PNG_BYTES,
source_sandbox_id="remote-old",
)
},
}
blocks = ViewImageMiddleware()._create_image_details_message(state)
assert _image_block_count(blocks) == 0
assert sandbox.downloads == ["/mnt/user-data/outputs/plot.png"]
def test_middleware_same_generation_failure_does_not_use_host_copy(tmp_path, monkeypatch):
host_path = tmp_path / "matching-host-copy.png"
host_path.write_bytes(PNG_BYTES)
sandbox = _FailingRemoteSandbox(PNG_BYTES)
provider = _Provider(sandbox)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
state = {
"sandbox": {"sandbox_id": sandbox.id},
"viewed_images": {
"/mnt/user-data/outputs/plot.png": _image_metadata(
host_path,
PNG_BYTES,
source_sandbox_id=sandbox.id,
)
},
}
blocks = ViewImageMiddleware()._create_image_details_message(state)
assert _image_block_count(blocks) == 0
assert sandbox.downloads == ["/mnt/user-data/outputs/plot.png"]
def test_middleware_sync_read_finishes_before_lease_release(tmp_path, monkeypatch):
sandbox = _BlockingRemoteSandbox(PNG_BYTES)
provider = _Provider(sandbox)
manager = SandboxLeaseManager(provider)
manager.retain(
"run-owner",
sandbox.id,
thread_id="thread-1",
user_id="user-1",
)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
request = _make_model_request(
{
"sandbox": {"sandbox_id": sandbox.id},
"viewed_images": {
"/mnt/user-data/outputs/plot.png": _image_metadata(
tmp_path / "missing-host-copy.png",
PNG_BYTES,
source_sandbox_id=sandbox.id,
)
},
}
)
handler_called = threading.Event()
invocation_errors: list[BaseException] = []
def handler(_prepared: ModelRequest) -> AIMessage:
handler_called.set()
return AIMessage(content="ok")
def invoke_under_lease() -> None:
try:
ViewImageMiddleware().wrap_model_call(request, handler)
except BaseException as error:
invocation_errors.append(error)
finally:
manager.release("run-owner")
worker = threading.Thread(target=invoke_under_lease)
try:
worker.start()
assert sandbox.download_started.wait(timeout=1)
assert worker.is_alive()
assert not handler_called.is_set()
assert provider.release_calls == []
assert sandbox.released_scopes == []
sandbox.allow_download.set()
worker.join(timeout=2)
assert not worker.is_alive()
assert invocation_errors == []
assert handler_called.is_set()
assert provider.release_calls == [sandbox.id]
assert sandbox.released_scopes == ["run-owner"]
finally:
sandbox.allow_download.set()
worker.join(timeout=2)
if manager.binding_for("run-owner") is not None:
manager.release("run-owner")
manager.close()
@pytest.mark.asyncio
async def test_middleware_cancellation_drains_sandbox_download_before_lease_release(tmp_path, monkeypatch):
sandbox = _BlockingRemoteSandbox(PNG_BYTES)
provider = _Provider(sandbox)
manager = SandboxLeaseManager(provider)
manager.retain(
"run-owner",
sandbox.id,
thread_id="thread-1",
user_id="user-1",
)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
request = _make_model_request(
{
"sandbox": {"sandbox_id": sandbox.id},
"viewed_images": {
"/mnt/user-data/outputs/plot.png": _image_metadata(
tmp_path / "missing-host-copy.png",
PNG_BYTES,
source_sandbox_id=sandbox.id,
)
},
}
)
handler_called = False
async def handler(_prepared: ModelRequest) -> AIMessage:
nonlocal handler_called
handler_called = True
return AIMessage(content="unexpected")
async def invoke_under_lease():
try:
return await ViewImageMiddleware().awrap_model_call(request, handler)
finally:
await manager.release_async("run-owner")
task = asyncio.create_task(invoke_under_lease())
try:
assert await asyncio.to_thread(sandbox.download_started.wait, 1)
task.cancel()
await asyncio.sleep(0)
assert not task.done()
assert provider.release_calls == []
assert sandbox.released_scopes == []
sandbox.allow_download.set()
with pytest.raises(asyncio.CancelledError):
await task
assert handler_called is False
assert provider.release_calls == [sandbox.id]
assert sandbox.released_scopes == ["run-owner"]
finally:
sandbox.allow_download.set()
if not task.done():
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
manager.close()

134
backend/tests/test_view_image_tool_cancellation.py Normal file
View File

@ -0,0 +1,134 @@
import asyncio
import base64
import threading
from pathlib import Path
import pytest
from langchain.tools import ToolRuntime
from deerflow.sandbox.lease import SandboxLeaseManager
from deerflow.tools.builtins.view_image_tool import view_image_tool
PNG_BYTES = base64.b64decode("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==")
class _BlockingRemoteSandbox:
id = "remote-1"
def __init__(self) -> None:
self.download_started = threading.Event()
self.allow_download = threading.Event()
self.downloads: list[str] = []
self.released_scopes: list[str] = []
def download_file(self, path: str) -> bytes:
self.downloads.append(path)
self.download_started.set()
assert self.allow_download.wait(timeout=5)
return PNG_BYTES
def release_command_scope(self, scope_id: str) -> None:
self.released_scopes.append(scope_id)
class _Provider:
def __init__(self, sandbox: _BlockingRemoteSandbox) -> None:
self.sandbox = sandbox
self.release_calls: list[str] = []
def get(self, sandbox_id: str):
return self.sandbox if sandbox_id == self.sandbox.id else None
def release(self, sandbox_id: str) -> None:
self.release_calls.append(sandbox_id)
def _thread_data(tmp_path: Path) -> dict[str, str]:
user_data = tmp_path / "threads" / "thread-1" / "user-data"
workspace = user_data / "workspace"
uploads = user_data / "uploads"
outputs = user_data / "outputs"
for directory in (workspace, uploads, outputs):
directory.mkdir(parents=True)
return {
"workspace_path": str(workspace),
"uploads_path": str(uploads),
"outputs_path": str(outputs),
}
def _runtime(tmp_path: Path, sandbox_id: str) -> ToolRuntime:
return ToolRuntime(
state={
"thread_data": _thread_data(tmp_path),
"sandbox": {"sandbox_id": sandbox_id},
},
context={"thread_id": "thread-1"},
config={"configurable": {"thread_id": "thread-1"}},
stream_writer=lambda _: None,
tools=[],
tool_call_id="tc-tool-cancel",
store=None,
)
@pytest.mark.asyncio
async def test_view_image_ainvoke_drains_download_before_lease_release(tmp_path, monkeypatch):
sandbox = _BlockingRemoteSandbox()
provider = _Provider(sandbox)
manager = SandboxLeaseManager(provider)
manager.retain(
"run-owner",
sandbox.id,
thread_id="thread-1",
user_id="user-1",
)
monkeypatch.setattr(
"deerflow.sandbox.sandbox_provider.get_sandbox_provider",
lambda: provider,
)
runtime = _runtime(tmp_path, sandbox.id)
async def invoke_under_lease():
try:
return await view_image_tool.ainvoke(
{
"args": {
"runtime": runtime,
"image_path": "/mnt/user-data/outputs/plot.png",
},
"name": "view_image",
"type": "tool_call",
"id": "tc-tool-cancel",
}
)
finally:
await manager.release_async("run-owner")
task = asyncio.create_task(invoke_under_lease())
try:
assert await asyncio.to_thread(sandbox.download_started.wait, 1)
task.cancel()
await asyncio.sleep(0)
assert not task.done()
assert provider.release_calls == []
assert sandbox.released_scopes == []
sandbox.allow_download.set()
with pytest.raises(asyncio.CancelledError):
await task
assert sandbox.downloads == ["/mnt/user-data/outputs/plot.png"]
assert provider.release_calls == [sandbox.id]
assert sandbox.released_scopes == ["run-owner"]
finally:
sandbox.allow_download.set()
if not task.done():
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
if manager.binding_for("run-owner") is not None:
await manager.release_async("run-owner")
manager.close()