EvoScientist-Multi/tests/test_recoverable_tools.py

from __future__ import annotations

from typing import Any, cast

import httpx
import pytest
from types import SimpleNamespace
from langchain_core.messages import ToolMessage

from EvoScientist.llm.contracts import EvoRuntimeError
from EvoScientist.middleware import recoverable_tools


def test_tavily_search_is_classified_read_only():
    assert recoverable_tools._effect_class("tavily_search") == "read_only"


def test_tool_effect_context_prefers_dedicated_grant(monkeypatch):
    monkeypatch.setattr(
        "langgraph.config.get_config",
        lambda: {
            "configurable": {
                "ai4sci_model_proxy": {
                    "gateway_url": "http://gateway",
                    "run_id": "model-run",
                    "envelope_signature": "model-signature",
                },
                "ai4sci_tool_effect": {
                    "gateway_url": "http://gateway",
                    "run_id": "tool-run",
                    "envelope_signature": "tool-signature",
                },
            },
            "metadata": {},
        },
    )

    proxy, _ = recoverable_tools._context()

    assert proxy == {
        "gateway_url": "http://gateway",
        "run_id": "tool-run",
        "envelope_signature": "tool-signature",
    }


def test_evomemory_never_falls_back_to_parent_model_proxy(monkeypatch):
    monkeypatch.setattr(
        "langgraph.config.get_config",
        lambda: {
            "configurable": {
                "ai4sci_model_proxy": {
                    "gateway_url": "http://gateway",
                    "run_id": "parent-run",
                    "envelope_signature": "parent-signature",
                },
            },
            "metadata": {"run_kind": "evomemory_turn_worker"},
        },
    )

    proxy, metadata = recoverable_tools._context()

    assert proxy is None
    assert metadata["run_kind"] == "evomemory_turn_worker"


@pytest.mark.asyncio
async def test_tool_effect_gateway_error_preserves_machine_code(monkeypatch):
    captured = {}

    class Client:
        async def __aenter__(self):
            return self

        async def __aexit__(self, *_):
            return None

        async def post(self, url, json, headers):
            captured["headers"] = headers
            return httpx.Response(
                409,
                json={"detail": {"code": "RUN_FENCE_LOST"}},
                request=httpx.Request("POST", url),
            )

    monkeypatch.setenv("AI4SCI_EVO_RUNTIME_GRANT_SECRET", "runtime-service-secret")
    monkeypatch.setattr(recoverable_tools.httpx, "AsyncClient", lambda **_: Client())

    with pytest.raises(EvoRuntimeError) as exc_info:
        await recoverable_tools._post(
            {
                "gateway_url": "http://gateway",
                "run_id": "run-1",
                "envelope_signature": "signature",
            },
            "prepare",
            {},
        )

    assert exc_info.value.code == "RUN_FENCE_LOST"
    assert captured["headers"] == {
        "X-Ai4Sci-Service-Token": "runtime-service-secret"
    }
    assert repr(exc_info.value) == "EvoRuntimeError(code='RUN_FENCE_LOST')"


@pytest.mark.asyncio
async def test_terminal_callback_timeout_returns_tool_error_instead_of_crashing_run(monkeypatch):
    request = SimpleNamespace(
        tool_call={"id": "tool-1", "name": "tavily_search", "args": {"query": "x"}}
    )
    monkeypatch.setattr(
        recoverable_tools,
        "_context",
        lambda: (
            {
                "gateway_url": "http://gateway",
                "run_id": "run-1",
                "envelope_signature": "signature",
            },
            {},
        ),
    )
    phases = []

    async def fake_post(_proxy, phase, _payload):
        phases.append(phase)
        if phase == "prepare":
            return {"action": "execute", "fencing_token": 7}
        raise httpx.ConnectTimeout("terminal callback timed out")

    async def handler(_request):
        return ToolMessage(
            content="one successful search result",
            tool_call_id="tool-1",
            name="tavily_search",
        )

    monkeypatch.setattr(recoverable_tools, "_post", fake_post)

    result = await recoverable_tools.RecoverableToolEffectMiddleware().awrap_tool_call(
        cast(Any, request),
        handler,
    )

    assert phases == ["prepare", "terminal"]
    assert isinstance(result, ToolMessage)
    assert result.status == "error"
    assert result.content == "TOOL_EFFECT_TERMINAL_UNAVAILABLE"


