Refactored so overriden logic lives in register_activity() so async works regardless of how we register it...

Signed-off-by: Patrick Assuied <[email protected]>

Author: passuied

ext/dapr-ext-workflow/dapr/ext/workflow/aio/workflow_runtime.py

@@ -61,6 +61,44 @@ def __init__(
         self._main_event_loop = main_event_loop
         self._logger = Logger('WorkflowRuntime', logger_options)
 
+    def register_activity(self, fn: Activity, *, name: Optional[str] = None):
+        """Registers an async workflow activity and ensures proper execution and metadata.
+        This mirrors the decorator behavior so direct registration behaves identically.
+        """
+        # Validate/prepare alternate name on the original function first (not the wrapper)
+        if hasattr(fn, '_activity_registered'):
+            alt_name = fn.__dict__['_dapr_alternate_name']
+            raise ValueError(f'Activity {fn.__name__} already registered as {alt_name}')
+        if hasattr(fn, '_dapr_alternate_name'):
+            alt_name = fn._dapr_alternate_name
+            if name is not None:
+                raise ValueError(f'Activity {fn.__name__} already has an alternate name {alt_name}')
+        else:
+            alt_name = name if name else fn.__name__
+            fn.__dict__['_dapr_alternate_name'] = alt_name
+
+        # Build the target function that super().register_activity will wrap.
+        if self._main_event_loop:
+
+            @wraps(fn)
+            def target_fn(*args, **kwargs):
+                result = fn(*args, **kwargs)
+                if inspect.isawaitable(result):
+                    future = asyncio.run_coroutine_threadsafe(result, self._main_event_loop)
+                    return future.result()
+                return result
+        else:
+
+            @wraps(fn)
+            def target_fn():
+                return fn
+
+        # Delegate to base registration without passing name because the wrapper already
+        # carries _dapr_alternate_name via @wraps(fn) copying fn.__dict__.
+        super().register_activity(target_fn, name=None)
+        fn.__dict__['_activity_registered'] = True
+        fn.__dict__['_dapr_alternate_name'] = alt_name
+
     def activity(self, __fn: Activity = None, *, name: Optional[str] = None):
         """Decorator to register an async activity function.
 
@@ -88,38 +126,20 @@ async def add(ctx, x: int, y: int) -> int:
         """
 
         def wrapper(fn: Activity):
-            # If a main event loop is provided, wrap the function so that any awaitable
-            # result is executed on that loop in a thread-safe way.
-            if self._main_event_loop:
-
-                @wraps(fn)
-                def sync_wrapper(*args, **kwargs):
-                    result = fn(*args, **kwargs)
-                    if inspect.isawaitable(result):
-                        future = asyncio.run_coroutine_threadsafe(result, self._main_event_loop)
-                        return future.result()
-                    return result
-
-                target_fn = sync_wrapper
-            else:
-                # No special handling needed; register the original function directly.
-                @wraps(fn)
-                def innerfn():
-                    return fn
-
-                target_fn = innerfn
+            # Register first to ensure original fn has flags set
+            self.register_activity(fn, name=name)
 
-            self.register_activity(target_fn, name=name)
+            @wraps(fn)
+            def innerfn():
+                return fn
 
+            # Mirror naming metadata on the returned decorator wrapper
             if hasattr(fn, '_dapr_alternate_name'):
-                target_fn.__dict__['_dapr_alternate_name'] = fn.__dict__['_dapr_alternate_name']
+                innerfn.__dict__['_dapr_alternate_name'] = fn.__dict__['_dapr_alternate_name']
             else:
-                target_fn.__dict__['_dapr_alternate_name'] = name if name else fn.__name__
-            target_fn.__signature__ = inspect.signature(fn)
-            # Copy attributes to fn so it doesn't get registered again when calling `register_activity()` again.
-            fn.__dict__['_activity_registered'] = target_fn.__dict__['_activity_registered']
-            fn.__dict__['_dapr_alternate_name'] = target_fn.__dict__['_dapr_alternate_name']
-            return target_fn
+                innerfn.__dict__['_dapr_alternate_name'] = name if name else fn.__name__
+            innerfn.__signature__ = inspect.signature(fn)
+            return innerfn
 
         if __fn:
             # This case is true when the decorator is used without arguments