respond to comments

Signed-off-by: Tim Li <[email protected]>

Author: timl3136

cadence/__init__.py

@@ -7,7 +7,7 @@
 # Import main client functionality
 from .client import Client
 from .worker import Registry
-from .workflow import workflow
+from . import workflow
 
 __version__ = "0.1.0"
 

cadence/_internal/workflow/workflow_engine.py

@@ -279,7 +279,7 @@ async def _execute_workflow_function(self, decision_task: PollForDecisionTaskRes
             workflow_input = await self._extract_workflow_input(decision_task)
 
             # Execute workflow function
-            result = self._execute_workflow_function_once(workflow_func, workflow_input)
+            result = await self._execute_workflow_function_once(workflow_func, workflow_input)
 
             # Check if workflow is complete
             if result is not None:
@@ -341,7 +341,7 @@ async def _extract_workflow_input(self, decision_task: PollForDecisionTaskRespon
         logger.warning("No WorkflowExecutionStarted event found in history")
         return None
 
-    def _execute_workflow_function_once(self, workflow_func: Callable, workflow_input: Any) -> Any:
+    async def _execute_workflow_function_once(self, workflow_func: Callable, workflow_input: Any) -> Any:
         """
         Execute the workflow function once (not during replay).
 
@@ -355,23 +355,9 @@ def _execute_workflow_function_once(self, workflow_func: Callable, workflow_inpu
         logger.debug(f"Executing workflow function with input: {workflow_input}")
         result = workflow_func(workflow_input)
 
-        # If the workflow function is async, we need to handle it properly
+        # If the workflow function is async, await it properly
         if asyncio.iscoroutine(result):
-            # For now, use asyncio.run for async workflow functions
-            # TODO: Implement proper deterministic event loop for workflow execution
-            try:
-                result = asyncio.run(result)
-            except RuntimeError:
-                # If we're already in an event loop, create a new task
-                loop = asyncio.get_event_loop()
-                if loop.is_running():
-                    # We can't use asyncio.run inside a running loop
-                    # For now, just get the result (this may not be deterministic)
-                    logger.warning("Async workflow function called within running event loop - may not be deterministic")
-                    # This is a workaround - in a real implementation, we'd need proper task scheduling
-                    result = None
-                else:
-                    result = loop.run_until_complete(result)
+            result = await result
 
         return result
 

cadence/worker/_registry.py

@@ -7,12 +7,15 @@
 """
 
 import logging
-from typing import Callable, Dict, Optional, Unpack, TypedDict, Sequence, overload, Type, Union
+from typing import Callable, Dict, Optional, Unpack, TypedDict, overload, Type, Union, TypeVar
 from cadence.activity import ActivityDefinitionOptions, ActivityDefinition, ActivityDecorator, P, T
 from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions
 
 logger = logging.getLogger(__name__)
 
+# TypeVar for workflow class types
+W = TypeVar('W')
+
 
 class RegisterWorkflowOptions(TypedDict, total=False):
     """Options for registering a workflow."""
@@ -35,9 +38,9 @@ def __init__(self) -> None:
 
     def workflow(
         self,
-        cls: Optional[Type] = None,
+        cls: Optional[Type[W]] = None,
         **kwargs: Unpack[RegisterWorkflowOptions]
-    ) -> Union[Type, Callable[[Type], Type]]:
+    ) -> Union[Type[W], Callable[[Type[W]], Type[W]]]:
         """
         Register a workflow class.
 
@@ -57,7 +60,7 @@ def workflow(
         """
         options = RegisterWorkflowOptions(**kwargs)
 
-        def decorator(target: Type) -> Type:
+        def decorator(target: Type[W]) -> Type[W]:
             workflow_name = options.get('name') or target.__name__
 
             if workflow_name in self._workflows:
@@ -194,7 +197,7 @@ def of(*args: 'Registry') -> 'Registry':
 
         return result
 
-def _find_activity_definitions(instance: object) -> Sequence[ActivityDefinition]:
+def _find_activity_definitions(instance: object) -> list[ActivityDefinition]:
     attr_to_def = {}
     for t in instance.__class__.__mro__:
         for attr in dir(t):
@@ -206,10 +209,7 @@ def _find_activity_definitions(instance: object) -> Sequence[ActivityDefinition]
                     raise ValueError(f"'{attr}' was overridden with a duplicate activity definition")
                 attr_to_def[attr] = value
 
-    # Create new definitions, copying the attributes from the declaring type but using the function
-    # from the specific object. This allows for the decorator to be applied to the base class and the
-    # function to be overridden
-    result = []
+    result: list[ActivityDefinition] = []
     for attr, definition in attr_to_def.items():
         result.append(ActivityDefinition(getattr(instance, attr), definition.name, definition.strategy, definition.params))
 

cadence/workflow.py

@@ -2,12 +2,12 @@
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
-from typing import Iterator, Callable, TypeVar, TypedDict, Type, cast, Any
-from functools import wraps
+from typing import Iterator, Callable, TypeVar, TypedDict, Type, cast, Any, Optional, Union
+import inspect
 
 from cadence.client import Client
 
-T = TypeVar('T')
+T = TypeVar('T', bound=Callable[..., Any])
 
 
 class WorkflowDefinitionOptions(TypedDict, total=False):
@@ -23,9 +23,10 @@ class WorkflowDefinition:
     Provides type safety and metadata for workflow classes.
     """
 
-    def __init__(self, cls: Type, name: str):
+    def __init__(self, cls: Type, name: str, run_method_name: str):
         self._cls = cls
         self._name = name
+        self._run_method_name = run_method_name
 
     @property
     def name(self) -> str:
@@ -39,13 +40,7 @@ def cls(self) -> Type:
 
     def get_run_method(self, instance: Any) -> Callable:
         """Get the workflow run method from an instance of the workflow class."""
-        for attr_name in dir(instance):
-            if attr_name.startswith('_'):
-                continue
-            attr = getattr(instance, attr_name)
-            if callable(attr) and hasattr(attr, '_workflow_run'):
-                return cast(Callable, attr)
-        raise ValueError(f"No @workflow.run method found in class {self._cls.__name__}")
+        return cast(Callable, getattr(instance, self._run_method_name))
 
     @staticmethod
     def wrap(cls: Type, opts: WorkflowDefinitionOptions) -> 'WorkflowDefinition':
@@ -66,8 +61,8 @@ def wrap(cls: Type, opts: WorkflowDefinitionOptions) -> 'WorkflowDefinition':
         if "name" in opts and opts["name"]:
             name = opts["name"]
 
-        # Validate that the class has exactly one run method
-        run_method_count = 0
+        # Validate that the class has exactly one run method and find it
+        run_method_name = None
         for attr_name in dir(cls):
             if attr_name.startswith('_'):
                 continue
@@ -78,40 +73,54 @@ def wrap(cls: Type, opts: WorkflowDefinitionOptions) -> 'WorkflowDefinition':
 
             # Check for workflow run method
             if hasattr(attr, '_workflow_run'):
-                run_method_count += 1
+                if run_method_name is not None:
+                    raise ValueError(f"Multiple @workflow.run methods found in class {cls.__name__}")
+                run_method_name = attr_name
 
-        if run_method_count == 0:
+        if run_method_name is None:
             raise ValueError(f"No @workflow.run method found in class {cls.__name__}")
-        elif run_method_count > 1:
-            raise ValueError(f"Multiple @workflow.run methods found in class {cls.__name__}")
 
-        return WorkflowDefinition(cls, name)
+        return WorkflowDefinition(cls, name, run_method_name) 
 
 
-def run(func: Callable[..., T]) -> Callable[..., T]:
+def run(func: Optional[T] = None) -> Union[T, Callable[[T], T]]:
     """
     Decorator to mark a method as the main workflow run method.
 
+    Can be used with or without parentheses:
+        @workflow.run
+        async def my_workflow(self):
+            ...
+
+        @workflow.run()
+        async def my_workflow(self):
+            ...
+
     Args:
         func: The method to mark as the workflow run method
 
     Returns:
         The decorated method with workflow run metadata
+    
+    Raises:
+        ValueError: If the function is not async
     """
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        return func(*args, **kwargs)
-
-    # Attach metadata to the function
-    wrapper._workflow_run = True  # type: ignore
-    return wrapper
-
-
-# Create a simple namespace object for the workflow decorators
-class _WorkflowNamespace:
-    run = staticmethod(run)
-
-workflow = _WorkflowNamespace()
+    def decorator(f: T) -> T:
+        # Validate that the function is async
+        if not inspect.iscoroutinefunction(f):
+            raise ValueError(f"Workflow run method '{f.__name__}' must be async")
+        
+        # Attach metadata to the function
+        f._workflow_run = True  # type: ignore
+        return f
+
+    # Support both @workflow.run and @workflow.run()
+    if func is None:
+        # Called with parentheses: @workflow.run()
+        return decorator
+    else:
+        # Called without parentheses: @workflow.run
+        return decorator(func)
 
 
 @dataclass

tests/cadence/_internal/workflow/test_workflow_engine_integration.py

@@ -9,7 +9,8 @@
 from cadence.api.v1.common_pb2 import Payload, WorkflowExecution, WorkflowType
 from cadence.api.v1.history_pb2 import History, HistoryEvent, WorkflowExecutionStartedEventAttributes
 from cadence._internal.workflow.workflow_engine import WorkflowEngine, DecisionResult
-from cadence.workflow import WorkflowInfo
+from cadence import workflow
+from cadence.workflow import WorkflowInfo, WorkflowDefinition, WorkflowDefinitionOptions
 from cadence.client import Client
 
 
@@ -38,11 +39,9 @@ def workflow_info(self):
     @pytest.fixture
     def mock_workflow_definition(self):
         """Create a mock workflow definition."""
-        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions, workflow
-
         class TestWorkflow:
             @workflow.run
-            def weird_name(self, input_data):
+            async def weird_name(self, input_data):
                 return f"processed: {input_data}"
 
         workflow_opts = WorkflowDefinitionOptions(name="test_workflow")
@@ -214,39 +213,42 @@ async def test_extract_workflow_input_deserialization_error(self, workflow_engin
         # Verify no input was extracted due to error
         assert input_data is None
 
-    def test_execute_workflow_function_sync(self, workflow_engine):
+    @pytest.mark.asyncio
+    async def test_execute_workflow_function_sync(self, workflow_engine):
         """Test synchronous workflow function execution."""
         input_data = "test-input"
 
         # Get the workflow function from the instance
         workflow_func = workflow_engine._workflow_definition.get_run_method(workflow_engine._workflow_instance)
 
         # Execute the workflow function
-        result = workflow_engine._execute_workflow_function_once(workflow_func, input_data)
+        result = await workflow_engine._execute_workflow_function_once(workflow_func, input_data)
 
         # Verify the result
         assert result == "processed: test-input"
 
-    def test_execute_workflow_function_async(self, workflow_engine):
+    @pytest.mark.asyncio
+    async def test_execute_workflow_function_async(self, workflow_engine):
         """Test asynchronous workflow function execution."""
         async def async_workflow_func(input_data):
             return f"async-processed: {input_data}"
 
         input_data = "test-input"
 
         # Execute the async workflow function
-        result = workflow_engine._execute_workflow_function_once(async_workflow_func, input_data)
+        result = await workflow_engine._execute_workflow_function_once(async_workflow_func, input_data)
 
         # Verify the result
         assert result == "async-processed: test-input"
 
-    def test_execute_workflow_function_none(self, workflow_engine):
+    @pytest.mark.asyncio
+    async def test_execute_workflow_function_none(self, workflow_engine):
         """Test workflow function execution with None function."""
         input_data = "test-input"
 
         # Execute with None workflow function - should raise TypeError
         with pytest.raises(TypeError, match="'NoneType' object is not callable"):
-            workflow_engine._execute_workflow_function_once(None, input_data)
+            await workflow_engine._execute_workflow_function_once(None, input_data)
 
     def test_workflow_engine_initialization(self, workflow_engine, workflow_info, mock_client, mock_workflow_definition):
         """Test WorkflowEngine initialization."""
@@ -281,11 +283,9 @@ async def test_workflow_engine_workflow_completion(self, workflow_engine, mock_c
         decision_task = self.create_mock_decision_task()
 
         # Create a workflow definition that returns a result (indicating completion)
-        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions, workflow
-
         class CompletingWorkflow:
             @workflow.run
-            def run(self, input_data):
+            async def run(self, input_data):
                 return "workflow-completed"
 
         workflow_opts = WorkflowDefinitionOptions(name="completing_workflow")

tests/cadence/worker/test_decision_task_handler.py

@@ -17,6 +17,8 @@
 from cadence.worker._decision_task_handler import DecisionTaskHandler
 from cadence.worker._registry import Registry
 from cadence._internal.workflow.workflow_engine import WorkflowEngine, DecisionResult
+from cadence import workflow
+from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions
 
 
 class TestDecisionTaskHandler:
@@ -83,8 +85,6 @@ def test_initialization(self, mock_client, mock_registry):
     async def test_handle_task_implementation_success(self, handler, sample_decision_task, mock_registry):
         """Test successful decision task handling."""
         # Create actual workflow definition
-        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions, workflow
-
         class MockWorkflow:
             @workflow.run
             async def run(self):
@@ -151,8 +151,6 @@ async def test_handle_task_implementation_workflow_not_found(self, handler, samp
     async def test_handle_task_implementation_caches_engines(self, handler, sample_decision_task, mock_registry):
         """Test that decision task handler caches workflow engines for same workflow execution."""
         # Create actual workflow definition
-        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions, workflow
-
         class MockWorkflow:
             @workflow.run
             async def run(self):
@@ -189,8 +187,6 @@ async def run(self):
     async def test_handle_task_implementation_different_executions_get_separate_engines(self, handler, mock_registry):
         """Test that different workflow executions get separate engines."""
         # Create actual workflow definition
-        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions, workflow
-
         class MockWorkflow:
             @workflow.run
             async def run(self):
@@ -348,8 +344,6 @@ async def test_respond_decision_task_completed_error(self, handler, sample_decis
     async def test_workflow_engine_creation_with_workflow_info(self, handler, sample_decision_task, mock_registry):
         """Test that WorkflowEngine is created with correct WorkflowInfo."""
         # Create actual workflow definition
-        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions, workflow
-
         class MockWorkflow:
             @workflow.run
             async def run(self):

tests/cadence/worker/test_decision_task_handler_integration.py

@@ -13,7 +13,7 @@
 from cadence.api.v1.decision_pb2 import Decision
 from cadence.worker._decision_task_handler import DecisionTaskHandler
 from cadence.worker._registry import Registry
-from cadence.workflow import workflow
+from cadence import workflow
 from cadence.client import Client
 
 

tests/cadence/worker/test_decision_worker_integration.py

@@ -11,7 +11,7 @@
 from cadence.api.v1.history_pb2 import History, HistoryEvent, WorkflowExecutionStartedEventAttributes
 from cadence.worker._decision import DecisionWorker
 from cadence.worker._registry import Registry
-from cadence.workflow import workflow
+from cadence import workflow
 from cadence.client import Client
 
 

tests/cadence/worker/test_registry.py

@@ -6,8 +6,9 @@
 import pytest
 
 from cadence import activity
+from cadence import workflow
 from cadence.worker import Registry
-from cadence.workflow import workflow, WorkflowDefinition
+from cadence.workflow import WorkflowDefinition
 from tests.cadence import common_activities