diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 89ecc4420..064418390 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -1032,4 +1032,13 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): traceback: Optional[TracebackType], ) -> Optional[bool]: # unwind stack - return await self.stack.__aexit__(exc_type, exc_value, traceback) + exit_task = asyncio.create_task( + self.stack.__aexit__(exc_type, exc_value, traceback) + ) + try: + return await exit_task + except asyncio.CancelledError as e: + # Bubble up the exit task upon cancellation to permit the API + # consumer to await it before e.g., re-using the DB connection. + e.args = (*e.args, exit_task) + raise