diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index 95273c78f..47d30de08 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -21,6 +21,7 @@ __all__ = ( "InvalidUpdateError", "GraphBubbleUp", "GraphInterrupt", + "NodeCancelledError", "NodeError", "NodeInterrupt", "NodeTimeoutError", @@ -164,6 +165,28 @@ class NodeError: """Exception raised by the failed node.""" +class NodeCancelledError(Exception): + """Raised when a node body raises ``asyncio.CancelledError`` itself. + + ``asyncio.CancelledError`` is a ``BaseException`` and the pregel runner + treats cancelled task futures as silent tear-down (e.g. when it stops + sibling tasks after a peer fails). That is the correct behaviour for + *framework-initiated* cancellation, but a user node that raises + ``asyncio.CancelledError`` from its own body should surface as a node + failure, the same way any other exception would. + + The retry layer converts user-raised ``asyncio.CancelledError`` into this + type so it flows through the normal error path and the run reports as + ``error`` instead of silently succeeding. + """ + + node: str + + def __init__(self, node: str, message: str | None = None) -> None: + super().__init__(message or f"Node {node!r} raised asyncio.CancelledError") + self.node = node + + class NodeTimeoutError(Exception): """Raised when a node invocation exceeds one of its configured timeouts. diff --git a/libs/langgraph/langgraph/pregel/_retry.py b/libs/langgraph/langgraph/pregel/_retry.py index acade9643..ecd485084 100644 --- a/libs/langgraph/langgraph/pregel/_retry.py +++ b/libs/langgraph/langgraph/pregel/_retry.py @@ -37,13 +37,23 @@ from langgraph._internal._constants import ( ) from langgraph._internal._runnable import create_task_in_config_context from langgraph._internal._timeout import sync_timeout_unsupported -from langgraph.errors import GraphBubbleUp, NodeTimeoutError, ParentCommand +from langgraph.errors import ( + GraphBubbleUp, + NodeCancelledError, + NodeTimeoutError, + ParentCommand, +) from langgraph.pregel.protocol import StreamProtocol from langgraph.runtime import ExecutionInfo, Runtime from langgraph.types import Command, PregelExecutableTask, RetryPolicy, TimeoutPolicy logger = logging.getLogger(__name__) SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11) +# `asyncio.Task.cancelling()` was added in Python 3.11. It reports the number of +# pending cancel requests on the task: ``0`` means no external code asked us to +# cancel — so a ``CancelledError`` observed here was raised by the task body +# itself (the user's node) rather than by pregel cancelling sibling tasks. +SUPPORTS_TASK_CANCELLING = sys.version_info >= (3, 11) def _timeout_secs(value: float | timedelta) -> float: @@ -302,6 +312,28 @@ class _IdleProgressCallbackHandler(BaseCallbackHandler): on_custom_event = _touch +def _is_user_raised_cancelled() -> bool: + """Return True if the in-flight ``CancelledError`` came from the task body. + + Pregel cancels sibling tasks via ``task.cancel()`` when a peer fails, which + increments ``asyncio.Task.cancelling()`` on the target before the cancel + actually fires. A user node that calls ``raise asyncio.CancelledError()`` + from inside its own body raises while ``cancelling() == 0``, which is the + signal we use to convert the exception into a regular + :class:`NodeCancelledError`. + + Returns ``False`` when we can't tell (``cancelling()`` unavailable, or no + current task — neither should happen in practice from ``arun_with_retry``) + so framework-initiated cancellation continues to propagate unchanged. + """ + if not SUPPORTS_TASK_CANCELLING: + return False + current = asyncio.current_task() + if current is None: + return False + return current.cancelling() == 0 + + def _drain_cancelled(task: asyncio.Task[Any]) -> None: # Mark the abandoned task's exception as retrieved so asyncio doesn't log it. with suppress(asyncio.CancelledError): @@ -600,6 +632,12 @@ def run_with_retry( except GraphBubbleUp: # if interrupted, end raise + except asyncio.CancelledError as exc: + # A sync node has no asyncio context, so any ``CancelledError`` that + # reaches here was raised by the node body itself. Surface it as a + # regular exception so the pregel runner panics the run instead of + # treating the task as a silent tear-down (LSD-1507). + raise NodeCancelledError(task.name) from exc except Exception as exc: if SUPPORTS_EXC_NOTES: exc.add_note(f"During task with name '{task.name}' and id '{task.id}'") @@ -736,6 +774,24 @@ async def arun_with_retry( # if interrupted, end _finish_timed_attempt(config, attempt_ctx) raise + except asyncio.CancelledError as exc: + # ``CancelledError`` reaches us in two very different shapes: + # 1. Pregel cancelled this task because a sibling failed + # (``asyncio.Task.cancelling() >= 1``). The framework already + # knows the run is failing and we must let cancellation + # propagate so the watchdog/cleanup code in the runner sees a + # cancelled future. + # 2. The node body itself raised ``asyncio.CancelledError`` ( + # ``cancelling() == 0``). The runner would otherwise treat + # this as silent tear-down and the run would report + # ``success`` even though the node failed (LSD-1507). Convert + # it into :class:`NodeCancelledError` so it follows the same + # path as any other node failure. + if _is_user_raised_cancelled(): + _finish_timed_attempt(config, attempt_ctx, exc) + raise NodeCancelledError(task.name) from exc + _finish_timed_attempt(config, attempt_ctx, exc) + raise except Exception as exc: _finish_timed_attempt(config, attempt_ctx, exc) if SUPPORTS_EXC_NOTES: diff --git a/libs/langgraph/tests/test_retry.py b/libs/langgraph/tests/test_retry.py index ee54f4b39..fcdde937b 100644 --- a/libs/langgraph/tests/test_retry.py +++ b/libs/langgraph/tests/test_retry.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import operator import sys import threading @@ -35,7 +36,13 @@ from langgraph._internal._runnable import RunnableCallable from langgraph._internal._timeout import coerce_timeout_policy from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue -from langgraph.errors import GraphInterrupt, NodeError, NodeTimeoutError, ParentCommand +from langgraph.errors import ( + GraphInterrupt, + NodeCancelledError, + NodeError, + NodeTimeoutError, + ParentCommand, +) from langgraph.func import entrypoint, task from langgraph.graph import END, START, StateGraph, add_messages from langgraph.pregel import NodeBuilder, Pregel @@ -63,6 +70,18 @@ NEEDS_CONTEXTVARS = pytest.mark.skipif( reason="Python 3.11+ is required for async contextvars support", ) +# `asyncio.Task.cancelling()` is Python 3.11+. The LSD-1507 fix in +# `langgraph/pregel/_retry.py` falls back to a no-op on 3.10 (preserves the +# existing CancelledError-as-silent-tear-down behaviour) because there is no +# reliable way to distinguish user-raised from framework-initiated +# cancellation without that API. Tests for the converted behaviour gate on +# the same Python version boundary. +NEEDS_TASK_CANCELLING = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="LSD-1507 user-cancellation conversion requires Python 3.11+ " + "(asyncio.Task.cancelling)", +) + def test_should_retry_on_single_exception(): """Test retry with a single exception type.""" @@ -2802,3 +2821,123 @@ def test_error_handler_resumes_after_crash_multiple_nodes(): assert call_count["handler_b"] == 2 # ran again on resume assert "recovered_a:a" in result["results"] assert "recovered_b:b" in result["results"] + + +@NEEDS_TASK_CANCELLING +@pytest.mark.anyio +async def test_arun_with_retry_user_raised_cancelled_becomes_node_cancelled(): + class UserCancelsProc: + async def ainvoke(self, input, config): + raise asyncio.CancelledError("nope") + + task = _make_task(UserCancelsProc(), name="user-cancel") + + with pytest.raises(NodeCancelledError) as excinfo: + await arun_with_retry(task, retry_policy=None) + assert excinfo.value.node == "user-cancel" + # original CancelledError chained for debugging + assert isinstance(excinfo.value.__cause__, asyncio.CancelledError) + + +@NEEDS_TASK_CANCELLING +@pytest.mark.anyio +async def test_arun_with_retry_user_raised_cancelled_with_timeout_policy(): + """The timeout path runs the node in a child task; the conversion must + still trigger for user-raised ``CancelledError``.""" + + class UserCancelsProc: + async def ainvoke(self, input, config): + raise asyncio.CancelledError + + task = _make_task( + UserCancelsProc(), timeout=_idle_timeout(1.0), name="user-cancel-timed" + ) + + with pytest.raises(NodeCancelledError) as excinfo: + await arun_with_retry(task, retry_policy=None) + assert excinfo.value.node == "user-cancel-timed" + + +def test_run_with_retry_sync_node_raising_cancelled_becomes_node_cancelled(): + class SyncUserCancelsProc: + def invoke(self, input, config): + raise asyncio.CancelledError("sync nope") + + task = _make_task(SyncUserCancelsProc(), timeout=None, name="sync-user-cancel") + + with pytest.raises(NodeCancelledError) as excinfo: + run_with_retry(task, retry_policy=None) + assert excinfo.value.node == "sync-user-cancel" + assert isinstance(excinfo.value.__cause__, asyncio.CancelledError) + + +@NEEDS_TASK_CANCELLING +@pytest.mark.anyio +async def test_arun_with_retry_external_cancel_propagates_as_cancelled(): + """When the asyncio task running ``arun_with_retry`` is cancelled from the + outside, the cancellation must still propagate as + ``asyncio.CancelledError``. Converting it to ``NodeCancelledError`` would + break the runner's ability to cancel sibling tasks during cleanup.""" + + started = asyncio.Event() + observed: list[BaseException] = [] + + class SlowProc: + async def ainvoke(self, input, config): + started.set() + await asyncio.sleep(10.0) + return "never" + + task = _make_task(SlowProc(), timeout=None, name="external-cancel") + + async def runner(): + try: + await arun_with_retry(task, retry_policy=None) + except BaseException as exc: + observed.append(exc) + raise + + bg = asyncio.create_task(runner()) + await started.wait() + bg.cancel() + # We expect the cancellation to surface to us as well; swallow it here so + # the test runner's own task isn't poisoned by the cancel. + with contextlib.suppress(asyncio.CancelledError): + await bg + assert observed, "runner did not observe any exception" + # Framework cancellation must remain a CancelledError, not be rewritten as + # NodeCancelledError. + assert isinstance(observed[0], asyncio.CancelledError) + assert not isinstance(observed[0], NodeCancelledError) + + +@NEEDS_TASK_CANCELLING +@pytest.mark.anyio +async def test_pregel_user_raised_cancellederror_fails_run(): + """End-to-end: a two-branch graph where one branch raises + ``asyncio.CancelledError`` must fail the run instead of returning + a partial-success state. This is the LSD-1507 customer scenario.""" + + class _S(TypedDict, total=False): + vals: Annotated[list[str], operator.add] + + async def ok(state: _S) -> _S: + return {"vals": ["ok"]} + + async def boom(state: _S) -> _S: + raise asyncio.CancelledError("user-raised in node body") + + graph = ( + StateGraph(_S) + .add_node("ok", ok) + .add_node("boom", boom) + .add_edge(START, "ok") + .add_edge(START, "boom") + .add_edge("ok", END) + .add_edge("boom", END) + .compile() + ) + + with pytest.raises(NodeCancelledError) as excinfo: + await graph.ainvoke({"vals": []}) + assert excinfo.value.node == "boom"