Merge pull request #2693 from langchain-ai/eugene/fix_test

langgraph[patch]: Fix unit test for Command(update)

libs/langgraph/langgraph/pregel/loop.py

@@ -422,18 +422,6 @@ class PregelLoop(LoopProtocol):
             self.status = "out_of_steps"
             return False
 
-        # apply NULL writes
-        if null_writes := [
-            w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID
-        ]:
-            mv_writes = apply_writes(
-                self.checkpoint,
-                self.channels,
-                [PregelTaskWrites((), INPUT, null_writes, [])],
-                self.checkpointer_get_next_version,
-            )
-            for key, values in mv_writes.items():
-                self._update_mv(key, values)
         # prepare next tasks
         self.tasks = prepare_next_tasks(
             self.checkpoint,
@@ -552,6 +540,18 @@ class PregelLoop(LoopProtocol):
             # save writes
             for tid, ws in writes.items():
                 self.put_writes(tid, ws)
+        # apply NULL writes
+        if null_writes := [
+            w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID
+        ]:
+            mv_writes = apply_writes(
+                self.checkpoint,
+                self.channels,
+                [PregelTaskWrites((), INPUT, null_writes, [])],
+                self.checkpointer_get_next_version,
+            )
+            for key, values in mv_writes.items():
+                self._update_mv(key, values)
         # proceed past previous checkpoint
         if is_resuming:
             self.checkpoint["versions_seen"].setdefault(INTERRUPT, {})

libs/langgraph/tests/test_pregel.py

@@ -14906,9 +14906,14 @@ def test_dict_mixed_return() -> None:
     assert graph.invoke({"foo": ""}) == {"foo": "ab"}
 
 
-def test_command_with_static_breakpoints() -> None:
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_command_with_static_breakpoints(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
     """Test that we can use Command to resume and update with static breakpoints."""
 
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
     class State(TypedDict):
         """The graph state."""
 
@@ -14930,20 +14935,13 @@ def test_command_with_static_breakpoints() -> None:
     builder.add_edge(START, "node1")
     builder.add_edge("node1", "node2")
 
-    # A checkpointer must be enabled for interrupts to work!
-    checkpointer = MemorySaver()
     graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"])
-
-    config = {
-        "configurable": {
-            "thread_id": uuid.uuid4(),
-        }
-    }
+    config = {"configurable": {"thread_id": str(uuid.uuid4())}}
 
     # Start the graph and interrupt at the first node
     graph.invoke({"foo": "abc"}, config)
-    result = graph.invoke(Command(resume="node1"), config)
-    assert result == {"foo": "abc|node-1|node-2"}
+    result = graph.invoke(Command(update={"foo": "def"}), config)
+    assert result == {"foo": "def|node-1|node-2"}
 
 
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)

libs/langgraph/tests/test_pregel_async.py

@@ -13201,6 +13201,41 @@ async def test_interrupt_loop(checkpointer_name: str):
         ]
 
 
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_command_with_static_breakpoints(checkpointer_name: str) -> None:
+    """Test that we can use Command to resume and update with static breakpoints."""
+
+    class State(TypedDict):
+        """The graph state."""
+
+        foo: str
+
+    def node1(state: State):
+        return {
+            "foo": state["foo"] + "|node-1",
+        }
+
+    def node2(state: State):
+        return {
+            "foo": state["foo"] + "|node-2",
+        }
+
+    builder = StateGraph(State)
+    builder.add_node("node1", node1)
+    builder.add_node("node2", node2)
+    builder.add_edge(START, "node1")
+    builder.add_edge("node1", "node2")
+
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+        graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"])
+        config = {"configurable": {"thread_id": str(uuid.uuid4())}}
+
+        # Start the graph and interrupt at the first node
+        await graph.ainvoke({"foo": "abc"}, config)
+        result = await graph.ainvoke(Command(update={"foo": "def"}), config)
+        assert result == {"foo": "def|node-1|node-2"}
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_multistep_plan(checkpointer_name: str):
     from langchain_core.messages import AnyMessage