style(langgraph): refactor and fml

libs/langgraph/langgraph/pregel/_loop.py

@@ -319,35 +319,36 @@ class PregelLoop:
             ] + list(writes)
             self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes)
         else:
-            writes_to_save: WritesT = []
-            for channel, value in writes:
-                # 
-                if channel == INTERRUPT: 
-                    new_interrupts = list(value) if isinstance(value, (list, tuple)) else [value]
-                    # aggregate existing interrupts for this task
-                    existing = next(
-                        (v for tid, ch, v in self.checkpoint_pending_writes
-                        if tid == task_id and ch == INTERRUPT),
-                        None
-                    )
-                    if existing is not None:
-                        # backwards compat: support resuming tasks where saved interrupt value is not a list
-                        existing_interrupts = list(existing) if isinstance(existing, (list, tuple)) else [existing]
-                        # append if same interrupt id (multiple interrupt() calls in same task execution),
-                        # otherwise replace (different PUSH tasks, each with unique interrupt id - see tests/test_pregel.py::test_interrupt_task_functional)
-                        value = (
-                            existing_interrupts + new_interrupts
-                            if existing_interrupts[0].id == new_interrupts[0].id
-                            else new_interrupts
-                        )
-                    else:
-                        value = new_interrupts
-                writes_to_save.append((channel, value))
-            # remove existing writes for this task
+
+            def _merge_interrupts(task_id: str, value: tuple[str, Any]) -> WritesT:
+                """
+                Normalize to list and merge with any existing INTERRUPT writes
+                for the same task_id. If the interrupt id matches, append; else replace.
+                """
+                new = value if isinstance(value, list) else list(value)
+                existing = next(
+                    (
+                        v
+                        for tid, ch, v in self.checkpoint_pending_writes
+                        if tid == task_id and ch == INTERRUPT
+                    ),
+                    None,
+                )
+                if existing is None:
+                    return new
+                old = existing if isinstance(existing, list) else list(existing)
+                return old + new if old and new and old[0].id == new[0].id else new
+
+            writes_to_save = [
+                # aggregate existing interrupts for this task
+                (ch, _merge_interrupts(task_id, v) if ch == INTERRUPT else v)
+                for ch, v in writes
+            ]
+
+            # replace all writes for this task_id in one shot
             self.checkpoint_pending_writes = [
                 w for w in self.checkpoint_pending_writes if w[0] != task_id
-            ]
-            self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes_to_save)
+            ] + [(task_id, c, v) for c, v in writes_to_save]
 
         if self.durability != "exit" and self.checkpointer_put_writes is not None:
             config = patch_configurable(
@@ -505,7 +506,7 @@ class PregelLoop:
                 if channel == INTERRUPT
                 # interrupts within a task are uncovered sequentially as resumes are provided,
                 # so we only need to check the last interrupt id
-                and (list(value) if isinstance(value, (list, tuple)) else [value])[-1].id in skipped_interrupt_ids
+                and value[-1].id in skipped_interrupt_ids
             }
         else:
             self.skipped_task_ids = set()
@@ -567,28 +568,29 @@ class PregelLoop:
                 if task_id in self.skipped_task_ids and channel == INTERRUPT:
                     # find resume count for this task
                     resumes = next(
-                        (v for tid, ch, v in self.checkpoint_pending_writes
-                        if tid == task_id and ch == RESUME),
-                        None
+                        (
+                            v
+                            for tid, ch, v in self.checkpoint_pending_writes
+                            if tid == task_id and ch == RESUME
+                        ),
+                        None,
                     )
                     resume_count = len(resumes) if resumes is not None else 0
-                    interrupt_list = list(value) if isinstance(value, (list, tuple)) else [value]
                     # only output unresumed interrupts
-                    if resume_count < len(interrupt_list):
-                        self.output_writes(task_id, [(INTERRUPT, interrupt_list[resume_count:])])
+                    if resume_count < len(value):
+                        self.output_writes(task_id, [(INTERRUPT, value[resume_count:])])
 
         return True
 
     def after_tick(self) -> None:
         if self.skipped_task_ids:
             # raise early GraphInterrupt for skipped tasks.
-            # since we know len(resumes) != len(interrupts) for these tasks, we 
+            # since we know len(resumes) != len(interrupts) for these tasks, we
             # can prevent unecessary node re-execution by raising preemptively
             interrupts = []
             for task_id, channel, value in self.checkpoint_pending_writes:
                 if channel == INTERRUPT and task_id in self.skipped_task_ids:
-                    interrupt_list = list(value) if isinstance(value, (list, tuple)) else [value]
-                    interrupts.extend(interrupt_list)
+                    interrupts.extend(value)
             if interrupts:
                 raise GraphInterrupt(interrupts)
 
@@ -651,10 +653,11 @@ class PregelLoop:
 
         for task_id, channel, value in self.checkpoint_pending_writes:
             if channel == INTERRUPT:
