From 12be3fac33a95e71ece4b68b8a4816929c3a3da5 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Mon, 20 Jan 2025 11:13:37 -0800
Subject: [PATCH] Lint

---
 libs/langgraph/langgraph/pregel/runner.py |  4 +++-
 libs/langgraph/langgraph/utils/future.py  | 14 ++++++++------
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py
index 02e97a8e0..af80ed715 100644
--- a/libs/langgraph/langgraph/pregel/runner.py
+++ b/libs/langgraph/langgraph/pregel/runner.py
@@ -415,7 +415,9 @@ class PregelRunner:
             # otherwise return a chained sync future
             if in_async:
                 if isinstance(fut, asyncio.Task):
-                    sfut = asyncio.Future(loop=loop)
+                    sfut: Union[asyncio.Future[Any], concurrent.futures.Future[Any]] = (
+                        asyncio.Future(loop=loop)
+                    )
                     loop.call_soon_threadsafe(chain_future, fut, sfut)
                     return sfut
                 else:
diff --git a/libs/langgraph/langgraph/utils/future.py b/libs/langgraph/langgraph/utils/future.py
index 50fc1d0eb..8445f5396 100644
--- a/libs/langgraph/langgraph/utils/future.py
+++ b/libs/langgraph/langgraph/utils/future.py
@@ -4,7 +4,7 @@ import contextvars
 import inspect
 import sys
 import types
-from typing import Coroutine, Optional, TypeVar, Union
+from typing import Awaitable, Coroutine, Generator, Optional, TypeVar, Union, cast
 
 T = TypeVar("T")
 AnyFuture = Union[asyncio.Future, concurrent.futures.Future]
@@ -137,7 +137,7 @@ def chain_future(source: AnyFuture, destination: AnyFuture) -> AnyFuture:
 
 
 def _ensure_future(
-    coro_or_future: Coroutine[None, None, T],
+    coro_or_future: Union[Coroutine[None, None, T], Awaitable[T]],
     *,
     loop: asyncio.AbstractEventLoop,
     name: Optional[str] = None,
@@ -146,7 +146,9 @@ def _ensure_future(
     called_wrap_awaitable = False
     if not asyncio.iscoroutine(coro_or_future):
         if inspect.isawaitable(coro_or_future):
-            coro_or_future = _wrap_awaitable(coro_or_future)
+            coro_or_future = cast(
+                Coroutine[None, None, T], _wrap_awaitable(coro_or_future)
+            )
             called_wrap_awaitable = True
         else:
             raise TypeError(
@@ -165,7 +167,7 @@ def _ensure_future(
 
 
 @types.coroutine
-def _wrap_awaitable(awaitable):
+def _wrap_awaitable(awaitable: Awaitable[T]) -> Generator[None, None, T]:
     """Helper for asyncio.ensure_future().
 
     Wraps awaitable (an object with __await__) into a coroutine
@@ -184,9 +186,9 @@ def run_coroutine_threadsafe(
 
     Return a asyncio.Future to access the result.
     """
-    future = asyncio.Future(loop=loop)
+    future: asyncio.Future[T] = asyncio.Future(loop=loop)
 
-    def callback():
+    def callback() -> None:
         try:
             chain_future(
                 _ensure_future(coro, loop=loop, name=name, context=context), future