EvoScientist/tests/test_read_file_images.py

"""Tests for the read_file image post-processing pipeline."""

from __future__ import annotations

import base64
import io

import pytest
from langchain_core.messages import ToolMessage
from PIL import Image

from EvoScientist.middleware.read_file_images import (
    MAX_IMAGE_EDGE,
    install_read_file_image_patch,
    process_image_tool_message,
)


def _b64_image(fmt: str, size: tuple[int, int], *, mode: str = "RGB") -> str:
    image = Image.new(mode, size, color=(120, 30, 200))
    buffer = io.BytesIO()
    image.save(buffer, fmt)
    return base64.standard_b64encode(buffer.getvalue()).decode()


def _image_message(
    b64: str,
    *,
    media_type: str = "image/png",
    path: str = "/artifacts/x.png",
    block_type: str = "image",
) -> ToolMessage:
    return ToolMessage(
        content=[{"type": block_type, "base64": b64, "mime_type": media_type}],
        name="read_file",
        tool_call_id="call-1",
        status="success",
        additional_kwargs={"read_file_path": path, "read_file_media_type": media_type},
    )


def _decode_block(message: ToolMessage) -> tuple[dict, Image.Image]:
    block = message.content[-1]
    raw = base64.standard_b64decode(block["base64"])
    image = Image.open(io.BytesIO(raw))
    image.load()
    return block, image


class TestPassthrough:
    def test_small_correct_type_keeps_bytes_adds_metadata(self):
        b64 = _b64_image("PNG", (100, 80))
        message = _image_message(b64)
        result = process_image_tool_message(message)

        assert result is not message
        assert result.status == "success"
        meta, block = result.content
        assert meta["type"] == "text"
        assert "100x80" in meta["text"]
        assert "PNG" in meta["text"]
        assert "downsampled" not in meta["text"]
        # Original bytes preserved — no re-encode.
        assert block["base64"] == b64
        assert block["mime_type"] == "image/png"
        assert result.additional_kwargs["read_file_media_type"] == "image/png"

    def test_svg_untouched(self):
        message = _image_message(
            base64.standard_b64encode(b"<svg/>").decode(),
            media_type="image/svg+xml",
            path="/artifacts/icon.svg",
        )
        assert process_image_tool_message(message) is message

    def test_text_result_untouched(self):
        message = ToolMessage(
            content="1 | hello",
            name="read_file",
            tool_call_id="call-1",
            status="success",
        )
        assert process_image_tool_message(message) is message

    def test_error_result_untouched(self):
        message = ToolMessage(
            content="Error: not found",
            name="read_file",
            tool_call_id="call-1",
            status="error",
        )
        assert process_image_tool_message(message) is message

    def test_non_image_binary_untouched(self):
        message = _image_message(
            base64.standard_b64encode(b"%PDF-1.4 fake").decode(),
            media_type="application/pdf",
            path="/artifacts/paper.pdf",
            block_type="file",
        )
        assert process_image_tool_message(message) is message


class TestDownsample:
    def test_oversized_rgb_becomes_jpeg_with_metadata(self):
        message = _image_message(_b64_image("PNG", (3000, 1200)))
        result = process_image_tool_message(message)

        assert result.status == "success"
        assert len(result.content) == 2
        meta, block = result.content
        assert meta["type"] == "text"
        assert "3000x1200" in meta["text"]
        assert "PNG" in meta["text"]
        assert "downsampled" in meta["text"]
        assert str(MAX_IMAGE_EDGE) in meta["text"]

        assert block["type"] == "image"
        assert block["mime_type"] == "image/jpeg"
        assert result.additional_kwargs["read_file_media_type"] == "image/jpeg"

        _, image = _decode_block(result)
        assert max(image.size) <= MAX_IMAGE_EDGE
        assert image.format == "JPEG"
        # Aspect ratio preserved: 3000x1200 -> 2048x819.
        assert image.size == (2048, 819)

    def test_oversized_alpha_png_stays_png(self):
        message = _image_message(_b64_image("PNG", (2500, 2500), mode="RGBA"))
        result = process_image_tool_message(message)

        block, image = _decode_block(result)
        assert block["mime_type"] == "image/png"
        assert image.format == "PNG"
        assert image.mode == "RGBA"
        assert max(image.size) <= MAX_IMAGE_EDGE

    def test_unscaled_returns_exact_cap_when_equal(self):
        # Exactly at the cap: no scaling — bytes kept, metadata attached.
        b64 = _b64_image("PNG", (MAX_IMAGE_EDGE, 10))
        result = process_image_tool_message(_image_message(b64))
        assert result.content[-1]["base64"] == b64
        assert "downsampled" not in result.content[0]["text"]


