EvoScientist-Multi/tests/test_provider_context_middleware.py

from __future__ import annotations

import base64
from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest
from langchain.agents import create_agent
from langchain.agents.middleware.types import ModelRequest, ModelResponse
from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.types import Overwrite

from EvoScientist.llm.contracts import EvoRuntimeError
from EvoScientist.middleware.provider_context import ProviderContextMediaMiddleware


class _Backend:
    def __init__(self, *, error: str | None = None) -> None:
        self.error = error
        self.uploads: list[tuple[str, bytes]] = []

    def upload_files(self, files):
        self.uploads.extend(files)
        return [SimpleNamespace(path=path, error=self.error) for path, _data in files]

    async def aupload_files(self, files):
        return self.upload_files(files)


class _FakeModel(FakeMessagesListChatModel):
    def bind_tools(self, _tools, *, tool_choice=None, **_kwargs):
        return self


def _request(messages):
    return ModelRequest(
        messages=list(messages),
        model=MagicMock(),
        state={},
        runtime=MagicMock(),
        system_message=MagicMock(),
    )


def _image_block(raw: bytes) -> dict:
    return {
        "type": "image",
        "base64": base64.b64encode(raw).decode("ascii"),
        "mime_type": "image/png",
    }


def test_externalizes_historical_and_new_assistant_media() -> None:
    backend = _Backend()
    middleware = ProviderContextMediaMiddleware(backend)
    old_raw = b"old-png"
    new_raw = b"new-png"
    historical = AIMessage(content=[_image_block(old_raw)])
    captured = {}

    def handler(request):
        captured["messages"] = request.messages
        return ModelResponse(
            result=[
                AIMessage(
                    content=[
                        {"type": "text", "text": "done"},
                        _image_block(new_raw),
                    ]
                )
            ]
        )

    response = middleware.wrap_model_call(_request([historical]), handler)

    provider_content = captured["messages"][0].content
    assert all("base64" not in block for block in provider_content)
    assert "generated_image" in provider_content[0]["text"]
    stored_content = response.result[0].content
    assert stored_content[1]["type"] == "image"
    assert stored_content[1]["url"].startswith("/artifacts/model-output/")
    assert "base64" not in stored_content[1]
    assert {data for _path, data in backend.uploads} == {old_raw, new_raw}


def test_before_model_durably_replaces_historical_inline_media() -> None:
    middleware = ProviderContextMediaMiddleware(_Backend())

    update = middleware.before_model(
        {"messages": [AIMessage(content=[_image_block(b"old-png")])]}, None
    )

    assert update is not None
    assert isinstance(update["messages"], Overwrite)
    content = update["messages"].value[0].content
    assert content[0]["url"].startswith("/artifacts/model-output/")
    assert "base64" not in content[0]


@pytest.mark.asyncio
async def test_real_langgraph_injects_runtime_and_repairs_checkpoint_media() -> None:
    backend = _Backend()
    raw = b"x" * 1_349_952
    agent = create_agent(
        model=_FakeModel(responses=[AIMessage(content="done")]),
        tools=[],
        middleware=[ProviderContextMediaMiddleware(backend)],
    )

    result = await agent.ainvoke(
        {
            "messages": [
                AIMessage(content=[_image_block(raw)]),
                HumanMessage(content="continue"),
            ]
        }
    )

    repaired = result["messages"][0].content[0]
    assert repaired["url"].startswith("/artifacts/model-output/")
    assert "base64" not in repaired
    assert any(data == raw for _path, data in backend.uploads)


@pytest.mark.asyncio
async def test_before_model_normalizes_internal_middleware_failure() -> None:
    class _BrokenBackend(_Backend):
        async def aupload_files(self, _files):
            raise TypeError("sensitive internal detail")

    middleware = ProviderContextMediaMiddleware(_BrokenBackend())

    with pytest.raises(EvoRuntimeError, match="AGENT_MIDDLEWARE_FAILED") as exc:
        await middleware.abefore_model(
            {"messages": [AIMessage(content=[_image_block(b"png")])]}, None
        )

    assert exc.value.details == (
        {
            "failure_stage": "agent_middleware",
            "middleware": "provider_context_media",
            "middleware_node": "provider_context_media.before_model",
            "agent_error_type": "TypeError",
            "agent_error_module": "builtins",
        },
    )
    assert "sensitive internal detail" not in str(exc.value.details)


@pytest.mark.asyncio
async def test_async_externalization_is_content_addressed() -> None:
    backend = _Backend()
    middleware = ProviderContextMediaMiddleware(backend)
    raw = b"same-png"

    async def handler(_request):
        return ModelResponse(result=[AIMessage(content=[_image_block(raw)])])

    first = await middleware.awrap_model_call(_request([]), handler)
    second = await middleware.awrap_model_call(_request(first.result), handler)

    assert first.result[0].content[0]["url"] == second.result[0].content[0]["url"]
    assert all(len(data) < 100 for _path, data in backend.uploads)


def test_upload_failure_is_terminal_and_does_not_return_inline_media() -> None:
    middleware = ProviderContextMediaMiddleware(_Backend(error="disk full"))

    with pytest.raises(EvoRuntimeError, match="MEDIA_PERSIST_FAILED"):
        middleware.wrap_model_call(
            _request([]),
            lambda _request: ModelResponse(
                result=[AIMessage(content=[_image_block(b"png")])]
            ),
        )