deer-flow/backend/tests/test_view_image_remote_sandbox.py

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