class TestMimeSniffing:
    def test_extension_lie_corrected(self):
        # JPEG bytes in a .png-named file: declared image/png, sniffed JPEG.
        message = _image_message(
            _b64_image("JPEG", (200, 200)), media_type="image/png"
        )
        result = process_image_tool_message(message)

        block, image = _decode_block(result)
        assert block["mime_type"] == "image/jpeg"
        assert result.additional_kwargs["read_file_media_type"] == "image/jpeg"
        assert image.format == "JPEG"
        meta = result.content[0]
        assert "re-encoded" in meta["text"]

    def test_gif_first_frame(self):
        # Small GIFs keep their bytes (plus metadata); an oversized GIF
        # exercises the transcode path (first frame -> JPEG).
        small_b64 = _b64_image("GIF", (300, 100))
        small = _image_message(
            small_b64,
            media_type="image/gif",
            path="/artifacts/anim.gif",
        )
        small_result = process_image_tool_message(small)
        assert small_result.content[-1]["base64"] == small_b64
        assert small_result.content[-1]["mime_type"] == "image/gif"

        message = _image_message(
            _b64_image("GIF", (3000, 100)),
            media_type="image/gif",
            path="/artifacts/anim.gif",
        )
        result = process_image_tool_message(message)

        block, image = _decode_block(result)
        assert block["mime_type"] == "image/jpeg"
        assert image.format == "JPEG"
        assert max(image.size) <= MAX_IMAGE_EDGE


class TestErrors:
    def test_corrupt_image_is_explicit_error(self):
        message = _image_message(base64.standard_b64encode(b"not an image").decode())
        result = process_image_tool_message(message)

        assert result.status == "error"
        assert "无法读取图像数据" in result.content
        assert "/artifacts/x.png" in result.content
        assert result.tool_call_id == "call-1"

    def test_invalid_base64_is_explicit_error(self):
        message = _image_message("!!!not-base64!!!")
        result = process_image_tool_message(message)

        assert result.status == "error"
        assert "无法读取图像数据" in result.content


class TestOpenAICompatDelivery:
    """The model must actually receive the image on OpenAI-compatible
    providers: tool-role messages only carry text there, so the image has to
    be hoisted into a following user message with the metadata text staying
    on the tool result."""

    def test_processed_message_hoists_image_keeps_metadata(self):
        from EvoScientist.llm.patches import _sanitize_messages

        message = process_image_tool_message(
            _image_message(_b64_image("PNG", (100, 80)))
        )
        out = _sanitize_messages([message], hoist_tool_media=True)

        assert len(out) == 2
        tool_msg, hoisted = out
        # Tool result keeps an informative text body — never empty.
        assert tool_msg.type == "tool"
        assert isinstance(tool_msg.content, str)
        assert "100x80" in tool_msg.content
        # The image rides a following user message.
        assert hoisted.type == "human"
        assert any(
            isinstance(b, dict) and b.get("type") == "image" for b in hoisted.content
        )

    def test_openai_payload_is_protocol_legal(self):
        from langchain_openai.chat_models.base import _convert_message_to_dict

        from EvoScientist.llm.patches import _sanitize_messages

        message = process_image_tool_message(
            _image_message(_b64_image("PNG", (100, 80)))
        )
        out = _sanitize_messages([message], hoist_tool_media=True)
        tool_dict = _convert_message_to_dict(out[0])
        user_dict = _convert_message_to_dict(out[1])

        # tool role: plain string content only.
        assert tool_dict["role"] == "tool"
        assert isinstance(tool_dict["content"], str)
        # user role: image_url data URI.
        assert user_dict["role"] == "user"
        assert any(
            isinstance(b, dict) and b.get("type") == "image_url"
            for b in user_dict["content"]
        )

    def test_chat_openai_builder_applies_hoist_patch(self):
        from EvoScientist.model_registry.factory import _build_chat_openai

        model = _build_chat_openai(
            {"model": "qwen3.7-plus", "api_key": "sk-test"}, None, None
        )
        # functools.wraps marks the wrapper; unpatched ChatOpenAI has none.
        assert getattr(model._generate, "__wrapped__", None) is not None
        assert getattr(model._astream, "__wrapped__", None) is not None


