Add more comments

libs/langgraph/langgraph/pregel/__init__.py

@@ -1417,7 +1417,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
                         loop.tasks.values(),
                         timeout=self.step_timeout,
                         retry_policy=self.retry_policy,
-                        extra=lambda: aioloop.create_task(stream.wait()),
+                        get_waiter=lambda: aioloop.create_task(stream.wait()),
                     ):
                         # emit output
                         for o in output():

libs/langgraph/langgraph/pregel/runner.py

@@ -100,15 +100,15 @@ class PregelRunner:
         reraise: bool = True,
         timeout: Optional[float] = None,
         retry_policy: Optional[RetryPolicy] = None,
-        extra: Optional[Callable[[], asyncio.Future[None]]] = None,
+        get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None,
     ) -> AsyncIterator[None]:
         loop = asyncio.get_event_loop()
         # give control back to the caller
         yield
-        # add extra task if requested
-        if extra is not None:
+        # add waiter task if requested
+        if get_waiter is not None:
             futures: dict[asyncio.Future, Optional[PregelExecutableTask]] = {
-                extra(): None
+                get_waiter(): None
             }
         else:
             futures = {}
@@ -130,7 +130,7 @@ class PregelRunner:
                 ] = task
         all_futures = futures.copy()
         end_time = timeout + loop.time() if timeout else None
-        while len(futures) > (1 if extra is not None else 0):
+        while len(futures) > (1 if get_waiter is not None else 0):
             done, _ = await asyncio.wait(
                 futures,
                 return_when=asyncio.FIRST_COMPLETED,
@@ -141,8 +141,8 @@ class PregelRunner:
             for fut in done:
                 task = futures.pop(fut)
                 if task is None:
-                    # extra task finished, schedule another
-                    futures[extra()] = None
+                    # waiter task finished, schedule another
+                    futures[get_waiter()] = None
                     continue
                 if exc := _exception(fut):
                     if isinstance(exc, GraphInterrupt):
@@ -168,7 +168,7 @@ class PregelRunner:
                 break
             # give control back to the caller
             yield
-        # cancel extra task
+        # cancel waiter task
         for fut in futures:
             fut.cancel()
         # panic on failure or timeout

libs/langgraph/langgraph/utils/aio.py

@@ -6,7 +6,11 @@ PY_310 = sys.version_info >= (3, 10)
 
 class Queue(asyncio.Queue):
     async def wait(self):
-        """If queue is empty, wait until an item is available."""
+        """If queue is empty, wait until an item is available.
+
+        Copied from Queue.get(), removing the call to .get_nowait(),
+        ie. this doesn't consume the item, just waits for it.
+        """
         while self.empty():
             if PY_310:
                 getter = self._get_loop().create_future()