Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 90 additions & 12 deletions src/agents/sandbox/session/dependencies.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

import asyncio
import inspect
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass
Expand Down Expand Up @@ -73,7 +74,10 @@ 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._owned_results: list[object] = []
self._close_task: asyncio.Task[None] | None = None
self._closed = False

@classmethod
Expand Down Expand Up @@ -144,6 +148,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()
Comment on lines +151 to +153

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Handle factory tasks cancelled before their first step

When an overwrite races immediately after require() creates the pending task but before _run_factory receives its first timeslice, this cancellation prevents the coroutine's try/finally from running at all. The stale waiter consequently receives a raw CancelledError instead of the intended DependenciesBindingError, and the completed task remains permanently in _active_tasks; uncached requests cancelled in the same pre-start window also leave retained tasks. Ensure active-set cleanup runs from a task completion callback and translate pre-start rebind cancellation consistently.

AGENTS.md reference: AGENTS.md:L115-L115

Useful? React with 👍 / 👎.


async def get(self, key: DependencyKey) -> object | None:
binding = self._bindings.get(key)
Expand Down Expand Up @@ -173,24 +180,86 @@ 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
Comment on lines +191 to +192

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Clear failures completed by an eager task factory

On Python 3.12+ when the event loop uses asyncio.eager_task_factory (or an equivalent custom eager factory), create_task() can run _run_factory to completion before returning. For a synchronously failing cached factory, its finally therefore sees no _pending[key], after which this assignment stores the already-failed task permanently; every later require() replays the first exception instead of invoking the factory again, so the documented failure-and-retry behavior breaks in that environment. Avoid publishing an already-completed task as pending, or remove it from _pending from a completion callback after publication.

Useful? React with 👍 / 👎.

return await asyncio.shield(task)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Recheck closure after awaiting the factory task

When an owned factory completes and another task calls aclose() before the shielded waiter resumes, _run_factory has already removed itself from _active_tasks, so shutdown closes and clears the produced value; this await then still returns that already-closed object. This is reproducible when the factory signals an event immediately before returning and a closer waits on that event, and differs from the intended in-flight-close path that raises DependenciesError; recheck the container state after the await so a pre-close request cannot receive a disposed dependency.

AGENTS.md reference: AGENTS.md:L115-L115

Useful? React with 👍 / 👎.


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:
self._owned_results.append(value)
raise DependenciesError(f"Dependencies container closed while resolving `{key}`")

if self._bindings.get(key) is not binding:
if binding.owns_result:
self._owned_results.append(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 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)

seen_ids: set[int] = set()
for value in reversed(self._owned_results):
Expand All @@ -199,3 +268,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()
Loading