tests: add test for multiple interrupts and tasks

libs/langgraph/tests/test_pregel.py

@@ -1,4 +1,4 @@
 import enum
 import json
 import logging
 import operator
@@ -5300,3 +5300,61 @@
         {"node_a": [{"foo": "a1"}, {"foo": "a2"}]},
         {"node_b": {"foo": "b"}},
     ]
+
+
+def test_falsy_return_from_task() -> None:
+    """Test with a falsy return from a task."""
+    checkpointer = MemorySaver()
+
+    @task
+    def falsy_task() -> bool:
+        return False
+
+    @entrypoint(checkpointer=checkpointer)
+    def graph(state: dict) -> dict:
+        """React tool."""
+        task_result = falsy_task().result()
+        human_value = interrupt("test")
+
+    configurable = {"configurable": {"thread_id": uuid.uuid4()}}
+    graph.invoke({"a": 5}, configurable)
+    graph.invoke(Command(resume="123"), configurable)
+
+
+def test_multiple_interrupts_imperative() -> None:
+    """Test multiple interrupts with an imperative API."""
+    from langgraph.checkpoint.memory import MemorySaver
+    from langgraph.func import entrypoint, task
+
+    checkpointer = MemorySaver()
+    counter = 0
+
+    @task
+    def double(x: int) -> int:
+        """Increment the counter."""
+        nonlocal counter
+        counter += 1
+        return 2 * x
+
+    @entrypoint(checkpointer=checkpointer)
+    def graph(state: dict) -> dict:
+        """React tool."""
+
+        values = []
+
+        for idx in [1, 2, 3]:
+            values.extend([double(idx).result(), interrupt({"a": "boo"})])
+
+        return {"values": values}
+
+    configurable = {"configurable": {"thread_id": uuid.uuid4()}}
+    graph.invoke({}, configurable)
+    # Currently fails when double accepts a variable
+    graph.invoke(Command(resume="a"), configurable)
+    graph.invoke(Command(resume="b"), configurable)
+    result = graph.invoke(Command(resume="c"), configurable)
+    # `double` value should be cached appropriately when used w/ `interrupt`
+    assert result == {
+        "values": [2, "a", 4, "b", 6, "c"],
+    }
+    assert counter == 3