diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py index 2072125d1..5eb14a6fa 100644 --- a/libs/langgraph/langgraph/stream/run_stream.py +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -132,13 +132,23 @@ class GraphRunStream: def abort(self) -> None: """Stop the run early. - Closes the mux and marks the stream exhausted. The graph - iterator is dropped; any in-flight nodes see the closure on - their next yield point. Idempotent. + Closes the underlying graph iterator (propagating `GeneratorExit` + so in-flight nodes and subgraphs are cancelled), closes the mux, + and marks the stream exhausted. Idempotent. """ if self._exhausted: return self._exhausted = True + graph_iter = self._graph_iter + self._graph_iter = None + if ( + graph_iter is not None + and (close := getattr(graph_iter, "close", None)) is not None + ): + try: + close() + except Exception: + pass try: self._mux.close() except Exception: @@ -348,6 +358,8 @@ class AsyncGraphRunStream: self._scope_list: list[str] = list(mux.scope) self._pump_cond = asyncio.Condition() self._pumping = False + self._anext_task: asyncio.Future[Any] | None = None + self._aborting = False for key in mux.native_keys: setattr(self, key, mux.extensions[key]) if wire_pump: @@ -407,7 +419,25 @@ class AsyncGraphRunStream: try: try: - part = await self._graph_aiter.__anext__() + # Run the pull as a child task so `abort()` can cancel it + # mid-flight. Cancelling propagates `CancelledError` into the + # graph generator frame -> Pregel loop -> nested subgraph + # nodes, which a bare `aclose()` cannot do while the generator + # is running ("asynchronous generator is already running"). + self._anext_task = asyncio.ensure_future(self._graph_aiter.__anext__()) + try: + part = await self._anext_task + except asyncio.CancelledError: + if self._aborting: + # Abort-initiated cancel: stop gracefully. + self._exhausted = True + return False + # Genuine external cancel of this task: also stop the + # in-flight pull, then propagate. + self._anext_task.cancel() + raise + finally: + self._anext_task = None event = convert_to_protocol_event(part) self._observe_event(event) await self._mux.apush(event) @@ -428,15 +458,40 @@ class AsyncGraphRunStream: async def abort(self) -> None: """Stop the run early. - Marks the stream exhausted, wakes any pump-waiters, and closes - the mux. Any `apush` blocked on backpressure wakes and returns - without appending. Idempotent. + Marks the stream exhausted and wakes any pump-waiters. Cancels an + in-flight pull if one is running, then closes the underlying graph + iterator, so running nodes and nested subgraphs are cancelled + whether or not a pump is mid-pull. Closes the mux; any `apush` + blocked on backpressure wakes and returns without appending. + Idempotent. """ async with self._pump_cond: if self._exhausted: return self._exhausted = True + self._aborting = True + graph_aiter = self._graph_aiter + self._graph_aiter = None + anext_task = self._anext_task self._pump_cond.notify_all() + # If a pump is mid-pull, cancel it so the cancellation propagates + # into running nodes and nested subgraphs. Once it settles the + # generator is no longer running, so the `aclose()` below is a safe + # final cleanup (and handles the no-in-flight-pull case directly). + if anext_task is not None and not anext_task.done(): + anext_task.cancel() + try: + await anext_task + except (asyncio.CancelledError, Exception): + pass + if ( + graph_aiter is not None + and (aclose := getattr(graph_aiter, "aclose", None)) is not None + ): + try: + await aclose() + except Exception: + pass try: await self._mux.aclose() except Exception: diff --git a/libs/langgraph/tests/test_pregel_stream_events_v3.py b/libs/langgraph/tests/test_pregel_stream_events_v3.py index 5ff4a03af..a0b038c09 100644 --- a/libs/langgraph/tests/test_pregel_stream_events_v3.py +++ b/libs/langgraph/tests/test_pregel_stream_events_v3.py @@ -607,6 +607,164 @@ class TestStreamV2Async: _ = await anext(aiter(run.values)) assert run._exhausted is True + async def test_abort_cancels_running_subgraph(self) -> None: + class CountState(TypedDict): + count: int + + runs: list[int] = [] + + async def sub_node(state: CountState) -> dict: + runs.append(state["count"] + 1) + await asyncio.sleep(0.05) + return {"count": state["count"] + 1} + + sub_graph = ( + StateGraph(CountState) + .add_node("sub_node", sub_node) + .set_entry_point("sub_node") + .add_conditional_edges( + "sub_node", + lambda s: END if s["count"] >= 10 else "sub_node", + ) + .compile() + ) + + async def main_node(state: CountState) -> None: + await sub_graph.ainvoke({"count": 0}) + + main_graph = ( + StateGraph(CountState) + .add_node("main_node", main_node) + .set_entry_point("main_node") + .compile() + ) + + run = await main_graph.astream_events({"count": 0}, version="v3") + async for e in run: + if ( + e["method"] == "values" + and e["params"]["namespace"] + and e["params"]["data"]["count"] >= 2 + ): + break + await run.abort() + runs_at_abort = len(runs) + # Give the (now-cancelled) subgraph a chance to keep looping. + await asyncio.sleep(0.3) + assert len(runs) == runs_at_abort + assert len(runs) < 10 + + async def test_abort_cancels_deeply_nested_subgraph(self) -> None: + class CountState(TypedDict): + count: int + + runs: list[int] = [] + + async def deep_node(state: CountState) -> dict: + runs.append(state["count"] + 1) + await asyncio.sleep(0.05) + return {"count": state["count"] + 1} + + # Deepest graph loops until count >= 10. + graph = ( + StateGraph(CountState) + .add_node("deep_node", deep_node) + .set_entry_point("deep_node") + .add_conditional_edges( + "deep_node", + lambda s: END if s["count"] >= 10 else "deep_node", + ) + .compile() + ) + + # Wrap it three times: graph -> subgraph -> subgraph -> subgraph. + for _ in range(3): + + async def caller(state: CountState, _child: Any = graph) -> dict: + return await _child.ainvoke({"count": 0}) + + graph = ( + StateGraph(CountState) + .add_node("caller", caller) + .set_entry_point("caller") + .compile() + ) + + run = await graph.astream_events({"count": 0}, version="v3") + async for e in run: + if ( + e["method"] == "values" + and e["params"]["namespace"] + and e["params"]["data"]["count"] >= 2 + ): + break + await run.abort() + runs_at_abort = len(runs) + # Give the (now-cancelled) nested subgraph a chance to keep looping. + await asyncio.sleep(0.3) + assert len(runs) == runs_at_abort + assert len(runs) < 10 + + async def test_abort_cancels_subgraph_during_inflight_pump(self) -> None: + class CountState(TypedDict): + count: int + + started = asyncio.Event() + cancelled = asyncio.Event() + + async def sub_node(state: CountState) -> dict: + started.set() + try: + # Long-running node: still in flight when abort fires. + await asyncio.sleep(5) + except asyncio.CancelledError: + cancelled.set() + raise + return {"count": state["count"] + 1} + + sub_graph = ( + StateGraph(CountState) + .add_node("sub_node", sub_node) + .set_entry_point("sub_node") + .compile() + ) + + async def main_node(state: CountState) -> None: + await sub_graph.ainvoke({"count": 0}) + + main_graph = ( + StateGraph(CountState) + .add_node("main_node", main_node) + .set_entry_point("main_node") + .compile() + ) + + run = await main_graph.astream_events({"count": 0}, version="v3") + + # A consumer task drives the pump. Once the subgraph node is + # running, no further event is produced, so the consumer parks + # inside _apump_next awaiting graph_aiter.__anext__() — the + # generator is "running" and a plain aclose() would raise. + async def consume() -> None: + async for _e in run: + pass + + consumer = asyncio.create_task(consume()) + try: + await asyncio.wait_for(started.wait(), timeout=2.0) + # Let the consumer drain and park in __anext__. + await asyncio.sleep(0.05) + # Abort from a different task while the consumer is in __anext__. + await run.abort() + # The in-flight subgraph node must observe cancellation. + await asyncio.wait_for(cancelled.wait(), timeout=2.0) + finally: + consumer.cancel() + try: + await consumer + except asyncio.CancelledError: + pass + async def test_extensions_has_native_keys(self) -> None: run = await _build_simple_graph().astream_events( {"value": "x", "items": []}, version="v3"