From 70746042049e684d751e8f150803589a69306a74 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 22 Aug 2024 14:23:48 -0700 Subject: [PATCH] Try to improve async stack traces for exceptions in tasks (#1442) * Try to improve async stack traces for exceptions in tasks * Lint --- libs/langgraph/langgraph/pregel/executor.py | 23 +++++++++++---------- 1 file changed, 12 insertions(+), 11 deletions(-) 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