fix: forgot to await system task wrapper

Author: fubuloubu

silverback/main.py

@@ -2,7 +2,7 @@
 import inspect
 from collections import defaultdict
 from datetime import timedelta
-from functools import update_wrapper, wraps
+from functools import wraps
 from typing import Any, Awaitable, Callable
 
 from ape.api.networks import LOCAL_NETWORK_NAME
@@ -180,11 +180,21 @@ def __register_system_task(
         self, task_type: TaskType, task_handler: Callable
     ) -> AsyncTaskiqDecoratedTask:
         assert str(task_type).startswith("system:"), "Can only add system tasks"
+
+        # NOTE: We need this as `.register_task` tries to update `.__name__` of `task_handler`,
+        #       but methods do not allow setting this attribute (raises AttributeError)
+        @wraps(task_handler)
+        async def call_task_handler(*args, **kwargs):
+            result = task_handler(*args, **kwargs)
+
+            if inspect.isawaitable(result):
+                return await result
+
+            return result
+
         # NOTE: This has to be registered with the broker in the worker
         return self.broker.register_task(
-            # NOTE: We need this as `.register_task` tries to update `.__name__` of `task_handler`,
-            #       but it is a method not a function (`update_wrapper` transforms it into one)
-            update_wrapper(lambda *args, **kwargs: task_handler(*args, **kwargs), task_handler),
+            call_task_handler,
             # NOTE: Name makes it impossible to conflict with user's handler fn names
             task_name=str(task_type),
             task_type=str(task_type),