class TestInstall:
    def test_install_is_idempotent(self):
        from deepagents.middleware.filesystem import FilesystemMiddleware

        first = FilesystemMiddleware._create_read_file_tool
        install_read_file_image_patch()
        install_read_file_image_patch()
        assert FilesystemMiddleware._create_read_file_tool is first

    @pytest.mark.asyncio
    async def test_wrapped_tool_processes_images(self, tmp_path):
        from deepagents.backends.filesystem import FilesystemBackend
        from deepagents.middleware.filesystem import FilesystemMiddleware

        install_read_file_image_patch()
        big = tmp_path / "big.png"
        Image.new("RGB", (3000, 500), color=(1, 2, 3)).save(big, "PNG")

        middleware = FilesystemMiddleware(
            backend=FilesystemBackend(root_dir=str(tmp_path), virtual_mode=True)
        )
        tool = next(t for t in middleware.tools if t.name == "read_file")

        class _Runtime:
            tool_call_id = "call-9"

        result = await tool.coroutine(
            file_path="/big.png", runtime=_Runtime(), offset=0, limit=2000
        )

        assert result.status == "success"
        assert isinstance(result.content, list) and len(result.content) == 2
        assert result.content[0]["type"] == "text"
        assert "3000x500" in result.content[0]["text"]
        _, image = _decode_block(result)
        assert max(image.size) <= MAX_IMAGE_EDGE

    def test_wrapped_tool_keeps_text_behavior(self, tmp_path):
        from deepagents.backends.filesystem import FilesystemBackend
        from deepagents.middleware.filesystem import FilesystemMiddleware

        install_read_file_image_patch()
        (tmp_path / "notes.txt").write_text("alpha\nbeta\n")

        middleware = FilesystemMiddleware(
            backend=FilesystemBackend(root_dir=str(tmp_path), virtual_mode=True)
        )
        tool = next(t for t in middleware.tools if t.name == "read_file")

        class _Runtime:
            tool_call_id = "call-10"

        result = tool.func(
            file_path="/notes.txt", runtime=_Runtime(), offset=0, limit=2000
        )

        assert result.status == "success"
        assert "alpha" in result.content
        assert "beta" in result.content

    def test_wrapped_tool_preserves_runtime_injection(self, tmp_path):
        """Regression: *args/**kwargs wrappers hid the runtime parameter, so
        StructuredTool/ToolNode never injected it and every call raised
        "missing 1 required positional argument: 'runtime'"."""
        import inspect

        from deepagents.backends.filesystem import FilesystemBackend
        from deepagents.middleware.filesystem import FilesystemMiddleware
        from langchain_core.tools.base import _is_injected_arg_type

        install_read_file_image_patch()
        middleware = FilesystemMiddleware(
            backend=FilesystemBackend(root_dir=str(tmp_path), virtual_mode=True)
        )
        tool = next(t for t in middleware.tools if t.name == "read_file")

        for fn in (tool.func, tool.coroutine):
            params = inspect.signature(fn).parameters
            assert "runtime" in params
            assert _is_injected_arg_type(params["runtime"].annotation)
        # The exact mechanism langchain uses to decide injection.
        assert "runtime" in tool._injected_args_keys