fix(langgraph): fix interrupt optimization for AsyncPregelLoop

libs/langgraph/langgraph/pregel/_loop.py

@@ -356,9 +356,7 @@ class PregelLoop:
                 )
         # output writes
         if hasattr(self, "tasks"):
-            print("OUTPUT WRITES CALL POS 0")
             self.output_writes(task_id, writes)
-            print("^^^")
 
     def _put_pending_writes(self) -> None:
         if self.checkpointer_put_writes is None:
@@ -449,8 +447,6 @@ class PregelLoop:
             True if more iterations are needed.
         """
 
-        print(f"----!!!! TICK: {self.step}")
-
         # check if iteration limit is reached
         if self.step > self.stop:
             self.status = "out_of_steps"
@@ -492,7 +488,6 @@ class PregelLoop:
                     task_ids_to_block.add(task_id)
 
         self.task_ids_to_block = task_ids_to_block
-        print("FIRST GO: self.task_ids_to_block: ", self.task_ids_to_block)
 
         # produce debug output
         if self._checkpointer_put_after_previous is not None:
@@ -522,9 +517,7 @@ class PregelLoop:
 
         # if there are pending writes from a previous loop, apply them
         if self.skip_done_tasks and self.checkpoint_pending_writes:
-            print(">>Wrapper")
             self._match_writes(self.tasks)
-            print(">>Wrapper done")
 
         # before execution, check if we should interrupt
         if self.interrupt_before and should_interrupt(
@@ -539,9 +532,7 @@ class PregelLoop:
         # print output for any tasks we applied previous writes to
         for task in self.tasks.values():
             if task.writes:
-                print("OUTPUT WRITES CALL POS 1")
                 self.output_writes(task.id, task.writes, cached=True)
-                print("^^^")
 
         subtractor = set()
         for task_id in self.task_ids_to_block:
@@ -549,19 +540,14 @@ class PregelLoop:
                 subtractor.add(task_id)
         self.task_ids_to_block = self.task_ids_to_block - subtractor
 
-        print("SECOND GO: self.task_ids_to_block: ", self.task_ids_to_block)
-
         for task_id, write_type, value in self.checkpoint_pending_writes:
             if task_id in self.task_ids_to_block:
-                print("OUTPUT WRITES CALL POS 2")
                 self.output_writes(task_id, [(write_type, value)])
-                print("^^^")
 
         return True
 
     def after_tick(self) -> None:
         if self.task_ids_to_block:
-            print(f"THIRD GO: self.task_ids_to_block: {self.task_ids_to_block}")
             raise GraphInterrupt(
                 tuple(
                     value[0]
@@ -956,7 +942,6 @@ class PregelLoop:
                         )
                     }
                 ]
-                print("Outputting interrupts: ", interrupts)
                 stream_modes = self.stream.modes if self.stream else []
                 if "updates" in stream_modes:
                     self._emit("updates", lambda: iter(interrupts))
@@ -1074,9 +1059,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
     ) -> PregelExecutableTask | None:
         if pushed := super().accept_push(task, write_idx, call):
             for task in self.match_cached_writes():
-                print("OUTPUT WRITES CALL POS 3")
                 self.output_writes(task.id, task.writes, cached=True)
-                print("^^^")
         return pushed
 
     def put_writes(self, task_id: str, writes: WritesT) -> None:
@@ -1252,9 +1235,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
     ) -> PregelExecutableTask | None:
         if pushed := super().accept_push(task, write_idx, call):
             for task in await self.amatch_cached_writes():
-                print("OUTPUT WRITES CALL POS 4")
                 self.output_writes(task.id, task.writes, cached=True)
-                print("^^^")
         return pushed
 
     def put_writes(self, task_id: str, writes: WritesT) -> None:

libs/langgraph/langgraph/pregel/main.py

@@ -2995,7 +2995,11 @@ class Pregel(
                         for task in await loop.amatch_cached_writes():
                             loop.output_writes(task.id, task.writes, cached=True)
                         async for _ in runner.atick(
-                            [t for t in loop.tasks.values() if not t.writes],
+                            [
+                                t
+                                for t in loop.tasks.values()
+                                if not t.writes and t.id not in loop.task_ids_to_block
+                            ],
                             timeout=self.step_timeout,
                             get_waiter=get_waiter,
                             schedule_task=loop.aaccept_push,

libs/langgraph/tests/test_interruption.py

@@ -328,9 +328,9 @@ async def test_interrupt_with_send_payloads_async(
     assert node_counter["map_node"] == 5
 
 
-@pytest.mark.xfail(reason="Duplicate interrupts written, still debugging this")
+# @pytest.mark.xfail(reason="Duplicate interrupts written, still debugging this")
 async def test_interrupt_with_send_payloads_sequential_resume_async(
-    async_checkpointer: BaseCheckpointSaver, durability: Durability
+    async_checkpointer: BaseCheckpointSaver,
 ) -> None:
     """Test interruption in map node with Send payloads and sequential resume."""