fix(langgraph): don't save null writes to checkpoint

libs/langgraph/langgraph/pregel/_loop.py

@@ -316,35 +316,36 @@ class PregelLoop:
             ]
             writes_to_save: WritesT = [
                 w[1:] for w in self.checkpoint_pending_writes if w[0] == task_id
-            ] + list(writes)
+            ] + [
+                (c, v) for c, v in writes if c != RESUME
+            ]
             self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes)
         else:
-            # build map of existing interrupts by interrupt id for this task
+            # build map of existing interrupts for this task for quick lookup
             existing_interrupts_by_id = {  # interrupt id -> list of interrupts
                 v[0].id: v
                 for tid, ch, v in self.checkpoint_pending_writes
                 if tid == task_id and ch == INTERRUPT
             }
-            # check if a resume write exists
-            has_resume = any(ch == RESUME for ch, _ in writes)
-
             writes_to_save = []
             for ch, v in writes:
                 if ch == INTERRUPT:
-                    # merge with existing interrupts if same interrupt id
+                    # we merge new interrupt writes with existing interrupts writes if they
+                    # occured within the same task (which means they have the same interrupt id)
                     new_interrupts = v if isinstance(v, list) else list(v)
                     if new_interrupts and (
                         existing := existing_interrupts_by_id.get(new_interrupts[0].id)
                     ):
-                        # found existing interrupts with same interrupt id
-                        # if a resume write exists, it means this is a new interrupt
-                        # so we are safe to merge it into the existing interrupt writes
-                        v = existing + new_interrupts if has_resume else existing
+                        # if the graph is invoked with None, we will hit the same interrupt
+                        # that was raised before, in this case we don't want to duplicate its write
+                        # so we just keep the existing checkpoint writes
+                        v = existing + new_interrupts if self.input is not None else existing
                     writes_to_save.append((ch, v))
                 else:
+                    # we add non-interrupt writes as-is
                     writes_to_save.append((ch, v))
 
-            # replace all writes for this task_id in one shot
+            # replace all writes for this task_id with the merged writes
             self.checkpoint_pending_writes = [
                 w for w in self.checkpoint_pending_writes if w[0] != task_id
             ] + [(task_id, c, v) for c, v in writes_to_save]
@@ -477,6 +478,8 @@ class PregelLoop:
             return False
 
         # prepare next tasks
+        print("checkpoint_pending_writes before: ", self.checkpoint_pending_writes)
+
         self.tasks = prepare_next_tasks(
             self.checkpoint,
             self.checkpoint_pending_writes,
@@ -496,6 +499,8 @@ class PregelLoop:
             cache_policy=self.cache_policy,
         )
 
+        print("checkpoint_pending_writes after: ", self.checkpoint_pending_writes)
+
         resume_map = self.config.get(CONF, {}).get(CONFIG_KEY_RESUME_MAP, {})
         if resume_map:
             skipped_interrupt_ids = self._pending_interrupts() - set(resume_map)

libs/langgraph/tests/test_interruption.py

@@ -582,3 +582,91 @@ async def test_node_with_multiple_interrupts_requires_full_resume_async(
     assert "input" in final_result
     assert final_result["input"] == "human_first-human_second-human_third"
     assert node_counter == 5
+
+def test_invoke_interrupted_graph_with_none(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    """Test that invoking an interrupted graph with None does not duplicate interrupt writes"""
+
+    node_counter = 0
+
+    class State(TypedDict):
+        input: str
+
+    def double_interrupt_node(state: State):
+        nonlocal node_counter
+        node_counter += 1
+        first = interrupt("first")
+        second = interrupt("second")
+        return {"input": f"{first}-{second}"}
+
+    builder = StateGraph(State)
+    builder.add_node("double_interrupt", double_interrupt_node)
+    builder.add_edge(START, "double_interrupt")
+    builder.add_edge("double_interrupt", END)
+
+    graph = builder.compile(checkpointer=sync_checkpointer)
+
+    config = {"configurable": {"thread_id": "test_none_resume"}}
+
+    result = graph.invoke({"input": "start"}, config=config)
+    first_history = list(graph.get_state_history(config))
+    interrupts = result.get("__interrupt__", [])
+    assert len(interrupts) == 1
+    assert node_counter == 1
+
+    # invoke with None. this should execute the node and the history should
+    # look the same as the first run
+    partial = graph.invoke(
+        None, config=config
+    )
+    second_history = list(graph.get_state_history(config))
+    remaining_interrupts = partial.get("__interrupt__", [])
+    assert len(remaining_interrupts) == 1
+    assert remaining_interrupts[0].value == "first"
+    assert node_counter == 2
+
+    # history should look the same for tasks and interrupts
+    print("first_history[0].interrupts: ", first_history[0].interrupts)
+    print("second_history[0].interrupts: ", second_history[0].interrupts)
+    print("first_history[0].tasks: ", first_history[0].tasks)
+    print("second_history[0].tasks: ", second_history[0].tasks)
+    assert first_history[0].interrupts == second_history[0].interrupts
+    assert first_history[0].tasks == second_history[0].tasks
+
+    # now resume the first interrupt with some value
+    partial = graph.invoke(
+        Command(resume="weet"), config=config
+    )
+    print("partial 3", partial)
+    third_history = list(graph.get_state_history(config))
+    remaining_interrupts = partial.get("__interrupt__", [])
+    assert len(remaining_interrupts) == 1
+    assert remaining_interrupts[0].value == "second"
+    assert node_counter == 3
+    
+    # invoke with None again. the history should look the same as
+    # the third run
+    partial = graph.invoke(
+        None, config=config
+    )
+    print("partial 4", partial)
+    fourth_history = list(graph.get_state_history(config))
+    remaining_interrupts = partial.get("__interrupt__", [])
+    assert len(remaining_interrupts) == 1
+    assert node_counter == 4
+
+    print("\nthird_history[0].interrupts: ", third_history[0].interrupts)
+    print("fourth_history[0].interrupts: ", fourth_history[0].interrupts)
+    print("third_history[0].tasks: ", third_history[0].tasks)
+    print("fourth_history[0].tasks: ", fourth_history[0].tasks)
+    assert third_history[0].interrupts == fourth_history[0].interrupts
+    assert third_history[0].tasks == fourth_history[0].tasks
+
+    # resume the graph once more with a real value
+    partial = graph.invoke(
+        Command(resume="bix"), config=config
+    )
+    remaining_interrupts = partial.get("__interrupt__", [])
+    assert len(remaining_interrupts) == 0
+    assert node_counter == 5