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=<AsyncQueue.wait() ...>
  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 <caspar@langchain.dev>

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,

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)