Fix stack not being unwound when suppressing interrupt

libs/langgraph/langgraph/pregel/executor.py

@@ -80,18 +80,18 @@ class BackgroundExecutor(ContextManager):
             if cancel:
                 task.cancel()
         # wait for all tasks to finish
-        concurrent.futures.wait({t for t in self.tasks if not t.done()})
+        if tasks := {t for t in self.tasks if not t.done()}:
+            concurrent.futures.wait(tasks)
         # shutdown the executor
         self.stack.__exit__(exc_type, exc_value, traceback)
-        # raise caught exception
-        if exc_type is not None:
-            raise exc_value
         # re-raise the first exception that occurred in a task
-        for task in self.tasks:
-            try:
-                task.result()
-            except concurrent.futures.CancelledError:
-                pass
+        if exc_type is None:
+            # if there's already an exception being raised, don't raise another one
+            for task in self.tasks:
+                try:
+                    task.result()
+                except concurrent.futures.CancelledError:
+                    pass
 
 
 class AsyncBackgroundExecutor(AsyncContextManager):
@@ -132,16 +132,27 @@ class AsyncBackgroundExecutor(AsyncContextManager):
     async def __aenter__(self) -> Submit:
         return self.submit
 
-    async def exit(self) -> None:
-        fut = asyncio.gather(*self.tasks, return_exceptions=True)
-        try:
-            rtns = await asyncio.shield(fut)
-        finally:
-            del self.tasks
-        for rtn in rtns:
-            # if this is ever changed to BaseException, need to ignore CancelledError
-            if isinstance(rtn, Exception):
-                raise rtn
+    async def exit(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        # cancel all tasks that should be cancelled
+        for task, cancel in self.tasks.items():
+            if cancel:
+                task.cancel(self.sentinel)
+        # 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 exc_type is None:
+            # if there's already an exception being raised, don't raise another one
+            for task in self.tasks:
+                try:
+                    task.result()
+                except asyncio.CancelledError:
+                    pass
 
     async def __aexit__(
         self,
@@ -151,8 +162,6 @@ class AsyncBackgroundExecutor(AsyncContextManager):
     ) -> Optional[bool]:
         # we cannot use `await` outside of asyncio.shield, as this code can run
         # after owning task is cancelled, so pulling async logic to separate method
-        for task, cancel in self.tasks.items():
-            if cancel:
-                task.cancel(self.sentinel)
+
         # wait for all background tasks to finish, shielded from cancellation
-        await asyncio.shield(self.exit())
+        await asyncio.shield(self.exit(exc_type, exc_value, traceback))

libs/langgraph/langgraph/pregel/loop.py

@@ -102,6 +102,23 @@ class PregelLoop:
 
     # public
 
+    def __init__(
+        self,
+        input: Optional[Any],
+        *,
+        config: RunnableConfig,
+        checkpointer: Optional[BaseCheckpointSaver],
+        graph: "Pregel",
+    ) -> None:
+        self.stream = deque()
+        self.input = input
+        self.config = config
+        self.checkpointer = checkpointer
+        self.graph = graph
+        # TODO if managed values no longer needs graph we can replace with
+        # managed_specs, channel_specs
+        self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {})
+
     def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None:
         """Mark tasks as scheduled, to be used by queue-based executors."""
         raise NotImplementedError
@@ -219,7 +236,7 @@ class PregelLoop:
         if should_interrupt(self.checkpoint, interrupt_before, self.tasks):
             self.status = "interrupt_before"
             if self.is_nested:
-                raise GraphInterrupt(self)
+                raise GraphInterrupt()
             else:
                 return False
 
@@ -318,6 +335,15 @@ class PregelLoop:
         # increment step
         self.step += 1
 
+    def _suppress_interrupt(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        if exc_type is GraphInterrupt and not self.is_nested:
+            return True
+
 
 class SyncPregelLoop(PregelLoop, ContextManager):
     def __init__(
@@ -328,24 +354,21 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         checkpointer: Optional[BaseCheckpointSaver],
         graph: "Pregel",
     ) -> None:
-        self.stream = deque()
+        super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
         self.stack = ExitStack()
-        self.input = input
-        self.config = config
-        self.checkpointer = checkpointer
-        self.checkpointer_get_next_version = (
-            checkpointer.get_next_version if checkpointer else increment
-        )
-        self.checkpointer_put_writes = checkpointer.put_writes if checkpointer else None
-        self.checkpointer_put = checkpointer.put if checkpointer else None
-        self.graph = graph
-        # TODO if managed values no longer needs graph we can replace with
-        # managed_specs, channel_specs
+        self.stack.push(self._suppress_interrupt)
+        if checkpointer:
+            self.checkpointer_get_next_version = checkpointer.get_next_version
+            self.checkpointer_put_writes = checkpointer.put_writes
+            self.checkpointer_put = checkpointer.put
+        else:
+            self.checkpointer_get_next_version = increment
+            self.checkpointer_put_writes = None
+            self.checkpointer_put = None
 
     # context manager
 
     def __enter__(self) -> Self:
-        self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {})
         saved = (
             self.checkpointer.get_tuple(self.config) if self.checkpointer else None
         ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, [])
@@ -381,16 +404,6 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         exc_value: Optional[BaseException],
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:
-        # handle interrupt
-        if exc_type is GraphInterrupt:
-            if exc_value.args[0] is self:
-                # interrupt raised by this loop
-                exc_value.args = (object(),)
-            if not self.is_nested:
-                # in outer graph, catch interrupt
-                del self.graph
-                return True or self.stack.__exit__(None, None, None)
-
         # unwind stack
         del self.graph
         return self.stack.__exit__(exc_type, exc_value, traceback)
@@ -405,26 +418,21 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         checkpointer: Optional[BaseCheckpointSaver],
         graph: "Pregel",
     ) -> None:
-        self.stream = deque()
+        super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
         self.stack = AsyncExitStack()
-        self.input = input
-        self.config = config
-        self.checkpointer = checkpointer
-        self.checkpointer_get_next_version = (
-            checkpointer.get_next_version if checkpointer else increment
-        )
-        self.checkpointer_put_writes = (
-            checkpointer.aput_writes if checkpointer else None
-        )
-        self.checkpointer_put = checkpointer.aput if checkpointer else None
-        self.graph = graph
-        # TODO if managed values no longer needs graph we can replace with
-        # managed_specs, channel_specs
+        self.stack.push(self._suppress_interrupt)
+        if checkpointer:
+            self.checkpointer_get_next_version = checkpointer.get_next_version
+            self.checkpointer_put_writes = checkpointer.aput_writes
+            self.checkpointer_put = checkpointer.aput
+        else:
+            self.checkpointer_get_next_version = increment
+            self.checkpointer_put_writes = None
+            self.checkpointer_put = None
 
     # context manager
 
     async def __aenter__(self) -> Self:
-        self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {})
         saved = (
             await self.checkpointer.aget_tuple(self.config)
             if self.checkpointer
@@ -462,18 +470,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         exc_value: Optional[BaseException],
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:
-        # handle interrupt
-        if exc_type is GraphInterrupt:
-            if exc_value.args[0] is self:
-                # interrupt raised by this loop
-                exc_value.args = (object(),)
-            if not self.is_nested:
-                # in outer graph, catch interrupt
-                del self.graph
-                return True or await asyncio.shield(
-                    self.stack.__aexit__(None, None, None)
-                )
-
         # unwind stack
         del self.graph
         return await asyncio.shield(