diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 5da78a26c..3e4243c61 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -120,14 +120,14 @@ class AsyncBackgroundExecutor(AsyncContextManager): def done(self, task: asyncio.Task) -> None: try: - task.result() - except GraphInterrupt: - # This exception is an interruption signal, not an error - # so we don't want to re-raise it on exit - self.tasks.pop(task) - except BaseException: - pass - else: + if exc := task.exception(): + # This exception is an interruption signal, not an error + # so we don't want to re-raise it on exit + if isinstance(exc, GraphInterrupt): + self.tasks.pop(task) + else: + self.tasks.pop(task) + except asyncio.CancelledError: self.tasks.pop(task) async def __aenter__(self) -> Submit: @@ -146,12 +146,13 @@ class AsyncBackgroundExecutor(AsyncContextManager): # wait for all tasks to finish if self.tasks: await asyncio.wait(self.tasks) - # re-raise the first exception that occurred in a task + # if there's already an exception being raised, don't raise another one if exc_type is None: - # if there's already an exception being raised, don't raise another one + # re-raise the first exception that occurred in a task for task in self.tasks: try: - task.result() + if exc := task.exception(): + raise exc except asyncio.CancelledError: pass