From 3a6e681dee58a942edcaca67077c08aadc673c95 Mon Sep 17 00:00:00 2001 From: Jun <84921700+Amazingjun-j@users.noreply.github.com> Date: Thu, 10 Sep 2026 15:48:55 +0800 Subject: [PATCH] fix(view-image): read active sandbox images from sandbox (#5306) * 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 --- .../middlewares/view_image_middleware.py | 196 ++++- .../harness/deerflow/agents/thread_state.py | 10 +- .../packages/harness/deerflow/tools/AGENTS.md | 4 +- .../tools/builtins/view_image_tool.py | 233 ++++-- ...iew_image_provider_error_classification.py | 130 ++++ .../tests/test_view_image_remote_sandbox.py | 704 ++++++++++++++++++ .../test_view_image_tool_cancellation.py | 134 ++++ 7 files changed, 1319 insertions(+), 92 deletions(-) create mode 100644 backend/tests/test_view_image_provider_error_classification.py create mode 100644 backend/tests/test_view_image_remote_sandbox.py create mode 100644 backend/tests/test_view_image_tool_cancellation.py diff --git a/backend/packages/harness/deerflow/agents/middlewares/view_image_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/view_image_middleware.py index 3f148cd54..de1c61853 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/view_image_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/view_image_middleware.py @@ -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) diff --git a/backend/packages/harness/deerflow/agents/thread_state.py b/backend/packages/harness/deerflow/agents/thread_state.py index a40238e54..b435f4389 100644 --- a/backend/packages/harness/deerflow/agents/thread_state.py +++ b/backend/packages/harness/deerflow/agents/thread_state.py @@ -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: diff --git a/backend/packages/harness/deerflow/tools/AGENTS.md b/backend/packages/harness/deerflow/tools/AGENTS.md index 1518036f4..57e6c4db5 100644 --- a/backend/packages/harness/deerflow/tools/AGENTS.md +++ b/backend/packages/harness/deerflow/tools/AGENTS.md @@ -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` \ No newline at end of file diff --git a/backend/packages/harness/deerflow/tools/builtins/view_image_tool.py b/backend/packages/harness/deerflow/tools/builtins/view_image_tool.py index 33d208fb6..6368cedac 100644 --- a/backend/packages/harness/deerflow/tools/builtins/view_image_tool.py +++ b/backend/packages/harness/deerflow/tools/builtins/view_image_tool.py @@ -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, +) diff --git a/backend/tests/test_view_image_provider_error_classification.py b/backend/tests/test_view_image_provider_error_classification.py new file mode 100644 index 000000000..192ded258 --- /dev/null +++ b/backend/tests/test_view_image_provider_error_classification.py @@ -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__ diff --git a/backend/tests/test_view_image_remote_sandbox.py b/backend/tests/test_view_image_remote_sandbox.py new file mode 100644 index 000000000..0e4bb0244 --- /dev/null +++ b/backend/tests/test_view_image_remote_sandbox.py @@ -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() diff --git a/backend/tests/test_view_image_tool_cancellation.py b/backend/tests/test_view_image_tool_cancellation.py new file mode 100644 index 000000000..d4e82179f --- /dev/null +++ b/backend/tests/test_view_image_tool_cancellation.py @@ -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()