diff --git a/src/agents/sandbox/session/dependencies.py b/src/agents/sandbox/session/dependencies.py index cb1cec7552..8bfd75ec10 100644 --- a/src/agents/sandbox/session/dependencies.py +++ b/src/agents/sandbox/session/dependencies.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import inspect from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass @@ -73,7 +74,11 @@ class Dependencies: def __init__(self) -> None: self._bindings: dict[DependencyKey, _Binding] = {} self._cache: dict[DependencyKey, object] = {} + self._pending: dict[DependencyKey, asyncio.Task[object]] = {} + self._active_tasks: set[asyncio.Task[object]] = set() + self._cleanup_tasks: set[asyncio.Task[None]] = set() self._owned_results: list[object] = [] + self._close_task: asyncio.Task[None] | None = None self._closed = False @classmethod @@ -144,6 +149,9 @@ def _bind( raise DependenciesBindingError(f"Dependency `{key}` is already bound") self._bindings[key] = binding self._cache.pop(key, None) + pending = self._pending.pop(key, None) + if pending is not None: + pending.cancel() async def get(self, key: DependencyKey) -> object | None: binding = self._bindings.get(key) @@ -173,24 +181,96 @@ async def _resolve(self, key: DependencyKey, binding: _Binding) -> object: return binding.value assert isinstance(binding, _FactoryBinding) + if self._closed: + raise DependenciesError(f"Dependencies container is closed; cannot resolve `{key}`") if binding.cache and key in self._cache: return self._cache[key] - produced = binding.factory(self) - value = ( - await cast(Awaitable[object], produced) if inspect.isawaitable(produced) else produced - ) - if binding.cache: - self._cache[key] = value - if binding.owns_result: - self._owned_results.append(value) - return value + task = self._pending.get(key) + if task is None: + task = self._create_factory_task(key, binding) + self._pending[key] = task + return await asyncio.shield(task) + + task = self._create_factory_task(key, binding) + return await task + + def _create_factory_task( + self, key: DependencyKey, binding: _FactoryBinding + ) -> asyncio.Task[object]: + task = asyncio.create_task(self._run_factory(key, binding)) + self._active_tasks.add(task) + task.add_done_callback(_consume_task_exception) + return task + + async def _run_factory(self, key: DependencyKey, binding: _FactoryBinding) -> object: + try: + produced = binding.factory(self) + value = ( + await cast(Awaitable[object], produced) + if inspect.isawaitable(produced) + else produced + ) + + if self._closed: + if binding.owns_result: + await self._discard_factory_result(value) + raise DependenciesError(f"Dependencies container closed while resolving `{key}`") + + if self._bindings.get(key) is not binding: + if binding.owns_result: + await self._discard_factory_result(value) + raise DependenciesBindingError( + f"Dependency `{key}` was rebound while its factory was resolving" + ) + + if binding.cache: + self._cache[key] = value + if binding.owns_result: + self._owned_results.append(value) + return value + except asyncio.CancelledError: + if self._closed: + raise DependenciesError( + f"Dependencies container closed while resolving `{key}`" + ) from None + if self._bindings.get(key) is not binding: + raise DependenciesBindingError( + f"Dependency `{key}` was rebound while its factory was resolving" + ) from None + raise + finally: + task = asyncio.current_task() + if task is not None: + self._active_tasks.discard(task) + if self._pending.get(key) is task: + self._pending.pop(key, None) + + async def _discard_factory_result(self, value: object) -> None: + task = asyncio.create_task(_close_best_effort(value)) + self._cleanup_tasks.add(task) + task.add_done_callback(self._cleanup_tasks.discard) + await asyncio.shield(task) async def aclose(self) -> None: - if self._closed: - return - self._closed = True + task = self._close_task + if task is None: + self._closed = True + task = asyncio.create_task(self._close()) + self._close_task = task + await asyncio.shield(task) + + async def _close(self) -> None: + active_tasks = tuple(self._active_tasks) + for task in active_tasks: + task.cancel() + if active_tasks: + await asyncio.gather(*active_tasks, return_exceptions=True) + + while self._cleanup_tasks: + cleanup_tasks = tuple(self._cleanup_tasks) + await asyncio.gather(*cleanup_tasks, return_exceptions=True) seen_ids: set[int] = set() for value in reversed(self._owned_results): @@ -199,3 +279,12 @@ async def aclose(self) -> None: continue seen_ids.add(value_id) await _close_best_effort(value) + + self._pending.clear() + self._cache.clear() + self._owned_results.clear() + + +def _consume_task_exception(task: asyncio.Task[object]) -> None: + if not task.cancelled(): + task.exception() diff --git a/tests/sandbox/test_dependencies.py b/tests/sandbox/test_dependencies.py index ed282cf3e1..60553d1fc4 100644 --- a/tests/sandbox/test_dependencies.py +++ b/tests/sandbox/test_dependencies.py @@ -1,10 +1,13 @@ from __future__ import annotations +import asyncio + import pytest from agents.sandbox.session import ( Dependencies, DependenciesBindingError, + DependenciesError, DependenciesMissingDependencyError, ) @@ -17,6 +20,20 @@ async def aclose(self) -> None: self.calls += 1 +class _BlockingAsyncClosable: + def __init__(self) -> None: + self.calls = 0 + self.completed = False + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def aclose(self) -> None: + self.calls += 1 + self.started.set() + await self.release.wait() + self.completed = True + + class _AsyncCloseMethod: def __init__(self) -> None: self.calls = 0 @@ -99,6 +116,222 @@ def _factory(_dependencies: Dependencies) -> str: assert calls == 1 +@pytest.mark.asyncio +async def test_dependencies_cached_factory_resolves_once_concurrently() -> None: + dependencies = Dependencies() + key = "tests.concurrent_cached_factory" + started = asyncio.Event() + release = asyncio.Event() + calls = 0 + + async def _factory(_dependencies: Dependencies) -> _AsyncClosable: + nonlocal calls + calls += 1 + started.set() + await release.wait() + return _AsyncClosable() + + dependencies.bind_factory(key, _factory, cache=True, owns_result=True) + tasks = [asyncio.create_task(dependencies.require(key)) for _ in range(3)] + + await started.wait() + await asyncio.sleep(0) + release.set() + values = await asyncio.gather(*tasks) + + assert calls == 1 + assert values[0] is values[1] is values[2] + + await dependencies.aclose() + assert isinstance(values[0], _AsyncClosable) + assert values[0].calls == 1 + + +@pytest.mark.asyncio +async def test_dependencies_cached_factory_survives_waiter_cancellation() -> None: + dependencies = Dependencies() + key = "tests.cancelled_waiter" + started = asyncio.Event() + release = asyncio.Event() + calls = 0 + + async def _factory(_dependencies: Dependencies) -> object: + nonlocal calls + calls += 1 + started.set() + await release.wait() + return object() + + dependencies.bind_factory(key, _factory, cache=True) + cancelled_waiter = asyncio.create_task(dependencies.require(key)) + surviving_waiter = asyncio.create_task(dependencies.require(key)) + + await started.wait() + await asyncio.sleep(0) + cancelled_waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled_waiter + + release.set() + value = await surviving_waiter + + assert calls == 1 + assert await dependencies.require(key) is value + + +@pytest.mark.asyncio +async def test_dependencies_cached_factory_failure_allows_retry() -> None: + dependencies = Dependencies() + key = "tests.failed_factory_retry" + started = asyncio.Event() + release = asyncio.Event() + calls = 0 + + async def _factory(_dependencies: Dependencies) -> str: + nonlocal calls + calls += 1 + if calls == 1: + started.set() + await release.wait() + raise RuntimeError("factory failed") + return "recovered" + + dependencies.bind_factory(key, _factory, cache=True) + first = asyncio.create_task(dependencies.require(key)) + second = asyncio.create_task(dependencies.require(key)) + + await started.wait() + await asyncio.sleep(0) + release.set() + + for task in (first, second): + with pytest.raises(RuntimeError, match="factory failed"): + await task + + assert await dependencies.require(key) == "recovered" + assert calls == 2 + + +@pytest.mark.asyncio +async def test_dependencies_aclose_discards_owned_in_flight_result() -> None: + dependencies = Dependencies() + key = "tests.close_in_flight" + started = asyncio.Event() + produced: list[_AsyncClosable] = [] + + async def _factory(_dependencies: Dependencies) -> _AsyncClosable: + started.set() + try: + await asyncio.Future() + raise AssertionError("Unreachable") + except asyncio.CancelledError: + value = _AsyncClosable() + produced.append(value) + return value + + dependencies.bind_factory(key, _factory, cache=True, owns_result=True) + resolve_task = asyncio.create_task(dependencies.require(key)) + + await started.wait() + await dependencies.aclose() + + with pytest.raises(DependenciesError, match="closed"): + await asyncio.wait_for(resolve_task, timeout=0.1) + assert len(produced) == 1 + assert produced[0].calls == 1 + + +@pytest.mark.asyncio +async def test_dependencies_rebind_discards_stale_owned_in_flight_result() -> None: + dependencies = Dependencies() + key = "tests.rebind_in_flight" + started = asyncio.Event() + produced: list[_AsyncClosable] = [] + + async def _factory(_dependencies: Dependencies) -> _AsyncClosable: + started.set() + try: + await asyncio.Future() + raise AssertionError("Unreachable") + except asyncio.CancelledError: + value = _AsyncClosable() + produced.append(value) + return value + + dependencies.bind_factory(key, _factory, cache=True, owns_result=True) + stale_resolve = asyncio.create_task(dependencies.require(key)) + + await started.wait() + dependencies.bind_factory(key, lambda _dependencies: "replacement", overwrite=True) + + with pytest.raises(DependenciesBindingError, match="rebound"): + await stale_resolve + assert await dependencies.require(key) == "replacement" + assert len(produced) == 1 + assert produced[0].calls == 1 + + +@pytest.mark.asyncio +async def test_dependencies_rebind_reports_binding_error_to_waiters() -> None: + dependencies = Dependencies() + key = "tests.rebind_cancelled_factory" + started = asyncio.Event() + + async def _factory(_dependencies: Dependencies) -> object: + started.set() + await asyncio.Future() + raise AssertionError("Unreachable") + + dependencies.bind_factory(key, _factory, cache=True) + stale_resolve = asyncio.create_task(dependencies.require(key)) + + await started.wait() + dependencies.bind_factory(key, lambda _dependencies: "replacement", overwrite=True) + + with pytest.raises(DependenciesBindingError, match="rebound"): + await stale_resolve + assert await dependencies.require(key) == "replacement" + + +@pytest.mark.asyncio +async def test_dependencies_aclose_waits_for_stale_owned_result_cleanup() -> None: + dependencies = Dependencies() + key = "tests.close_during_stale_cleanup" + started = asyncio.Event() + produced = asyncio.Event() + values: list[_BlockingAsyncClosable] = [] + + async def _factory(_dependencies: Dependencies) -> _BlockingAsyncClosable: + started.set() + try: + await asyncio.Future() + raise AssertionError("Unreachable") + except asyncio.CancelledError: + value = _BlockingAsyncClosable() + values.append(value) + produced.set() + return value + + dependencies.bind_factory(key, _factory, cache=True, owns_result=True) + stale_resolve = asyncio.create_task(dependencies.require(key)) + + await started.wait() + dependencies.bind_factory(key, lambda _dependencies: "replacement", overwrite=True) + await produced.wait() + value = values[0] + await value.started.wait() + + close_waiter = asyncio.create_task(dependencies.aclose()) + await asyncio.sleep(0) + value.release.set() + await close_waiter + + with pytest.raises(DependenciesError, match="closed"): + await stale_resolve + assert value.calls == 1 + assert value.completed + + @pytest.mark.asyncio async def test_dependencies_uncached_factory_resolves_every_time() -> None: dependencies = Dependencies() @@ -156,6 +389,27 @@ async def test_dependencies_aclose_closes_owned_results_and_is_idempotent() -> N assert isinstance(v3b, _SyncClosable) and v3b.calls == 1 +@pytest.mark.asyncio +async def test_dependencies_aclose_continues_after_waiter_cancellation() -> None: + dependencies = Dependencies() + key = "tests.cancelled_close" + value = _BlockingAsyncClosable() + dependencies.bind_factory(key, lambda _dependencies: value, owns_result=True) + _ = await dependencies.require(key) + + close_waiter = asyncio.create_task(dependencies.aclose()) + await value.started.wait() + close_waiter.cancel() + + with pytest.raises(asyncio.CancelledError): + await close_waiter + + value.release.set() + await dependencies.aclose() + assert value.calls == 1 + assert value.completed + + @pytest.mark.asyncio async def test_dependencies_bound_values_are_not_closed() -> None: dependencies = Dependencies()