Repository: PrefectHQ/prefect

Issue:
Wrong type hints in `@flow`'s on_completion and on_failure
### First check

- [X] I added a descriptive title to this issue.
- [X] I used the GitHub search to find a similar issue and didn't find it.
- [X] I searched the Prefect documentation for this issue.
- [X] I checked that this issue is related to Prefect and not one of its dependencies.

### Bug summary

Current type hints of `@flow`'s on_completion and on_failure is

https://github.com/PrefectHQ/prefect/blob/c32bf18a8cbf226ff68e733b5f48123a7c3ad2f1/src/prefect/tasks.py#L385-L386

These type hints won't accept async functions. As reported by pylance/pyright:
> "Coroutine[Any, Any, None]" is incompatible with "None"

---

Async functions indeed supported since the initial PR. #8605

https://github.com/PrefectHQ/prefect/blob/c32bf18a8cbf226ff68e733b5f48123a7c3ad2f1/src/prefect/engine.py#L2758-L2778

### Reproduction

```python3
# run `pyright` on this file

from prefect import flow

async def f(x, y, z):
    pass

@flow(on_completion=([f]))
def g():
    pass
```


### Error

```python3
error: Argument of type "list[(x: Unknown, y: Unknown, z: Unknown) -> Coroutine[Any, Any, None]]" cannot be assigned to parameter "on_completion" of type "List[(Flow, FlowRun, State[Unknown]) -> None] | None" in function "flow"
    Type "(x: Unknown, y: Unknown, z: Unknown) -> Coroutine[Any, Any, None]" cannot be assigned to type "(Flow, FlowRun, State[Unknown]) -> None"
      Function return type "Coroutine[Any, Any, None]" is incompatible with type "None"
        "Coroutine[Any, Any, None]" is incompatible with "None" (reportArgumentType)
```


### Versions

```Text
Git commit:          d147b180
Built:               Thu, Sep 14, 2023 2:54 PM
OS/Arch:             linux/x86_64
Profile:             default
Server type:         ephemeral
Server:
  Database:          sqlite
  SQLite version:    3.39.3
```


### Additional context

_No response_


Discussion:
Hi, can I be assigned on this one? Thanks in advanced!

Resolving patch:
diff --git a/src/prefect/flows.py b/src/prefect/flows.py
--- a/src/prefect/flows.py
+++ b/src/prefect/flows.py
@@ -1345,8 +1345,12 @@ def flow(
     result_serializer: Optional[ResultSerializer] = None,
     cache_result_in_memory: bool = True,
     log_prints: Optional[bool] = None,
-    on_completion: Optional[List[Callable[[FlowSchema, FlowRun, State], None]]] = None,
-    on_failure: Optional[List[Callable[[FlowSchema, FlowRun, State], None]]] = None,
+    on_completion: Optional[
+        List[Callable[[FlowSchema, FlowRun, State], Union[Awaitable[None], None]]]
+    ] = None,
+    on_failure: Optional[
+        List[Callable[[FlowSchema, FlowRun, State], Union[Awaitable[None], None]]]
+    ] = None,
     on_cancellation: Optional[
         List[Callable[[FlowSchema, FlowRun, State], None]]
     ] = None,
@@ -1373,8 +1377,12 @@ def flow(
     result_serializer: Optional[ResultSerializer] = None,
     cache_result_in_memory: bool = True,
     log_prints: Optional[bool] = None,
-    on_completion: Optional[List[Callable[[FlowSchema, FlowRun, State], None]]] = None,
-    on_failure: Optional[List[Callable[[FlowSchema, FlowRun, State], None]]] = None,
+    on_completion: Optional[
+        List[Callable[[FlowSchema, FlowRun, State], Union[Awaitable[None], None]]]
+    ] = None,
+    on_failure: Optional[
+        List[Callable[[FlowSchema, FlowRun, State], Union[Awaitable[None], None]]]
+    ] = None,
     on_cancellation: Optional[
         List[Callable[[FlowSchema, FlowRun, State], None]]
     ] = None,
diff --git a/src/prefect/tasks.py b/src/prefect/tasks.py
--- a/src/prefect/tasks.py
+++ b/src/prefect/tasks.py
@@ -391,8 +391,12 @@ def with_options(
         timeout_seconds: Union[int, float] = None,
         log_prints: Optional[bool] = NotSet,
         refresh_cache: Optional[bool] = NotSet,
-        on_completion: Optional[List[Callable[["Task", TaskRun, State], None]]] = None,
-        on_failure: Optional[List[Callable[["Task", TaskRun, State], None]]] = None,
+        on_completion: Optional[
+            List[Callable[["Task", TaskRun, State], Union[Awaitable[None], None]]]
+        ] = None,
+        on_failure: Optional[
+            List[Callable[["Task", TaskRun, State], Union[Awaitable[None], None]]]
+        ] = None,
         retry_condition_fn: Optional[Callable[["Task", TaskRun, State], bool]] = None,
         viz_return_value: Optional[Any] = None,
     ):


Test patch: