Add WorkflowDefinition

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

Author: timl3136

cadence/worker/_decision_task_handler.py

@@ -76,7 +76,7 @@ async def _handle_task_implementation(self, task: PollForDecisionTaskResponse) -
         )
 
         try:
-            workflow_func = self._registry.get_workflow(workflow_type_name)
+            workflow_definition = self._registry.get_workflow(workflow_type_name)
         except KeyError:
             logger.error(
                 "Workflow type not found in registry",
@@ -105,7 +105,7 @@ async def _handle_task_implementation(self, task: PollForDecisionTaskResponse) -
                 workflow_engine = WorkflowEngine(
                     info=workflow_info, 
                     client=self._client, 
-                    workflow_func=workflow_func
+                    workflow_func=workflow_definition.fn
                 )
                 self._workflow_engines[cache_key] = workflow_engine
 

cadence/worker/_registry.py

@@ -9,6 +9,7 @@
 import logging
 from typing import Callable, Dict, Optional, Unpack, TypedDict, Sequence, overload
 from cadence.activity import ActivityDefinitionOptions, ActivityDefinition, ActivityDecorator, P, T
+from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions
 
 logger = logging.getLogger(__name__)
 
@@ -28,7 +29,7 @@ class Registry:
 
     def __init__(self) -> None:
         """Initialize the registry."""
-        self._workflows: Dict[str, Callable] = {}
+        self._workflows: Dict[str, WorkflowDefinition] = {}
         self._activities: Dict[str, ActivityDefinition] = {}
         self._workflow_aliases: Dict[str, str] = {}  # alias -> name mapping
 
@@ -60,7 +61,10 @@ def decorator(f: Callable) -> Callable:
             if workflow_name in self._workflows:
                 raise KeyError(f"Workflow '{workflow_name}' is already registered")
 
-            self._workflows[workflow_name] = f
+            # Create WorkflowDefinition with type information
+            workflow_opts = WorkflowDefinitionOptions(name=workflow_name)
+            workflow_def = WorkflowDefinition.wrap(f, workflow_opts)
+            self._workflows[workflow_name] = workflow_def
 
             # Register alias if provided
             alias = options.get('alias')
@@ -135,15 +139,15 @@ def _register_activity(self, defn: ActivityDefinition) -> None:
         self._activities[defn.name] = defn
 
 
-    def get_workflow(self, name: str) -> Callable:
+    def get_workflow(self, name: str) -> WorkflowDefinition:
         """
         Get a registered workflow by name.
         
         Args:
             name: Name or alias of the workflow
             
         Returns:
-            The workflow function
+            The workflow definition with type information
             
         Raises:
             KeyError: If workflow is not found

cadence/workflow.py

@@ -2,10 +2,147 @@
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
-from typing import Iterator
+from functools import update_wrapper
+from inspect import signature, Parameter
+from typing import Iterator, Callable, TypeVar, ParamSpec, Generic, TypedDict, Unpack, overload, get_type_hints, Type, Any
 
 from cadence.client import Client
 
+
+@dataclass(frozen=True)
+class WorkflowParameter:
+    """Parameter information for a workflow function."""
+    name: str
+    type_hint: Type | None
+    default_value: Any | None
+
+
+class WorkflowDefinitionOptions(TypedDict, total=False):
+    """Options for defining a workflow."""
+    name: str
+
+
+P = ParamSpec('P')
+T = TypeVar('T')
+
+
+class WorkflowDefinition(Generic[P, T]):
+    """
+    Definition of a workflow function with metadata.
+
+    Similar to ActivityDefinition but for workflows.
+    Provides type safety and metadata for workflow functions.
+    """
+
+    def __init__(self, wrapped: Callable[P, T], name: str, params: list[WorkflowParameter]):
+        self._wrapped = wrapped
+        self._name = name
+        self._params = params
+        update_wrapper(self, wrapped)
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
+        return self._wrapped(*args, **kwargs)
+
+    @property
+    def name(self) -> str:
+        """Get the workflow name."""
+        return self._name
+
+    @property
+    def params(self) -> list[WorkflowParameter]:
+        """Get the workflow parameters."""
+        return self._params
+
+    @property
+    def fn(self) -> Callable[P, T]:
+        """Get the underlying workflow function."""
+        return self._wrapped
+
+    @classmethod
+    def wrap(cls, fn: Callable[P, T], opts: WorkflowDefinitionOptions) -> 'WorkflowDefinition[P, T]':
+        """
+        Wrap a function as a WorkflowDefinition.
+
+        Args:
+            fn: The workflow function to wrap
+            opts: Options for the workflow definition
+
+        Returns:
+            A WorkflowDefinition instance
+        """
+        name = fn.__qualname__
+        if "name" in opts and opts["name"]:
+            name = opts["name"]
+
+        params = _get_workflow_params(fn)
+        return cls(fn, name, params)
+
+
+WorkflowDecorator = Callable[[Callable[P, T]], WorkflowDefinition[P, T]]
+
+
+@overload
+def defn(fn: Callable[P, T]) -> WorkflowDefinition[P, T]:
+    ...
+
+
+@overload
+def defn(**kwargs: Unpack[WorkflowDefinitionOptions]) -> WorkflowDecorator:
+    ...
+
+
+def defn(fn: Callable[P, T] | None = None, **kwargs: Unpack[WorkflowDefinitionOptions]) -> WorkflowDecorator | WorkflowDefinition[P, T]:
+    """
+    Decorator to define a workflow function.
+
+    Usage:
+        @defn
+        def my_workflow(input_data: str) -> str:
+            return f"processed: {input_data}"
+
+        @defn(name="custom_workflow_name")
+        def my_other_workflow(input_data: str) -> str:
+            return f"custom: {input_data}"
+
+    Args:
+        fn: The workflow function (when used without parentheses)
+        **kwargs: Workflow definition options
+
+    Returns:
+        Either a WorkflowDefinition (direct decoration) or a decorator function
+    """
+    opts = WorkflowDefinitionOptions(**kwargs)
+
+    def decorator(inner_fn: Callable[P, T]) -> WorkflowDefinition[P, T]:
+        return WorkflowDefinition.wrap(inner_fn, opts)
+
+    if fn is not None:
+        return decorator(fn)
+
+    return decorator
+
+
+def _get_workflow_params(fn: Callable) -> list[WorkflowParameter]:
+    """Extract parameter information from a workflow function."""
+    args = signature(fn).parameters
+    hints = get_type_hints(fn)
+    result = []
+    for name, param in args.items():
+        # Filter out self parameter
+        if param.name == "self":
+            continue
+        default = None
+        if param.default != Parameter.empty:
+            default = param.default
+        if param.kind in (Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD):
+            type_hint = hints.get(name, None)
+            result.append(WorkflowParameter(name, type_hint, default))
+        else:
+            raise ValueError(f"Parameters must be positional. {name} is {param.kind}, and not valid")
+
+    return result
+
+
 @dataclass
 class WorkflowInfo:
     workflow_type: str

tests/cadence/worker/test_decision_task_handler.py

@@ -82,9 +82,14 @@ def test_initialization(self, mock_client, mock_registry):
     @pytest.mark.asyncio
     async def test_handle_task_implementation_success(self, handler, sample_decision_task, mock_registry):
         """Test successful decision task handling."""
-        # Mock workflow function
-        mock_workflow_func = Mock()
-        mock_registry.get_workflow.return_value = mock_workflow_func
+        # Create actual workflow definition
+        def mock_workflow_func():
+            return "test_result"
+
+        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions
+        workflow_opts = WorkflowDefinitionOptions(name="test_workflow")
+        workflow_definition = WorkflowDefinition.wrap(mock_workflow_func, workflow_opts)
+        mock_registry.get_workflow.return_value = workflow_definition
 
         # Mock workflow engine
         mock_engine = Mock(spec=WorkflowEngine)
@@ -142,9 +147,14 @@ async def test_handle_task_implementation_workflow_not_found(self, handler, samp
     @pytest.mark.asyncio
     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."""
-        # Mock workflow function
-        mock_workflow_func = Mock()
-        mock_registry.get_workflow.return_value = mock_workflow_func
+        # Create actual workflow definition
+        def mock_workflow_func():
+            return "test_result"
+
+        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions
+        workflow_opts = WorkflowDefinitionOptions(name="test_workflow")
+        workflow_definition = WorkflowDefinition.wrap(mock_workflow_func, workflow_opts)
+        mock_registry.get_workflow.return_value = workflow_definition
 
         # Mock workflow engine
         mock_engine = Mock(spec=WorkflowEngine)
@@ -324,18 +334,20 @@ 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."""
         mock_workflow_func = Mock()
-        mock_registry.get_workflow.return_value = mock_workflow_func
-        
+        mock_workflow_definition = Mock()
+        mock_workflow_definition.fn = mock_workflow_func
+        mock_registry.get_workflow.return_value = mock_workflow_definition
+
         mock_engine = Mock(spec=WorkflowEngine)
         mock_engine._is_workflow_complete = False  # Add missing attribute
         mock_decision_result = Mock(spec=DecisionResult)
         mock_decision_result.decisions = []
         mock_engine.process_decision = AsyncMock(return_value=mock_decision_result)
-        
+
         with patch('cadence.worker._decision_task_handler.WorkflowEngine', return_value=mock_engine) as mock_workflow_engine_class:
             with patch('cadence.worker._decision_task_handler.WorkflowInfo') as mock_workflow_info_class:
                 await handler._handle_task_implementation(sample_decision_task)
-                
+
                 # Verify WorkflowInfo was created with correct parameters (called once for engine)
                 assert mock_workflow_info_class.call_count == 1
                 for call in mock_workflow_info_class.call_args_list:
@@ -345,7 +357,7 @@ async def test_workflow_engine_creation_with_workflow_info(self, handler, sample
                         'workflow_id': "test_workflow_id",
                         'workflow_run_id': "test_run_id"
                     }
-                
+
                 # Verify WorkflowEngine was created with correct parameters
                 mock_workflow_engine_class.assert_called_once()
                 call_args = mock_workflow_engine_class.call_args

tests/cadence/worker/test_decision_task_handler_integration.py

@@ -35,12 +35,12 @@ def mock_client(self):
     def registry(self):
         """Create a registry with a test workflow."""
         reg = Registry()
-        
-        @reg.workflow
+
+        @reg.workflow(name="test_workflow")
         def test_workflow(input_data):
             """Simple test workflow that returns the input."""
             return f"processed: {input_data}"
-        
+
         return reg
 
     @pytest.fixture

tests/cadence/worker/test_registry.py

@@ -30,16 +30,21 @@ def test_basic_registration_and_retrieval(self, registration_type):
             @reg.workflow
             def test_func():
                 return "test"
-            
-            func = reg.get_workflow("test_func")
+
+            # Registry stores WorkflowDefinition internally
+            func_def = reg.get_workflow(test_func.__name__)
+            # WorkflowDefinition can be called directly
+            assert func_def() == "test"
+            # Verify it's actually a WorkflowDefinition
+            from cadence.workflow import WorkflowDefinition
+            assert isinstance(func_def, WorkflowDefinition)
         else:
             @reg.activity
             def test_func():
                 return "test"
 
             func = reg.get_activity(test_func.name)
-        
-        assert func() == "test"
+            assert func() == "test"
 
     def test_direct_call_behavior(self):
         reg = Registry()