-                interrupt_list = list(value) if isinstance(value, (list, tuple)) else [value]
-                pending_interrupts[task_id] = (interrupt_list[0].id, len(interrupt_list))
+                pending_interrupts[task_id] = (
+                    value[0].id,
+                    len(value),
+                )
             elif channel == RESUME:
-                # count resume values for this task
                 resume_list = value if isinstance(value, list) else [value]
                 pending_resumes[task_id] = len(resume_list)
 
@@ -1095,7 +1098,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
 
     def put_writes(self, task_id: str, writes: WritesT) -> None:
         """Put writes for a task, to be read by the next tick."""
-        
+
         super().put_writes(task_id, writes)
         if not writes or self.cache is None or not hasattr(self, "tasks"):
             return

libs/langgraph/tests/test_interruption.py

@@ -467,7 +467,8 @@ def test_node_with_multiple_interrupts_requires_full_resume(
     # invoke with an interrupt map that DOES NOT match double_interrupt_node.
     # this should not execute the node because the optimization kicks in
     partial = graph.invoke(
-        Command(resume={"00000000000000000000000000000000": "nothing_burger"}), config=config
+        Command(resume={"00000000000000000000000000000000": "nothing_burger"}),
+        config=config,
     )
     remaining_interrupts = partial.get("__interrupt__", [])
     assert len(remaining_interrupts) == 1
@@ -475,18 +476,14 @@ def test_node_with_multiple_interrupts_requires_full_resume(
     assert node_counter == 2
 
     # invoke with None resume. this should execute the node
-    partial = graph.invoke(
-        None, config=config
-    )
+    partial = graph.invoke(None, config=config)
     remaining_interrupts = partial.get("__interrupt__", [])
     assert len(remaining_interrupts) == 1
     assert remaining_interrupts[0].value == "second"
     assert node_counter == 3
 
     # invoke with nonspecific resume. this should execute the node
-    partial = graph.invoke(
-        Command(resume="human_second"), config=config
-    )
+    partial = graph.invoke(Command(resume="human_second"), config=config)
     remaining_interrupts = partial.get("__interrupt__", [])
     assert len(remaining_interrupts) == 1
     print("REMAINING INTERRUPTS: ", remaining_interrupts)
@@ -495,9 +492,7 @@ def test_node_with_multiple_interrupts_requires_full_resume(
 
     # finally, invoke with an interrupt map that matches double_interrupt_node.
     # this should execute the node and all interrupts should be resolved
-    final_result = graph.invoke(
-        Command(resume="human_third"), config=config
-    )
+    final_result = graph.invoke(Command(resume="human_third"), config=config)
     assert "input" in final_result
     assert final_result["input"] == "human_first-human_second-human_third"
     assert node_counter == 5
@@ -510,7 +505,7 @@ async def test_node_with_multiple_interrupts_requires_full_resume_async(
 
     Ensures that a node is not re-executed until valid resume values have been provided to all
     discovered interrupts"""
-    
+
     node_counter = 0
 
     class State(TypedDict):
@@ -553,7 +548,8 @@ async def test_node_with_multiple_interrupts_requires_full_resume_async(
     # invoke with an interrupt map that DOES NOT match double_interrupt_node.
     # this should not execute the node because the optimization kicks in
     partial = await graph.ainvoke(
-        Command(resume={"00000000000000000000000000000000": "nothing_burger"}), config=config
+        Command(resume={"00000000000000000000000000000000": "nothing_burger"}),
+        config=config,
     )
     remaining_interrupts = partial.get("__interrupt__", [])
     assert len(remaining_interrupts) == 1
@@ -561,18 +557,14 @@ async def test_node_with_multiple_interrupts_requires_full_resume_async(
     assert node_counter == 2
 
     # invoke with None resume. this should execute the node
-    partial = await graph.ainvoke(
-        None, config=config
-    )
+    partial = await graph.ainvoke(None, config=config)
     remaining_interrupts = partial.get("__interrupt__", [])
     assert len(remaining_interrupts) == 1
     assert remaining_interrupts[0].value == "second"
     assert node_counter == 3
 
     # invoke with nonspecific resume. this should execute the node
-    partial = await graph.ainvoke(
-        Command(resume="human_second"), config=config
-    )
+    partial = await graph.ainvoke(Command(resume="human_second"), config=config)
     remaining_interrupts = partial.get("__interrupt__", [])
     assert len(remaining_interrupts) == 1
     print("REMAINING INTERRUPTS: ", remaining_interrupts)
@@ -581,9 +573,7 @@ async def test_node_with_multiple_interrupts_requires_full_resume_async(
 
     # finally, invoke with an interrupt map that matches double_interrupt_node.
     # this should execute the node and all interrupts should be resolved
-    final_result = await graph.ainvoke(
-        Command(resume="human_third"), config=config
-    )
+    final_result = await graph.ainvoke(Command(resume="human_third"), config=config)
     assert "input" in final_result
     assert final_result["input"] == "human_first-human_second-human_third"
     assert node_counter == 5