Improve typings for task, it now returns a mixed sync/async future protocol (#3186)

- note this type is never instantiated, it is only used for typing (we
cannot make it a protocol as it inherits from concurrent.futures.Future)

libs/langgraph/langgraph/func/__init__.py

@@ -22,7 +22,13 @@ from langgraph.channels.last_value import LastValue
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.constants import END, PREVIOUS, START, TAG_HIDDEN
 from langgraph.pregel import Pregel
-from langgraph.pregel.call import P, T, call, get_runnable_for_entrypoint
+from langgraph.pregel.call import (
+    P,
+    SyncAsyncFuture,
+    T,
+    call,
+    get_runnable_for_entrypoint,
+)
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.store.base import BaseStore
@@ -32,25 +38,13 @@ from langgraph.types import _DC_KWARGS, RetryPolicy, StreamMode, StreamWriter
 @overload
 def task(
     *, retry: Optional[RetryPolicy] = None
-) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, asyncio.Future[T]]]: ...
-
-
-@overload
-def task(  # type: ignore[overload-cannot-match]
-    *, retry: Optional[RetryPolicy] = None
-) -> Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]]: ...
+) -> Callable[[Callable[P, T]], Callable[P, SyncAsyncFuture[T]]]: ...
 
 
 @overload
 def task(
     __func_or_none__: Callable[P, T],
-) -> Callable[P, concurrent.futures.Future[T]]: ...
-
-
-@overload
-def task(
-    __func_or_none__: Callable[P, Awaitable[T]],
-) -> Callable[P, asyncio.Future[T]]: ...
+) -> Callable[P, SyncAsyncFuture[T]]: ...
 
 
 def task(
@@ -58,10 +52,8 @@ def task(
     *,
     retry: Optional[RetryPolicy] = None,
 ) -> Union[
-    Callable[[Callable[P, Awaitable[T]]], Callable[P, asyncio.Future[T]]],
-    Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]],
-    Callable[P, asyncio.Future[T]],
-    Callable[P, concurrent.futures.Future[T]],
+    Callable[[Callable[P, T]], Callable[P, SyncAsyncFuture[T]]],
+    Callable[P, SyncAsyncFuture[T]],
 ]:
     """Define a LangGraph task using the `task` decorator.
 

libs/langgraph/langgraph/pregel/call.py

@@ -1,12 +1,11 @@
 """Utility to convert a user provided function into a Runnable with a ChannelWrite."""
 
-import asyncio
 import concurrent.futures
 import functools
 import inspect
 import sys
 import types
-from typing import Any, Callable, Optional, TypeVar, Union
+from typing import Any, Callable, Generator, Generic, Optional, TypeVar, cast
 
 from langchain_core.runnables import Runnable
 from typing_extensions import ParamSpec
@@ -208,12 +207,17 @@ P1 = TypeVar("P1")
 T = TypeVar("T")
 
 
+class SyncAsyncFuture(Generic[T], concurrent.futures.Future[T]):
+    def __await__(self) -> Generator[T, None, T]:
+        yield cast(T, ...)
+
+
 def call(
     func: Callable[P, T],
     *args: Any,
     retry: Optional[RetryPolicy] = None,
     **kwargs: Any,
-) -> Union[concurrent.futures.Future[T], asyncio.Future[T]]:
+) -> SyncAsyncFuture[T]:
     config = get_config()
     impl = config[CONF][CONFIG_KEY_CALL]
     fut = impl(func, (args, kwargs), retry=retry, callbacks=config["callbacks"])