mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-10 11:47:51 +02:00
fix(langgraph): [LSD-1507] Distinguish between user cancelled and other cancellations (#7920)
- Distinguish between Node cancellations and other cancellations - Use python 3.11+ feature where `cancelling() == 0` when it is the node cancelling - Bubble up the node cancellation example so the client can take care of it instead of silently failing without reporting it. Read the full contributing guidelines: https://docs.langchain.com/oss/python/contributing/overview > **All contributions must be in English.** See the [language policy](https://docs.langchain.com/oss/python/contributing/overview#language-policy). If you paste a large clearly AI generated description here your PR may be IGNORED or CLOSED! Thank you for contributing to LangGraph! Follow these steps to have your pull request considered as ready for review. 1. PR title: Should follow the format: TYPE(SCOPE): DESCRIPTION - feat(langgraph): add multi-tenant support - Allowed TYPE and SCOPE values: https://github.com/langchain-ai/langgraph/blob/main/.github/workflows/pr_lint.yml#L19-L43 2. PR description: - Write 1-2 sentences summarizing the change. - The `Fixes #xx` line at the top is **required** for external contributions — update the issue number and keep the keyword. This links your PR to the approved issue and auto-closes it on merge. - If there are any breaking changes, please clearly describe them. - If this PR depends on another PR being merged first, please include "Depends on #PR_NUMBER" in the description. 3. Run `make format`, `make lint` and `make test` from the root of the package(s) you've modified. - We will not consider a PR unless these three are passing in CI. 4. How did you verify your code works? - Existing unit tests pass - New unit tests added - Used langgraph deployment to make sure the feature is working as expected. Additional guidelines: - All external PRs must link to an issue or discussion where a solution has been approved by a maintainer, and you must be assigned to that issue. PRs without prior approval will be closed. - PRs should not touch more than one package unless absolutely necessary. - Do not update the `uv.lock` files or add dependencies to `pyproject.toml` files (even optional ones) unless you have explicit permission to do so by a maintainer. ## Social handles (optional) <!-- If you'd like a shoutout on release, add your socials below --> Twitter: @ LinkedIn: https://linkedin.com/in/
This commit is contained in:
@@ -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.
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user