diff --git a/src/agents/sandbox/session/dependencies.py b/src/agents/sandbox/session/dependencies.py index 1a3f1fd40d..10a54ecbe9 100644 --- a/src/agents/sandbox/session/dependencies.py +++ b/src/agents/sandbox/session/dependencies.py @@ -268,14 +268,22 @@ async def _close(self) -> None: await asyncio.gather(*active_tasks, return_exceptions=True) seen_ids: set[int] = set() + cancellation: asyncio.CancelledError | None = None for value in reversed(self._owned_results): value_id = id(value) if value_id in seen_ids: continue seen_ids.add(value_id) - await _close_best_effort(value) + try: + await _close_best_effort(value) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc self._pending.clear() self._active_tasks.clear() self._cache.clear() self._owned_results.clear() + + if cancellation is not None: + raise cancellation diff --git a/tests/sandbox/test_dependencies.py b/tests/sandbox/test_dependencies.py index b0d37a94b0..cbb2274185 100644 --- a/tests/sandbox/test_dependencies.py +++ b/tests/sandbox/test_dependencies.py @@ -44,6 +44,15 @@ async def close(self) -> None: self.calls += 1 +class _CancellingAsyncClosable: + def __init__(self) -> None: + self.calls = 0 + + async def aclose(self) -> None: + self.calls += 1 + raise asyncio.CancelledError("dependency close cancelled") + + class _SyncClosable: def __init__(self) -> None: self.calls = 0 @@ -435,6 +444,34 @@ async def test_dependencies_aclose_continues_after_waiter_cancellation() -> None assert value.completed +@pytest.mark.asyncio +async def test_dependencies_aclose_finishes_owned_cleanup_before_propagating_cancellation() -> None: + dependencies = Dependencies() + earlier = _AsyncClosable() + cancelling = _CancellingAsyncClosable() + dependencies.bind_factory( + "tests.earlier_owned", lambda _dependencies: earlier, owns_result=True + ) + dependencies.bind_factory( + "tests.cancelling_owned", lambda _dependencies: cancelling, owns_result=True + ) + + _ = await dependencies.require("tests.earlier_owned") + _ = await dependencies.require("tests.cancelling_owned") + + with pytest.raises(asyncio.CancelledError): + await dependencies.aclose() + + assert cancelling.calls == 1 + assert earlier.calls == 1 + + with pytest.raises(asyncio.CancelledError): + await dependencies.aclose() + + assert cancelling.calls == 1 + assert earlier.calls == 1 + + @pytest.mark.asyncio async def test_dependencies_bound_values_are_not_closed() -> None: dependencies = Dependencies()