EvoScientist-Multi/tests/test_subagent_timeout.py

from __future__ import annotations

import asyncio
import math
import time
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.parametrize("value", [math.nan, math.inf, -math.inf])
@pytest.mark.parametrize("field", ["timeout_seconds", "cancellation_grace_seconds"])
def test_subagent_timeout_rejects_non_finite_bounds(field, value):
    kwargs = {field: value}
    with pytest.raises(ValueError, match="finite and positive"):
        SubagentTimeoutMiddleware(**kwargs)


@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


@pytest.mark.anyio
async def test_deadline_detaches_handler_that_ignores_cancellation():
    middleware = SubagentTimeoutMiddleware(
        timeout_seconds=0.01, cancellation_grace_seconds=0.01
    )
    cleanup_started = asyncio.Event()
    release_cleanup = asyncio.Event()
    handler_done = asyncio.Event()

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

    started = time.monotonic()
    result = await middleware.awrap_tool_call(_request(), handler)
    elapsed = time.monotonic() - started

    assert cleanup_started.is_set()
    assert elapsed < 0.1
    assert result.additional_kwargs["error_code"] == "SUBAGENT_TIMEOUT"
    release_cleanup.set()
    await asyncio.wait_for(handler_done.wait(), timeout=0.2)


@pytest.mark.anyio
async def test_parent_cancel_detaches_handler_that_ignores_cancellation():
    middleware = SubagentTimeoutMiddleware(
        timeout_seconds=10, cancellation_grace_seconds=0.01
    )
    handler_started = asyncio.Event()
    cleanup_started = asyncio.Event()
    release_cleanup = asyncio.Event()
    handler_done = asyncio.Event()

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

    invocation = asyncio.create_task(
        middleware.awrap_tool_call(_request(), handler)
    )
    await handler_started.wait()
    started = time.monotonic()
    invocation.cancel()
    with pytest.raises(asyncio.CancelledError):
        await invocation
    elapsed = time.monotonic() - started

    assert cleanup_started.is_set()
    assert elapsed < 0.1
    release_cleanup.set()
    await asyncio.wait_for(handler_done.wait(), timeout=0.2)