EvoScientist-Multi/tests/test_subagent_timeout.py

from __future__ import annotations

import asyncio
from unittest.mock import MagicMock

import pytest
from langchain_core.messages import ToolMessage

from EvoScientist.middleware.subagent_timeout import SubagentTimeoutMiddleware


def _request(name: str = "task"):
    request = MagicMock()
    request.tool_call = {"id": "call-1", "name": name, "args": {}}
    return request


@pytest.mark.anyio
async def test_subagent_timeout_returns_stable_tool_error():
    middleware = SubagentTimeoutMiddleware(timeout_seconds=0.01)

    async def handler(_request):
        await asyncio.sleep(1)
        return ToolMessage(content="late", tool_call_id="call-1", name="task")

    result = await middleware.awrap_tool_call(_request(), handler)

    assert isinstance(result, ToolMessage)
    assert result.status == "error"
    assert result.name == "task"
    assert result.additional_kwargs["error_code"] == "SUBAGENT_TIMEOUT"
    assert "SUBAGENT_TIMEOUT" in result.content


@pytest.mark.anyio
async def test_subagent_timeout_passes_success_through():
    middleware = SubagentTimeoutMiddleware(timeout_seconds=1)
    expected = ToolMessage(content="done", tool_call_id="call-1", name="task")

    async def handler(_request):
        return expected

    assert await middleware.awrap_tool_call(_request(), handler) is expected


@pytest.mark.anyio
async def test_subagent_timeout_does_not_bound_other_tools():
    middleware = SubagentTimeoutMiddleware(timeout_seconds=0.01)
    expected = ToolMessage(content="done", tool_call_id="call-1", name="read_file")

    async def handler(_request):
        await asyncio.sleep(0.02)
        return expected

    assert await middleware.awrap_tool_call(_request("read_file"), handler) is expected


@pytest.mark.anyio
async def test_subagent_internal_timeout_error_is_not_reclassified():
    middleware = SubagentTimeoutMiddleware(timeout_seconds=1)

    async def handler(_request):
        raise TimeoutError("provider timed out immediately")

    with pytest.raises(TimeoutError, match="provider timed out immediately"):
        await middleware.awrap_tool_call(_request(), handler)


@pytest.mark.anyio
async def test_parent_cancellation_cancels_subagent_handler():
    middleware = SubagentTimeoutMiddleware(timeout_seconds=10)
    handler_cancelled = asyncio.Event()

    async def handler(_request) -> ToolMessage:
        try:
            await asyncio.Event().wait()
        finally:
            handler_cancelled.set()
        return ToolMessage(content="done", tool_call_id="call-1", name="task")

    invocation = asyncio.create_task(
        middleware.awrap_tool_call(_request(), handler)
    )
    await asyncio.sleep(0)
    invocation.cancel()

    with pytest.raises(asyncio.CancelledError):
        await invocation
    assert handler_cancelled.is_set()


@pytest.mark.anyio
async def test_parent_cancellation_wins_over_handler_cleanup_error():
    middleware = SubagentTimeoutMiddleware(timeout_seconds=10)

    async def handler(_request) -> ToolMessage:
        try:
            await asyncio.Event().wait()
        except asyncio.CancelledError as exc:
            raise RuntimeError("cleanup failed") from exc
        return ToolMessage(content="done", tool_call_id="call-1", name="task")

    invocation = asyncio.create_task(
        middleware.awrap_tool_call(_request(), handler)
    )
    await asyncio.sleep(0)
    invocation.cancel()

    with pytest.raises(asyncio.CancelledError):
        await invocation


@pytest.mark.anyio
async def test_deadline_wins_over_handler_cleanup_error():
    middleware = SubagentTimeoutMiddleware(timeout_seconds=0.01)

    async def handler(_request) -> ToolMessage:
        try:
            await asyncio.Event().wait()
        except asyncio.CancelledError as exc:
            raise RuntimeError("cleanup failed") from exc
        return ToolMessage(content="done", tool_call_id="call-1", name="task")

    result = await middleware.awrap_tool_call(_request(), handler)

    assert isinstance(result, ToolMessage)
    assert result.additional_kwargs["error_code"] == "SUBAGENT_TIMEOUT"


@pytest.mark.anyio
async def test_parent_cancellation_during_deadline_cleanup_is_not_swallowed():
    middleware = SubagentTimeoutMiddleware(timeout_seconds=0.01)
    cleanup_started = asyncio.Event()
    release_cleanup = asyncio.Event()

    async def handler(_request) -> ToolMessage:
        try:
            await asyncio.Event().wait()
        except asyncio.CancelledError:
            cleanup_started.set()
            await release_cleanup.wait()
            raise
        return ToolMessage(content="done", tool_call_id="call-1", name="task")

    invocation = asyncio.create_task(
        middleware.awrap_tool_call(_request(), handler)
    )
    await cleanup_started.wait()
    invocation.cancel()
    release_cleanup.set()
    with pytest.raises(asyncio.CancelledError):
        await invocation


@pytest.mark.anyio
async def test_parent_cancellation_after_cleanup_before_timeout_return_wins(monkeypatch):
    middleware = SubagentTimeoutMiddleware(timeout_seconds=0.01)

    async def handler(_request) -> ToolMessage:
        await asyncio.Event().wait()
        return ToolMessage(content="done", tool_call_id="call-1", name="task")

    async def finish_cleanup_then_cancel_parent(_task):
        current = asyncio.current_task()
        assert current is not None
        current.cancel()

    monkeypatch.setattr(middleware, "_cancel_task", finish_cleanup_then_cancel_parent)

    invocation = asyncio.create_task(
        middleware.awrap_tool_call(_request(), handler)
    )
    with pytest.raises(asyncio.CancelledError):
        await invocation