diff --git a/swe_af/execution/coding_loop.py b/swe_af/execution/coding_loop.py index 9790e9e..394d5e2 100644 --- a/swe_af/execution/coding_loop.py +++ b/swe_af/execution/coding_loop.py @@ -35,8 +35,8 @@ async def _call_with_timeout(coro, timeout: int = 2700, label: str = ""): """Wrap a coroutine with asyncio.wait_for timeout.""" try: return await asyncio.wait_for(coro, timeout=timeout) - except asyncio.TimeoutError: - raise TimeoutError(f"Agent call '{label}' timed out after {timeout}s") + except asyncio.TimeoutError as exc: + raise TimeoutError(f"Agent call '{label}' timed out after {timeout}s") from exc # --------------------------------------------------------------------------- diff --git a/swe_af/execution/dag_executor.py b/swe_af/execution/dag_executor.py index 212000b..46480cf 100644 --- a/swe_af/execution/dag_executor.py +++ b/swe_af/execution/dag_executor.py @@ -42,10 +42,10 @@ async def _call_with_timeout(coro, timeout: int = 2700, label: str = ""): """ try: return await asyncio.wait_for(coro, timeout=timeout) - except asyncio.TimeoutError: + except asyncio.TimeoutError as exc: raise TimeoutError( f"Agent call '{label}' timed out after {timeout}s" - ) + ) from exc # --------------------------------------------------------------------------- diff --git a/tests/test_call_with_timeout.py b/tests/test_call_with_timeout.py new file mode 100644 index 0000000..1bc1d47 --- /dev/null +++ b/tests/test_call_with_timeout.py @@ -0,0 +1,23 @@ +"""Unit tests for the timeout wrappers used by the execution engine.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from swe_af.execution.coding_loop import _call_with_timeout as coding_loop_timeout +from swe_af.execution.dag_executor import _call_with_timeout as dag_executor_timeout + + +@pytest.mark.parametrize("fn", [coding_loop_timeout, dag_executor_timeout]) +async def test_call_with_timeout_chains_original_exception(fn): + """TimeoutError must carry the original asyncio.TimeoutError as its cause.""" + + async def slow(): + raise asyncio.TimeoutError("simulated") + + with pytest.raises(TimeoutError) as exc_info: + await fn(slow(), timeout=1, label="test") + + assert isinstance(exc_info.value.__cause__, asyncio.TimeoutError)