diff --git a/src/saturn_engine/utils/asyncutils.py b/src/saturn_engine/utils/asyncutils.py index 3772027e..2601e199 100644 --- a/src/saturn_engine/utils/asyncutils.py +++ b/src/saturn_engine/utils/asyncutils.py @@ -441,4 +441,7 @@ async def __call__(self, *args: t.Any, **kwargs: t.Any) -> R: self.func(*args, **kwargs), name=f"cancellable({self.func})" ) self._tasks.add(task) - return await task + try: + return await task + finally: + self._tasks.remove(task)