Add test for cancellation

libs/langgraph/langgraph/pregel/algo.py

@@ -43,6 +43,7 @@ from langgraph.constants import (
     CONFIG_KEY_TASK_ID,
     CONFIG_KEY_WRITES,
     EMPTY_SEQ,
+    ERROR,
     INTERRUPT,
     NO_WRITES,
     NS_END,
@@ -270,7 +271,7 @@ def apply_writes(
     pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list)
     for task in tasks:
         for chan, val in task.writes:
-            if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN):
+            if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR):
                 pass
             elif chan == TASKS:  # TODO: remove branch in 1.0
                 checkpoint["pending_sends"].append(val)

libs/langgraph/langgraph/pregel/runner.py

@@ -441,7 +441,12 @@ class PregelRunner:
     ) -> None:
         if fut is not None:
             exception = _exception(fut)
-        if exception:
+        if isinstance(exception, asyncio.CancelledError):
+            # for cancelled tasks, also save error in task,
+            # so loop can finish super-step
+            task.writes.append((ERROR, exception))
+            self.put_writes(task.id, task.writes)
+        elif exception:
             if isinstance(exception, GraphInterrupt):
                 # save interrupt to checkpointer
                 if interrupts := [(INTERRUPT, i) for i in exception.args[0]]:
@@ -472,11 +477,12 @@ def _should_stop_others(
     GraphInterrupts are not considered failures."""
     for fut in done:
         if fut.cancelled():
-            return True
-        if exc := fut.exception():
-            return not isinstance(exc, GraphBubbleUp)
-    else:
-        return False
+            continue
+        elif exc := fut.exception():
+            if not isinstance(exc, GraphBubbleUp):
+                return True
+
+    return False
 
 
 def _exception(
@@ -502,7 +508,9 @@ def _panic_or_proceed(
     done: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set()
     inflight: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set()
     for fut in futs:
-        if fut.done():
+        if fut.cancelled():
+            continue
+        elif fut.done():
             done.add(fut)
         else:
             inflight.add(fut)
@@ -515,8 +523,6 @@ def _panic_or_proceed(
             # raise the exception
             if panic:
                 raise exc
-            else:
-                return
     if inflight:
         # if we got here means we timed out
         while inflight:

libs/langgraph/tests/test_pregel_async.py

@@ -2648,6 +2648,10 @@ async def test_send_sequences(checkpointer_name: str) -> None:
         ]
 
 
+@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_imp_task(checkpointer_name: str) -> None:
     async with awith_checkpointer(checkpointer_name) as checkpointer:
@@ -2690,6 +2694,64 @@ async def test_imp_task(checkpointer_name: str) -> None:
         assert mapper_calls == 2
 
 
+@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_imp_task_cancel(checkpointer_name: str) -> None:
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+        mapper_calls = 0
+        mapper_cancels = 0
+
+        @task()
+        async def mapper(input: int) -> str:
+            nonlocal mapper_calls, mapper_cancels
+            mapper_calls += 1
+            try:
+                await asyncio.sleep(1)
+            except asyncio.CancelledError:
+                mapper_cancels += 1
+                raise
+            return str(input) * 2
+
+        @entrypoint(checkpointer=checkpointer)
+        async def graph(input: list[int]) -> list[str]:
+            futures = [mapper(i) for i in input]
+            await asyncio.sleep(0.1)
+            futures.pop().cancel()  # cancel one
+            mapped = await asyncio.gather(*futures)
+            answer = interrupt("question")
+            return [m + answer for m in mapped]
+
+        thread1 = {"configurable": {"thread_id": "1"}}
+        assert [c async for c in graph.astream([0, 1], thread1)] == [
+            {"mapper": "00"},
+            {
+                "__interrupt__": (
+                    Interrupt(
+                        value="question",
+                        resumable=True,
+                        ns=[AnyStr("graph:")],
+                        when="during",
+                    ),
+                )
+            },
+        ]
+        assert mapper_calls == 2
+        assert mapper_cancels == 1
+
+        assert await graph.ainvoke(Command(resume="answer"), thread1) == [
+            "00answer",
+        ]
+        assert mapper_calls == 3
+        assert mapper_cancels == 2
+
+
+@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_imp_sync_from_async(checkpointer_name: str) -> None:
     async with awith_checkpointer(checkpointer_name) as checkpointer:
@@ -2722,6 +2784,10 @@ async def test_imp_sync_from_async(checkpointer_name: str) -> None:
         ]
 
 
+@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_imp_stream_order(checkpointer_name: str) -> None:
     async with awith_checkpointer(checkpointer_name) as checkpointer: