diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py
index a3d012e55..335d7b78e 100644
--- a/libs/langgraph/langgraph/func/__init__.py
+++ b/libs/langgraph/langgraph/func/__init__.py
@@ -36,23 +36,31 @@ from langgraph.types import _DC_KWARGS, RetryPolicy, StreamMode
 
 @overload
 def task(
-    *, name: Optional[str] = None, retry: Optional[RetryPolicy] = None
-) -> Callable[[Callable[P, T]], Callable[P, SyncAsyncFuture[T]]]: ...
+    *,
+    name: Optional[str] = None,
+    retry: Optional[RetryPolicy] = None,
+) -> Callable[
+    [Union[Callable[P, Awaitable[T]], Callable[P, T]]],
+    Callable[P, SyncAsyncFuture[T]],
+]: ...
 
 
 @overload
 def task(
-    __func_or_none__: Callable[P, T],
+    __func_or_none__: Union[Callable[P, Awaitable[T]], Callable[P, T]],
 ) -> Callable[P, SyncAsyncFuture[T]]: ...
 
 
 def task(
-    __func_or_none__: Optional[Union[Callable[P, T], Callable[P, Awaitable[T]]]] = None,
+    __func_or_none__: Optional[Union[Callable[P, Awaitable[T]], Callable[P, T]]] = None,
     *,
     name: Optional[str] = None,
     retry: Optional[RetryPolicy] = None,
 ) -> Union[
-    Callable[[Callable[P, T]], Callable[P, SyncAsyncFuture[T]]],
+    Callable[
+        [Union[Callable[P, Awaitable[T]], Callable[P, T]]],
+        Callable[P, SyncAsyncFuture[T]],
+    ],
     Callable[P, SyncAsyncFuture[T]],
 ]:
     """Define a LangGraph task using the `task` decorator.