From c6ee807de543a42f1aba1d7d79828c8964156aaf Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Wed, 1 Jan 2025 05:41:57 -0800
Subject: [PATCH] Bubble up stack close task

---
 libs/langgraph/langgraph/pregel/loop.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

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