tests: add test for multiple interrupts and tasks (#2941)

libs/langgraph/langgraph/pregel/io.py

@@ -10,6 +10,7 @@ from langgraph.constants import (
     EMPTY_SEQ,
     ERROR,
     INTERRUPT,
+    MISSING,
     NULL_TASK_ID,
     RESUME,
     RETURN,
@@ -173,7 +174,8 @@ def map_output_updates(
         return
     updated: list[tuple[str, Any]] = []
     for task, writes in output_tasks:
-        if rtn := next((value for chan, value in writes if chan == RETURN), None):
+        rtn = next((value for chan, value in writes if chan == RETURN), MISSING)
+        if rtn is not MISSING:
             updated.append((task.name, rtn))
         elif isinstance(output_channels, str):
             updated.extend(

libs/langgraph/langgraph/pregel/runner.py

@@ -1,6 +1,5 @@
 import asyncio
 import concurrent.futures
-import threading
 import time
 from functools import partial
 from typing import (
@@ -22,9 +21,11 @@ from langchain_core.callbacks import Callbacks
 from langgraph.constants import (
     CONF,
     CONFIG_KEY_CALL,
+    CONFIG_KEY_SCRATCHPAD,
     CONFIG_KEY_SEND,
     ERROR,
     INTERRUPT,
+    MISSING,
     NO_WRITES,
     PUSH,
     RESUME,
@@ -70,8 +71,6 @@ class PregelRunner:
         retry_policy: Optional[RetryPolicy] = None,
         get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None,
     ) -> Iterator[None]:
-        locks: dict[str, threading.Lock] = {}
-
         def writer(
             task: PregelExecutableTask,
             writes: Sequence[tuple[str, Any]],
@@ -81,24 +80,19 @@ class PregelRunner:
             if all(w[0] != PUSH for w in writes):
                 return task.config[CONF][CONFIG_KEY_SEND](writes)
 
-            if task.id not in locks:
-                locks[task.id] = threading.Lock()
-            with locks[task.id]:
-                prev_length = len(task.writes)
-                # delegate to the underlying writer
-                task.config[CONF][CONFIG_KEY_SEND](writes)
-                # confirm no other concurrent writes were added
-                assert len(task.writes) == prev_length + len(writes)
             # schedule PUSH tasks, collect futures
+            scratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD]
+            scratchpad.setdefault("call_counter", 0)
             rtn: dict[int, Optional[concurrent.futures.Future]] = {}
-            for idx, w in enumerate(writes, start=prev_length):
+            for idx, w in enumerate(writes):
                 # bail if not a PUSH write
                 if w[0] != PUSH:
                     continue
                 # schedule the next task, if the callback returns one
-                if next_task := self.schedule_task(
-                    task, idx, calls[idx - prev_length] if calls else None
-                ):
+                wcall = calls[idx] if calls else None
+                cnt = scratchpad["call_counter"]
+                scratchpad["call_counter"] += 1
+                if next_task := self.schedule_task(task, cnt, wcall):
                     if fut := next(
                         (
                             f
@@ -109,13 +103,19 @@ class PregelRunner:
                     ):
                         # if the parent task was retried,
                         # the next task might already be running
-                        rtn[idx - prev_length] = fut
+                        rtn[idx] = fut
                     elif next_task.writes:
                         # if it already ran, return the result
                         fut = concurrent.futures.Future()
-                        if val := next(v for c, v in next_task.writes if c == RETURN):
+                        if (
+                            val := next(
+                                (v for c, v in next_task.writes if c == RETURN), MISSING
+                            )
+                        ) and val is not MISSING:
                             fut.set_result(val)
-                        elif exc := next(v for c, v in next_task.writes if c == ERROR):
+                        elif exc := next(
+                            (v for c, v in next_task.writes if c == ERROR), None
+                        ):
                             fut.set_exception(
                                 exc
                                 if isinstance(exc, BaseException)
@@ -123,7 +123,7 @@ class PregelRunner:
                             )
                         else:
                             fut.set_result(None)
-                        rtn[idx - prev_length] = fut
+                        rtn[idx] = fut
                     else:
                         # schedule the next task
                         fut = self.submit(
@@ -141,7 +141,7 @@ class PregelRunner:
                         )
                         fut.add_done_callback(partial(self.commit, next_task))
                         futures[fut] = next_task
-                        rtn[idx - prev_length] = fut
+                        rtn[idx] = fut
             return [rtn.get(i) for i in range(len(writes))]
 
         def call(
@@ -189,6 +189,8 @@ class PregelRunner:
                     raise
             if not futures:  # maybe `t` schuduled another task
                 return
+            else:
+                tasks = ()  # don't reschedule this task
         # add waiter task if requested
         if get_waiter is not None:
             futures[get_waiter()] = None
@@ -255,8 +257,6 @@ class PregelRunner:
         retry_policy: Optional[RetryPolicy] = None,
         get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None,
     ) -> AsyncIterator[None]:
-        locks: dict[str, threading.Lock] = {}
-
         def writer(
             task: PregelExecutableTask,
             writes: Sequence[tuple[str, Any]],
@@ -266,23 +266,19 @@ class PregelRunner:
             if all(w[0] != PUSH for w in writes):
                 return task.config[CONF][CONFIG_KEY_SEND](writes)
 
-            if task.id not in locks:
-                locks[task.id] = threading.Lock()
-            with locks[task.id]:
-                prev_length = len(task.writes)
-                # delegate to the underlying writer
-                task.config[CONF][CONFIG_KEY_SEND](writes)
-                # confirm no other concurrent writes were added
-                assert len(task.writes) == prev_length + len(writes)
             # schedule PUSH tasks, collect futures
+            scratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD]
+            scratchpad.setdefault("call_counter", 0)
             rtn: dict[int, Optional[asyncio.Future]] = {}
-            for idx, w in enumerate(writes, start=prev_length):
+            for idx, w in enumerate(writes):
                 # bail if not a PUSH write
                 if w[0] != PUSH:
                     continue
                 # schedule the next task, if the callback returns one
-                wcall = calls[idx - prev_length] if calls is not None else None
-                if next_task := self.schedule_task(task, idx, wcall):
+                wcall = calls[idx] if calls is not None else None
+                cnt = scratchpad["call_counter"]
+                scratchpad["call_counter"] += 1
+                if next_task := self.schedule_task(task, cnt, wcall):
                     # if the parent task was retried,
                     # the next task might already be running
                     if fut := next(
@@ -295,13 +291,19 @@ class PregelRunner:
                     ):
                         # if the parent task was retried,
                         # the next task might already be running
-                        rtn[idx - prev_length] = fut
+                        rtn[idx] = fut
                     elif next_task.writes:
                         # if it already ran, return the result
                         fut = asyncio.Future()
-                        if val := next(v for c, v in next_task.writes if c == RETURN):
+                        if (
+                            val := next(
+                                (v for c, v in next_task.writes if c == RETURN), MISSING
+                            )
+                        ) and val is not MISSING:
                             fut.set_result(val)
-                        elif exc := next(v for c, v in next_task.writes if c == ERROR):
+                        elif exc := next(
+                            (v for c, v in next_task.writes if c == ERROR), None
+                        ):
                             fut.set_exception(
                                 exc
                                 if isinstance(exc, BaseException)
@@ -309,7 +311,7 @@ class PregelRunner:
                             )
                         else:
                             fut.set_result(None)
-                        rtn[idx - prev_length] = fut
+                        rtn[idx] = fut
                     else:
                         # schedule the next task
                         fut = cast(
@@ -333,7 +335,7 @@ class PregelRunner:
                         )
                         fut.add_done_callback(partial(self.commit, next_task))
                         futures[fut] = next_task
-                        rtn[idx - prev_length] = fut
+                        rtn[idx] = fut
             return [rtn.get(i) for i in range(len(writes))]
 
         def call(
@@ -388,6 +390,8 @@ class PregelRunner:
                     raise
             if not futures:  # maybe `t` schuduled another task
                 return
+            else:
+                tasks = ()  # don't reschedule this task
         # add waiter task if requested
         if get_waiter is not None:
             futures[get_waiter()] = None

libs/langgraph/langgraph/types.py

@@ -342,9 +342,12 @@ class LoopProtocol:
 
 
 class PregelScratchpad(TypedDict, total=False):
+    # interrupt
     interrupt_counter: int
     used_null_resume: bool
     resume: list[Any]
+    # call
+    call_counter: int
 
 
 def interrupt(value: Any) -> Any:

libs/langgraph/tests/test_pregel.py

@@ -5260,3 +5260,60 @@ def test_multiple_updates() -> None:
         {"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."""
+        falsy_task().result()
+        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)
+    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

libs/langgraph/tests/test_pregel_async.py

@@ -6693,3 +6693,69 @@ async def test_multiple_updates() -> None:
         {"node_a": [{"foo": "a1"}, {"foo": "a2"}]},
         {"node_b": {"foo": "b"}},
     ]
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="Python 3.11+ is required for async contextvars support",
+)
+async def test_falsy_return_from_task() -> None:
+    """Test with a falsy return from a task."""
+    checkpointer = MemorySaver()
+
+    @task
+    async def falsy_task() -> bool:
+        return False
+
+    @entrypoint(checkpointer=checkpointer)
+    async def graph(state: dict) -> dict:
+        """React tool."""
+        await falsy_task()
+        interrupt("test")
+
+    configurable = {"configurable": {"thread_id": uuid.uuid4()}}
+    await graph.ainvoke({"a": 5}, configurable)
+    await graph.ainvoke(Command(resume="123"), configurable)
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="Python 3.11+ is required for async contextvars support",
+)
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_multiple_interrupts_imperative(checkpointer_name: str) -> None:
+    """Test multiple interrupts with an imperative API."""
+    from langgraph.func import entrypoint, task
+
+    counter = 0
+
+    @task
+    async def double(x: int) -> int:
+        """Increment the counter."""
+        nonlocal counter
+        counter += 1
+        return 2 * x
+
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+
+        @entrypoint(checkpointer=checkpointer)
+        async def graph(state: dict) -> dict:
+            """React tool."""
+
+            values = []
+
+            for idx in [1, 2, 3]:
+                values.extend([await double(idx), interrupt({"a": "boo"})])
+
+            return {"values": values}
+
+        configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
+        await graph.ainvoke({}, configurable)
+        await graph.ainvoke(Command(resume="a"), configurable)
+        await graph.ainvoke(Command(resume="b"), configurable)
+        result = await graph.ainvoke(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