EvoScientist-Multi/tests/test_web_search.py

from __future__ import annotations

import asyncio
import socket

import httpx
import pytest


@pytest.mark.anyio
async def test_fetch_webpage_rejects_loopback_before_connecting(monkeypatch):
    from EvoScientist.tools import search

    connected = asyncio.Event()

    async def serve(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        connected.set()
        await reader.read(4096)
        writer.write(
            b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\n"
            b"Content-Length: 6\r\nConnection: close\r\n\r\nsecret"
        )
        await writer.drain()
        writer.close()
        await writer.wait_closed()

    server = await asyncio.start_server(serve, "127.0.0.1", 0)
    port = server.sockets[0].getsockname()[1]
    monkeypatch.setitem(search._ALLOWED_WEB_PORTS, "http", port)
    try:
        result = await search.fetch_webpage_content(
            f"http://127.0.0.1:{port}/private"
        )
    finally:
        server.close()
        await server.wait_closed()

    assert result.startswith("Error fetching content from ")
    assert "not publicly routable" in result
    assert not connected.is_set()


@pytest.mark.anyio
async def test_fetch_webpage_connects_to_the_validated_ip_not_rebound_dns(monkeypatch):
    from EvoScientist.tools import search

    original_getaddrinfo = socket.getaddrinfo
    captured: dict[str, object] = {}

    def rebind_getaddrinfo(host, requested_port, *args, **kwargs):
        if host != "rebinding.example":
            return original_getaddrinfo(host, requested_port, *args, **kwargs)
        return [(2, 1, 6, "", ("93.184.216.34", requested_port))]

    async def send(client, request, **_kwargs):
        captured["trust_env"] = client._trust_env
        captured["host"] = request.url.host
        captured["host_header"] = request.headers["host"]
        captured["sni_hostname"] = request.extensions.get("sni_hostname")

        class PublicStream(httpx.AsyncByteStream):
            async def __aiter__(self):
                yield b"public"

        return httpx.Response(
            200,
            headers={"content-type": "text/plain"},
            stream=PublicStream(),
            request=request,
        )

    monkeypatch.setattr(socket, "getaddrinfo", rebind_getaddrinfo)
    monkeypatch.setattr(httpx.AsyncClient, "send", send)

    result = await search.fetch_webpage_content("https://rebinding.example/public")

    assert result == "public"
    assert captured == {
        "trust_env": False,
        "host": "93.184.216.34",
        "host_header": "rebinding.example",
        "sni_hostname": "rebinding.example",
    }


@pytest.mark.anyio
@pytest.mark.parametrize(
    "url",
    [
        "ftp://example.com/file",
        "https://user@example.com/file",
        "https://example.com:8443/file",
        "https://[::1]/file",
        "http://0.0.0.0/file",
        "http://224.0.0.1/file",
        "http://192.0.2.1/file",
    ],
)
async def test_fetch_webpage_rejects_unsafe_url_targets(url):
    from EvoScientist.tools.search import fetch_webpage_content

    result = await fetch_webpage_content(url)

    assert result.startswith("Error fetching content from ")


@pytest.mark.anyio
async def test_fetch_webpage_rejects_hostname_when_any_dns_answer_is_private(
    monkeypatch,
):
    from EvoScientist.tools import search

    def mixed_dns(host, port, *_args, **_kwargs):
        assert host == "mixed.example"
        return [
            (2, 1, 6, "", ("93.184.216.34", port)),
            (2, 1, 6, "", ("127.0.0.1", port)),
        ]

    monkeypatch.setattr(socket, "getaddrinfo", mixed_dns)

    result = await search.fetch_webpage_content("https://mixed.example/")

    assert "not publicly routable" in result


@pytest.mark.anyio
async def test_fetch_webpage_revalidates_redirect_targets(monkeypatch):
    from EvoScientist.tools import search

    requests: list[httpx.Request] = []

    def public_dns(host, port, *_args, **_kwargs):
        assert host == "public.example"
        return [(2, 1, 6, "", ("93.184.216.34", port))]

    async def send(_client, request, **_kwargs):
        requests.append(request)
        return httpx.Response(
            302,
            headers={"location": "http://127.0.0.1/private"},
            request=request,
        )

    monkeypatch.setattr(socket, "getaddrinfo", public_dns)
    monkeypatch.setattr(httpx.AsyncClient, "send", send)

    result = await search.fetch_webpage_content("https://public.example/start")

    assert "not publicly routable" in result
    assert len(requests) == 1


@pytest.mark.anyio
async def test_cross_host_redirect_uses_fresh_tls_connection_pool(monkeypatch):
    from EvoScientist.tools import search

    client_ids: list[int] = []

    class PublicStream(httpx.AsyncByteStream):
        async def __aiter__(self):
            yield b"second host"

    def same_ip_dns(host, port, *_args, **_kwargs):
        assert host in {"first.example", "second.example"}
        return [(2, 1, 6, "", ("93.184.216.34", port))]

    async def send(client, request, **_kwargs):
        client_ids.append(id(client))
        if request.headers["host"] == "first.example":
            return httpx.Response(
                302,
                headers={"location": "https://second.example/page"},
                request=request,
            )
        return httpx.Response(
            200,
            headers={"content-type": "text/plain"},
            stream=PublicStream(),
            request=request,
        )

    monkeypatch.setattr(socket, "getaddrinfo", same_ip_dns)
    monkeypatch.setattr(httpx.AsyncClient, "send", send)

    result = await search.fetch_webpage_content("https://first.example/start")

    assert result == "second host"
    assert len(client_ids) == 2
    assert client_ids[0] != client_ids[1]


@pytest.mark.anyio
async def test_fetch_webpage_stops_streaming_at_byte_limit(monkeypatch):
    from EvoScientist.tools import search

    chunks_read = 0

    class CountingStream(httpx.AsyncByteStream):
        async def __aiter__(self):
            nonlocal chunks_read
            for _ in range(100):
                chunks_read += 1
                yield b"a" * 100

    def public_dns(_host, port, *_args, **_kwargs):
        return [(2, 1, 6, "", ("93.184.216.34", port))]

    async def send(_client, request, **_kwargs):
        return httpx.Response(
            200,
            headers={"content-type": "text/plain"},
            stream=CountingStream(),
            request=request,
        )

    monkeypatch.setattr(search, "MAX_PAGE_RESPONSE_BYTES", 250)
    monkeypatch.setattr(socket, "getaddrinfo", public_dns)
    monkeypatch.setattr(httpx.AsyncClient, "send", send)

    result = await search.fetch_webpage_content("https://public.example/page")

    assert result == "a" * 250
    assert chunks_read <= 4


@pytest.mark.anyio
async def test_fetch_webpage_rejects_compressed_response_without_decoding(
    monkeypatch,
):
    from EvoScientist.tools import search

    raw_reads = 0

    class CompressedStream(httpx.AsyncByteStream):
        async def __aiter__(self):
            nonlocal raw_reads
            raw_reads += 1
            yield b"compressed-body-must-not-be-decoded"

    def public_dns(_host, port, *_args, **_kwargs):
        return [(2, 1, 6, "", ("93.184.216.34", port))]

    async def send(_client, request, **_kwargs):
        assert request.headers["accept-encoding"] == "identity"
        return httpx.Response(
            200,
            headers={"content-type": "text/plain", "content-encoding": "gzip"},
            stream=CompressedStream(),
            request=request,
        )

    monkeypatch.setattr(socket, "getaddrinfo", public_dns)
    monkeypatch.setattr(httpx.AsyncClient, "send", send)

    result = await search.fetch_webpage_content("https://public.example/page")

    assert "compressed responses are not allowed" in result
    assert raw_reads == 0


@pytest.mark.anyio
async def test_fetch_webpage_enforces_wall_clock_timeout_for_slow_drip(
    monkeypatch,
):
    from EvoScientist.tools import search

    class SlowDripStream(httpx.AsyncByteStream):
        async def __aiter__(self):
            for _ in range(20):
                await asyncio.sleep(0.04)
                yield b"x"

    def public_dns(_host, port, *_args, **_kwargs):
        return [(2, 1, 6, "", ("93.184.216.34", port))]

    async def send(_client, request, **_kwargs):
        return httpx.Response(
            200,
            headers={"content-type": "text/plain"},
            stream=SlowDripStream(),
            request=request,
        )

    monkeypatch.setattr(socket, "getaddrinfo", public_dns)
    monkeypatch.setattr(httpx.AsyncClient, "send", send)

    started = asyncio.get_running_loop().time()
    result = await search.fetch_webpage_content(
        "https://public.example/slow",
        timeout=0.05,
    )
    elapsed = asyncio.get_running_loop().time() - started

    assert "wall-clock timeout" in result
    assert elapsed < 0.15


@pytest.mark.anyio
async def test_tavily_search_keeps_indexed_summary_when_source_fetch_fails(monkeypatch):
    from EvoScientist.tools import search

    class _Client:
        def search(self, *_args, **_kwargs):
            return {
                "results": [
                    {
                        "title": "AIR staff profile",
                        "url": "https://air.cas.cn/example",
                        "content": "Indexed staff-profile summary.",
                    }
                ]
            }

    recorded: list[tuple[str, str]] = []

    async def fetch_failed(_url: str, timeout: float = 10.0) -> str:
        return "Error fetching content from https://air.cas.cn/example: DNS failed"

    async def record(service: str, action: str) -> None:
        recorded.append((service, action))

    monkeypatch.setattr(search, "_get_tavily_client", lambda: _Client())
    monkeypatch.setattr(search, "fetch_webpage_content", fetch_failed)
    monkeypatch.setattr("EvoScientist.runtime_integrations.record_service_usage", record)

    result = await search.tavily_search.ainvoke({"query": "高铭 空天院"})

    assert "Indexed staff-profile summary." in result
    assert "https://air.cas.cn/example" in result
    assert "Tavily-indexed summary" in result
    assert recorded == [("tavily", "search")]


@pytest.mark.anyio
async def test_tavily_search_bounds_fetched_page_content(monkeypatch):
    from EvoScientist.tools import search

    class _Client:
        def search(self, *_args, **_kwargs):
            return {
                "results": [
                    {
                        "title": f"Result {index}",
                        "url": f"https://example.com/{index}",
                        "content": f"Indexed summary {index}",
                    }
                    for index in range(3)
                ]
            }

    async def huge_page(_url: str, timeout: float = 10.0) -> str:
        return "page-content " * 10_000

    monkeypatch.setattr(search, "_get_tavily_client", lambda: _Client())
    monkeypatch.setattr(search, "fetch_webpage_content", huge_page)

    result = await search.tavily_search.ainvoke({"query": "bounded search"})

    assert len(result) <= search.MAX_SEARCH_RESULT_CHARS
    for index in range(3):
        assert f"https://example.com/{index}" in result
    assert "[page content truncated]" in result


@pytest.mark.anyio
async def test_tavily_search_preserves_every_result_url_under_total_budget(monkeypatch):
    from EvoScientist.tools import search

    class _Client:
        def search(self, *_args, **_kwargs):
            return {
                "results": [
                    {
                        "title": f"Result {index} " + ("very-long-title " * 800),
                        "url": f"https://example.com/result-{index}",
                        "content": f"Indexed summary {index}",
                    }
                    for index in range(3)
                ]
            }

    async def page(_url: str, timeout: float = 10.0) -> str:
        return "page-content " * 1_000

    monkeypatch.setattr(search, "_get_tavily_client", lambda: _Client())
    monkeypatch.setattr(search, "fetch_webpage_content", page)

    result = await search.tavily_search.ainvoke({"query": "preserve urls"})

    assert len(result) <= search.MAX_SEARCH_RESULT_CHARS
    for index in range(3):
        assert f"https://example.com/result-{index}" in result
    assert "[search result content truncated to preserve all titles and URLs]" in result


@pytest.mark.anyio
async def test_tavily_search_bounds_maliciously_long_url(monkeypatch):
    from EvoScientist.tools import search

    long_url = "https://example.com/" + ("a" * 20_000)

    class _Client:
        def search(self, *_args, **_kwargs):
            return {
                "results": [
                    {
                        "title": "Long URL result",
                        "url": long_url,
                        "content": "Indexed summary",
                    }
                ]
            }

    async def page(_url: str, timeout: float = 10.0) -> str:
        return "page"

    monkeypatch.setattr(search, "_get_tavily_client", lambda: _Client())
    monkeypatch.setattr(search, "fetch_webpage_content", page)

    result = await search.tavily_search.ainvoke({"query": "long url"})

    assert len(result) <= search.MAX_SEARCH_RESULT_CHARS
    assert "https://example.com/" in result
    assert "[URL truncated]" in result