From 6139dacef96a4dacb4290892a628df674c0c62cc Mon Sep 17 00:00:00 2001 From: dmsthd Date: Wed, 24 Sep 2025 09:11:17 +0900 Subject: [PATCH] fix(langgraph): cleanup orphaned waiter task in AsyncPregelLoop (#6167) ### Summary This PR fixes an issue where `AsyncPregelLoop` could leave behind an orphaned `stream.wait()` task, resulting in warnings like: ``` Task was destroyed but it is pending! ``` ### Related Discussion This PR is in response to: [langchain-ai/langgraph#6163](https://github.com/langchain-ai/langgraph/discussions/6163) ### Problem * In the async path, `get_waiter()` was creating a new `asyncio.Task` via ```python aioloop.create_task(stream.wait()) ``` but never tracked or cleaned it up. * On cancellation or shutdown, these tasks remained pending and produced warnings. ### Solution * Changed `get_waiter()` to: * Maintain a **single waiter task** (similar to the sync path). * Auto-clear the reference when the task finishes. * Added `_cleanup_waiter()`: * On exit, attempt to wake the waiter (`stream._count.release()` if available). * Otherwise, cancel and `await` the pending task to ensure proper cleanup. * Wrapped the `while loop.tick():` block in a `try/finally` to guarantee `_cleanup_waiter()` runs on exit. * Added missing `import contextlib`. ### Impact * Prevents orphaned `stream.wait()` tasks. * Removes noisy `"Task was destroyed but it is pending!"` warnings. * Behavior of async streaming remains unchanged, only lifecycle management improved. ### Test Plan * Reproduced the issue by running async streaming with cancellation. * Verified warnings no longer appear after the fix. * Ran existing test suite (all passing). ### Notes * Sync and Async implementations now follow the same principle: *only one waiter at a time, always cleaned up on exit*. * Backwards-compatible; no API changes. ### Repro & Verification To confirm the issue and the fix I used the following minimal repro snippet: ```python # lg_repro.py import asyncio import os # Enable asyncio debug logs to surface pending task warnings os.environ.setdefault("PYTHONASYNCIODEBUG", "1") from langgraph.graph import START, END, StateGraph State = dict # Slow async node: processes once, then sleeps to keep the waiter alive async def slow_node(state: State) -> State: await asyncio.sleep(0.2) # simulate work state["count"] = state.get("count", 0) + 1 await asyncio.sleep(1.0) # keep stream.wait() waiter active return state # Build simple graph: START -> slow_node -> END builder = StateGraph(State) builder.add_node("slow", slow_node) builder.add_edge(START, "slow") builder.add_edge("slow", END) graph = builder.compile() async def run_and_cancel(): # astream with messages mode triggers internal stream.wait() waiter async def consumer(): async for _ in graph.astream({"msg": "hi"}, stream_mode="messages"): await asyncio.sleep(0.05) t = asyncio.create_task(consumer(), name="astream-consumer") # Allow the stream to start, then cancel the consumer await asyncio.sleep(0.1) t.cancel() try: await t except asyncio.CancelledError: pass # Let loop settle to show pending waiter task if not cleaned await asyncio.sleep(0.05) def main(): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) loop.set_debug(True) try: loop.run_until_complete(run_and_cancel()) finally: # If the internal waiter is not cleaned, closing the loop will warn loop.close() if __name__ == "__main__": main() ```` **How to run** ```powershell # Before (main branch) git checkout main pip install -e libs/langgraph $env:PYTHONASYNCIODEBUG=1; python lg_repro.py # After (patched branch) git checkout async-waiter-cleanup pip install -e libs/langgraph $env:PYTHONASYNCIODEBUG=1; python lg_repro.py ``` **Observed results** * **main branch (before fix):** Shows warnings like: ``` Task was destroyed but it is pending! ... coro= created at langgraph/pregel/main.py:2927 ``` * **patched branch (after fix):** No warnings. The single waiter is properly cleaned up on exit via `_cleanup_waiter()` (release semaphore if available, then cancel/await). --- This confirms that the patch removes the orphaned `stream.wait()` task and prevents `"Task was destroyed but it is pending!"` warnings during cancellation/shutdown. --------- Co-authored-by: Caspar Broekhuizen --- libs/langgraph/langgraph/pregel/main.py | 96 ++++++++++++++++------- libs/langgraph/tests/test_pregel_async.py | 55 +++++++++++++ 2 files changed, 123 insertions(+), 28 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index a47ff3f4b..2e3fdd926 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -3,15 +3,24 @@ from __future__ import annotations import asyncio import concurrent import concurrent.futures +import contextlib import queue import warnings import weakref from collections import defaultdict, deque -from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Iterator, Mapping, Sequence from dataclasses import is_dataclass from functools import partial from inspect import isclass -from typing import Any, Callable, Generic, Optional, Union, cast, get_type_hints +from typing import ( + Any, + Callable, + Generic, + Optional, + Union, + cast, + get_type_hints, +) from uuid import UUID, uuid5 from langchain_core.globals import get_debug @@ -2612,6 +2621,7 @@ class Pregel( if subgraphs: loop.config[CONF][CONFIG_KEY_STREAM] = loop.stream # enable concurrent streaming + get_waiter: Callable[[], concurrent.futures.Future[None]] | None = None if ( self.stream_eager or subgraphs @@ -2634,8 +2644,6 @@ class Pregel( else: return waiter - else: - get_waiter = None # type: ignore[assignment] # Similarly to Bulk Synchronous Parallel / Pregel model # computation proceeds in steps, while there are channel updates. # Channel updates from step N are only visible in step N+1 @@ -2916,45 +2924,77 @@ class Pregel( stream_put, stream_modes ) # enable concurrent streaming + get_waiter: Callable[[], asyncio.Task[None]] | None = None + _cleanup_waiter: Callable[[], Awaitable[None]] | None = None if ( self.stream_eager or subgraphs or "messages" in stream_modes or "custom" in stream_modes ): + # Keep a single waiter task alive; ensure cleanup on exit. + waiter: asyncio.Task[None] | None = None def get_waiter() -> asyncio.Task[None]: - return aioloop.create_task(stream.wait()) + nonlocal waiter + if waiter is None or waiter.done(): + waiter = aioloop.create_task(stream.wait()) + + def _clear(t: asyncio.Task[None]) -> None: + nonlocal waiter + if waiter is t: + waiter = None + + waiter.add_done_callback(_clear) + return waiter + + async def _cleanup_waiter() -> None: + """Wake pending waiter and/or cancel+await to avoid pending tasks.""" + nonlocal waiter + # Try to wake via semaphore like SyncPregelLoop + with contextlib.suppress(Exception): + if hasattr(stream, "_count"): + stream._count.release() + t = waiter + waiter = None + if t is not None and not t.done(): + t.cancel() + with contextlib.suppress(asyncio.CancelledError): + await t - else: - get_waiter = None # type: ignore[assignment] # Similarly to Bulk Synchronous Parallel / Pregel model # computation proceeds in steps, while there are channel updates # channel updates from step N are only visible in step N+1 # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps - while loop.tick(): - for task in await loop.amatch_cached_writes(): - loop.output_writes(task.id, task.writes, cached=True) - async for _ in runner.atick( - [t for t in loop.tasks.values() if not t.writes], - timeout=self.step_timeout, - get_waiter=get_waiter, - schedule_task=loop.aaccept_push, - ): - # emit output - for o in _output( - stream_mode, - print_mode, - subgraphs, - stream.get_nowait, - asyncio.QueueEmpty, + try: + while loop.tick(): + for task in await loop.amatch_cached_writes(): + loop.output_writes(task.id, task.writes, cached=True) + async for _ in runner.atick( + [t for t in loop.tasks.values() if not t.writes], + timeout=self.step_timeout, + get_waiter=get_waiter, + schedule_task=loop.aaccept_push, ): - yield o - loop.after_tick() - # wait for checkpoint - if durability_ == "sync": - await cast(asyncio.Future, loop._put_checkpoint_fut) + # emit output + for o in _output( + stream_mode, + print_mode, + subgraphs, + stream.get_nowait, + asyncio.QueueEmpty, + ): + yield o + loop.after_tick() + # wait for checkpoint + if durability_ == "sync": + await cast(asyncio.Future, loop._put_checkpoint_fut) + finally: + # ensure waiter doesn't remain pending on cancel/shutdown + if _cleanup_waiter is not None: + await _cleanup_waiter() + # emit output for o in _output( stream_mode, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index a1d552e6a..a6296a0c4 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -41,6 +41,7 @@ from syrupy import SnapshotAssertion from typing_extensions import NotRequired, TypedDict from langgraph._internal._constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL +from langgraph._internal._queue import AsyncQueue from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic @@ -9156,3 +9157,57 @@ async def test_null_resume_disallowed_with_multiple_interrupts( "text_1": "resume for prompt: original text 1", "text_2": "resume for prompt: original text 2", } + + +async def test_astream_waiter_cleanup_on_cancel( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that AsyncPregelLoop cleans up waiter tasks after cancellation.""" + + recorded_tasks: list[asyncio.Task[None]] = [] + finished_tasks: list[asyncio.Task[None]] = [] + + original_wait = AsyncQueue.wait + + async def tracked_wait(self: AsyncQueue) -> None: + task = asyncio.current_task() + assert task is not None + recorded_tasks.append(task) + try: + await original_wait(self) + finally: + finished_tasks.append(task) + + monkeypatch.setattr(AsyncQueue, "wait", tracked_wait) + + class State(TypedDict, total=False): + count: int + + async def slow_node(state: State) -> State: + await asyncio.sleep(0.05) + state = dict(state) + state["count"] = state.get("count", 0) + 1 + await asyncio.sleep(0.1) + return state + + builder = StateGraph(State) + builder.add_node("slow", slow_node) + builder.add_edge(START, "slow") + builder.add_edge("slow", END) + graph = builder.compile() + + async def consumer() -> None: + async for _ in graph.astream({"msg": "hi"}, stream_mode="messages"): + await asyncio.sleep(0.01) + + task = asyncio.create_task(consumer()) + await asyncio.sleep(0.05) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + await asyncio.sleep(0.05) + + assert recorded_tasks, "expected stream.wait() task to be created" + assert set(finished_tasks) == set(recorded_tasks) + assert all(t.done() for t in recorded_tasks)