@pytest.mark.asyncio
async def test_failed_tool_keeps_original_error_when_terminal_callback_times_out(monkeypatch):
    request = SimpleNamespace(
        tool_call={"id": "tool-2", "name": "tavily_search", "args": {"query": "x"}}
    )
    monkeypatch.setattr(
        recoverable_tools,
        "_context",
        lambda: (
            {
                "gateway_url": "http://gateway",
                "run_id": "run-1",
                "envelope_signature": "signature",
            },
            {},
        ),
    )

    async def fake_post(_proxy, phase, _payload):
        if phase == "prepare":
            return {"action": "execute", "fencing_token": 8}
        raise httpx.ConnectTimeout("terminal callback timed out")

    async def handler(_request):
        raise ValueError("original tool failure")

    monkeypatch.setattr(recoverable_tools, "_post", fake_post)

    with pytest.raises(ValueError, match="original tool failure"):
        await recoverable_tools.RecoverableToolEffectMiddleware().awrap_tool_call(
            cast(Any, request),
            handler,
        )


@pytest.mark.asyncio
async def test_terminal_callback_semantic_error_remains_fail_closed(monkeypatch):
    request = SimpleNamespace(
        tool_call={"id": "tool-3", "name": "tavily_search", "args": {"query": "x"}}
    )
    monkeypatch.setattr(
        recoverable_tools,
        "_context",
        lambda: (
            {
                "gateway_url": "http://gateway",
                "run_id": "run-1",
                "envelope_signature": "signature",
            },
            {},
        ),
    )

    async def fake_post(_proxy, phase, _payload):
        if phase == "prepare":
            return {"action": "execute", "fencing_token": 9}
        raise EvoRuntimeError("RUN_FENCE_LOST")

    async def handler(_request):
        return ToolMessage(
            content="result",
            tool_call_id="tool-3",
            name="tavily_search",
        )

    monkeypatch.setattr(recoverable_tools, "_post", fake_post)

    with pytest.raises(EvoRuntimeError) as exc_info:
        await recoverable_tools.RecoverableToolEffectMiddleware().awrap_tool_call(
            cast(Any, request),
            handler,
        )

    assert exc_info.value.code == "RUN_FENCE_LOST"


@pytest.mark.asyncio
async def test_non_idempotent_terminal_transport_failure_is_fail_closed(monkeypatch):
    request = SimpleNamespace(
        tool_call={"id": "tool-4", "name": "send_message", "args": {"text": "hello"}}
    )
    monkeypatch.setattr(
        recoverable_tools,
        "_context",
        lambda: (
            {
                "gateway_url": "http://gateway",
                "run_id": "run-1",
                "envelope_signature": "signature",
            },
            {},
        ),
    )

    async def fake_post(_proxy, phase, _payload):
        if phase == "prepare":
            return {"action": "execute", "fencing_token": 10}
        raise httpx.ConnectTimeout("terminal callback timed out")

    async def handler(_request):
        return ToolMessage(
            content="message sent",
            tool_call_id="tool-4",
            name="send_message",
        )

    monkeypatch.setattr(recoverable_tools, "_post", fake_post)

    with pytest.raises(EvoRuntimeError) as exc_info:
        await recoverable_tools.RecoverableToolEffectMiddleware().awrap_tool_call(
            cast(Any, request),
            handler,
        )

    assert exc_info.value.code == "TOOL_EFFECT_TERMINAL_UNAVAILABLE"