[Checkpointers] MemorySaver: refrain from overwriting writes (#2399)

libs/checkpoint/langgraph/checkpoint/memory/__init__.py

@@ -364,8 +364,12 @@ class MemorySaver(
         checkpoint_ns = config["configurable"]["checkpoint_ns"]
         checkpoint_id = config["configurable"]["checkpoint_id"]
         outer_key = (thread_id, checkpoint_ns, checkpoint_id)
+        outer_writes_ = self.writes.get(outer_key)
         for idx, (c, v) in enumerate(writes):
             inner_key = (task_id, WRITES_IDX_MAP.get(c, idx))
+            if inner_key[1] >= 0 and outer_writes_ and inner_key in outer_writes_:
+                continue
+
             self.writes[outer_key][inner_key] = (task_id, c, self.serde.dumps_typed(v))
 
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:

libs/checkpoint/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint"
-version = "2.0.2"
+version = "2.0.3"
 description = "Library with base interfaces for LangGraph checkpoint savers."
 authors = []
 license = "MIT"

libs/langgraph/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.2.45"
+version = "0.2.46"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"

libs/langgraph/tests/test_pregel.py

@@ -1115,6 +1115,86 @@ def test_fork_always_re_runs_nodes(
     ]
 
 
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_run_from_checkpoint_id_retains_previous_writes(
+    request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture
+) -> None:
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    class MyState(TypedDict):
+        myval: Annotated[int, operator.add]
+        otherval: bool
+
+    class Anode:
+        def __init__(self):
+            self.switch = False
+
+        def __call__(self, state: MyState):
+            self.switch = not self.switch
+            return {"myval": 2 if self.switch else 1, "otherval": self.switch}
+
+    builder = StateGraph(MyState)
+    thenode = Anode()  # Fun.
+    builder.add_node("node_one", thenode)
+    builder.add_node("node_two", thenode)
+    builder.add_edge(START, "node_one")
+
+    def _getedge(src: str):
+        swap = "node_one" if src == "node_two" else "node_two"
+
+        def _edge(st: MyState) -> Literal["__end__", "node_one", "node_two"]:
+            if st["myval"] > 3:
+                return END
+            if st["otherval"]:
+                return swap
+            return src
+
+        return _edge
+
+    builder.add_conditional_edges("node_one", _getedge("node_one"))
+    builder.add_conditional_edges("node_two", _getedge("node_two"))
+    graph = builder.compile(checkpointer=checkpointer)
+
+    thread_id = uuid.uuid4()
+    thread1 = {"configurable": {"thread_id": str(thread_id)}}
+
+    result = graph.invoke({"myval": 1}, thread1)
+    assert result["myval"] == 4
+    history = [c for c in graph.get_state_history(thread1)]
+
+    assert len(history) == 4
+    assert history[-1].values == {"myval": 0}
+    assert history[0].values == {"myval": 4, "otherval": False}
+
+    second_run_config = {
+        **thread1,
+        "configurable": {
+            **thread1["configurable"],
+            "checkpoint_id": history[1].config["configurable"]["checkpoint_id"],
+        },
+    }
+    second_result = graph.invoke(None, second_run_config)
+    assert second_result == {"myval": 5, "otherval": True}
+
+    new_history = [
+        c
+        for c in graph.get_state_history(
+            {"configurable": {"thread_id": str(thread_id), "checkpoint_ns": ""}}
+        )
+    ]
+
+    assert len(new_history) == len(history) + 1
+    for original, new in zip(history, new_history[1:]):
+        assert original.values == new.values
+        assert original.next == new.next
+        assert original.metadata["step"] == new.metadata["step"]
+
+    def _get_tasks(hist: list, start: int):
+        return [h.tasks for h in hist[start:]]
+
+    assert _get_tasks(new_history, 1) == _get_tasks(history, 0)
+
+
 def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")

libs/langgraph/tests/test_pregel_async.py

@@ -1960,6 +1960,85 @@ async def test_pending_writes_resume(
         )
 
 
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_run_from_checkpoint_id_retains_previous_writes(
+    request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture
+) -> None:
+    class MyState(TypedDict):
+        myval: Annotated[int, operator.add]
+        otherval: bool
+
+    class Anode:
+        def __init__(self):
+            self.switch = False
+
+        async def __call__(self, state: MyState):
+            self.switch = not self.switch
+            return {"myval": 2 if self.switch else 1, "otherval": self.switch}
+
+    builder = StateGraph(MyState)
+    thenode = Anode()  # Fun.
+    builder.add_node("node_one", thenode)
+    builder.add_node("node_two", thenode)
+    builder.add_edge(START, "node_one")
+
+    def _getedge(src: str):
+        swap = "node_one" if src == "node_two" else "node_two"
+
+        def _edge(st: MyState) -> Literal["__end__", "node_one", "node_two"]:
+            if st["myval"] > 3:
+                return END
+            if st["otherval"]:
+                return swap
+            return src
+
+        return _edge
+
+    builder.add_conditional_edges("node_one", _getedge("node_one"))
+    builder.add_conditional_edges("node_two", _getedge("node_two"))
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+        graph = builder.compile(checkpointer=checkpointer)
+
+        thread_id = uuid.uuid4()
+        thread1 = {"configurable": {"thread_id": str(thread_id)}}
+
+        result = await graph.ainvoke({"myval": 1}, thread1)
+        assert result["myval"] == 4
+        history = [c async for c in graph.aget_state_history(thread1)]
+
+        assert len(history) == 4
+        assert history[-1].values == {"myval": 0}
+        assert history[0].values == {"myval": 4, "otherval": False}
+
+        second_run_config = {
+            **thread1,
+            "configurable": {
+                **thread1["configurable"],
+                "checkpoint_id": history[1].config["configurable"]["checkpoint_id"],
+            },
+        }
+        second_result = await graph.ainvoke(None, second_run_config)
+        assert second_result == {"myval": 5, "otherval": True}
+
+        new_history = [
+            c
+            async for c in graph.aget_state_history(
+                {"configurable": {"thread_id": str(thread_id), "checkpoint_ns": ""}}
+            )
+        ]
+
+        assert len(new_history) == len(history) + 1
+        for original, new in zip(history, new_history[1:]):
+            assert original.values == new.values
+            assert original.next == new.next
+            assert original.metadata["step"] == new.metadata["step"]
+
+        def _get_tasks(hist: list, start: int):
+            return [h.tasks for h in hist[start:]]
+
+        assert _get_tasks(new_history, 1) == _get_tasks(history, 0)
+
+
 async def test_cond_edge_after_send() -> None:
     class Node:
         def __init__(self, name: str):