Implement entity execution and worker support with comprehensive tests

Co-authored-by: berndverst <4535280+berndverst@users.noreply.github.com>

Author: Copilot

durabletask/worker.py

@@ -75,10 +75,12 @@ def __init__(
 class _Registry:
     orchestrators: dict[str, task.Orchestrator]
     activities: dict[str, task.Activity]
+    entities: dict[str, task.Entity]
 
     def __init__(self):
         self.orchestrators = {}
         self.activities = {}
+        self.entities = {}
 
     def add_orchestrator(self, fn: task.Orchestrator) -> str:
         if fn is None:
@@ -118,6 +120,25 @@ def add_named_activity(self, name: str, fn: task.Activity) -> None:
     def get_activity(self, name: str) -> Optional[task.Activity]:
         return self.activities.get(name)
 
+    def add_entity(self, fn: task.Entity) -> str:
+        if fn is None:
+            raise ValueError("An entity function argument is required.")
+
+        name = task.get_name(fn)
+        self.add_named_entity(name, fn)
+        return name
+
+    def add_named_entity(self, name: str, fn: task.Entity) -> None:
+        if not name:
+            raise ValueError("A non-empty entity name is required.")
+        if name in self.entities:
+            raise ValueError(f"A '{name}' entity already exists.")
+
+        self.entities[name] = fn
+
+    def get_entity(self, name: str) -> Optional[task.Entity]:
+        return self.entities.get(name)
+
 
 class OrchestratorNotRegisteredError(ValueError):
     """Raised when attempting to start an orchestration that is not registered"""
@@ -131,6 +152,12 @@ class ActivityNotRegisteredError(ValueError):
     pass
 
 
+class EntityNotRegisteredError(ValueError):
+    """Raised when attempting to call an entity that is not registered"""
+
+    pass
+
+
 class TaskHubGrpcWorker:
     """A gRPC-based worker for processing durable task orchestrations and activities.
 
@@ -279,6 +306,14 @@ def add_activity(self, fn: task.Activity) -> str:
             )
         return self._registry.add_activity(fn)
 
+    def add_entity(self, fn: task.Entity) -> str:
+        """Registers an entity function with the worker."""
+        if self._is_running:
+            raise RuntimeError(
+                "Entities cannot be added while the worker is running."
+            )
+        return self._registry.add_entity(fn)
+
     def start(self):
         """Starts the worker on a background thread and begins listening for work items."""
         if self._is_running:
@@ -434,6 +469,13 @@ def stream_reader():
                                 stub,
                                 work_item.completionToken,
                             )
+                        elif work_item.HasField("entityRequest"):
+                            self._async_worker_manager.submit_activity(
+                                self._execute_entity,
+                                work_item.entityRequest,
+                                stub,
+                                work_item.completionToken,
+                            )
                         elif work_item.HasField("healthPing"):
                             pass
                         else:
@@ -569,6 +611,34 @@ def _execute_activity(
                 f"Failed to deliver activity response for '{req.name}#{req.taskId}' of orchestration ID '{instance_id}' to sidecar: {ex}"
             )
 
+    def _execute_entity(
+            self,
+            req: pb.EntityBatchRequest,
+            stub: stubs.TaskHubSidecarServiceStub,
+            completionToken,
+    ):
+        instance_id = req.instanceId
+        try:
+            executor = _EntityExecutor(self._registry, self._logger)
+            result = executor.execute(req)
+            result.completionToken = completionToken
+        except Exception as ex:
+            self._logger.exception(
+                f"An error occurred while trying to execute entity '{instance_id}': {ex}"
+            )
+            failure_details = ph.new_failure_details(ex)
+            result = pb.EntityBatchResult(
+                failureDetails=failure_details,
+                completionToken=completionToken,
+            )
+
+        try:
+            stub.CompleteEntityTask(result)
+        except Exception as ex:
+            self._logger.exception(
+                f"Failed to deliver entity response for entity '{instance_id}' to sidecar: {ex}"
+            )
+
 
 class _RuntimeOrchestrationContext(task.OrchestrationContext):
     _generator: Optional[Generator[task.Task, Any, Any]]
@@ -858,6 +928,36 @@ def continue_as_new(self, new_input, *, save_events: bool = False) -> None:
 
         self.set_continued_as_new(new_input, save_events)
 
+    def signal_entity(self, entity_id: str, operation_name: str, *,
+                      input: Optional[Any] = None) -> task.Task:
+        # Create a signal entity action
+        action = pb.OrchestratorAction()
+        action.sendEntitySignal.CopyFrom(pb.SendSignalAction(
+            instanceId=entity_id,
+            name=operation_name,
+            input=ph.get_string_value(shared.to_json(input)) if input is not None else None
+        ))
+
+        # Entity signals don't return values, so we create a completed task
+        signal_task = task.CompletableTask()
+        
+        # Store the action to be executed
+        task_id = self._next_task_id()
+        self._pending_actions[task_id] = action
+        self._pending_tasks[task_id] = signal_task
+        
+        # Mark as complete since signals don't have return values
+        signal_task.complete(None)
+        
+        return signal_task
+
+    def call_entity(self, entity_id: str, operation_name: str, *,
+                    input: Optional[Any] = None,
+                    retry_policy: Optional[task.RetryPolicy] = None) -> task.Task:
+        # For now, entity calls are not directly supported in orchestrations
+        # This would require additional protobuf support
+        raise NotImplementedError("Direct entity calls from orchestrations are not yet supported. Use signal_entity instead.")
+
 
 class ExecutionResults:
     actions: list[pb.OrchestratorAction]
@@ -1260,6 +1360,81 @@ def execute(
         return encoded_output
 
 
+class _EntityExecutor:
+    def __init__(self, registry: _Registry, logger: logging.Logger):
+        self._registry = registry
+        self._logger = logger
+
+    def execute(self, req: pb.EntityBatchRequest) -> pb.EntityBatchResult:
+        """Executes entity operations and returns the batch result."""
+        instance_id = req.instanceId
+        self._logger.debug(f"Executing entity batch for '{instance_id}' with {len(req.operations)} operation(s)...")
+
+        # Parse current entity state
+        current_state = shared.from_json(req.entityState.value) if not ph.is_empty(req.entityState) else None
+        
+        # Extract entity type from instance ID (format: entitytype@key)
+        entity_type = "Unknown"
+        if "@" in instance_id:
+            entity_type = instance_id.split("@")[0]
+        
+        results = []
+        actions = []
+        
+        for operation in req.operations:
+            try:
+                # Get the entity function using the entity type from instanceId
+                fn = self._registry.get_entity(entity_type)
+                if not fn:
+                    raise EntityNotRegisteredError(f"Entity function named '{entity_type}' was not registered!")
+
+                # Create entity context
+                ctx = task.EntityContext(
+                    instance_id=instance_id,
+                    operation_name=operation.operation,
+                    is_new_entity=(current_state is None)
+                )
+                ctx.set_state(current_state)
+
+                # Parse operation input
+                operation_input = shared.from_json(operation.input.value) if not ph.is_empty(operation.input) else None
+
+                # Execute the entity operation
+                operation_output = fn(ctx, operation_input)
+
+                # Update state for next operation
+                current_state = ctx.get_state()
+
+                # Create operation result
+                result = pb.OperationResult()
+                if operation_output is not None:
+                    result.success.CopyFrom(pb.OperationResultSuccess(
+                        result=ph.get_string_value(shared.to_json(operation_output))
+                    ))
+                else:
+                    result.success.CopyFrom(pb.OperationResultSuccess())
+                
+                results.append(result)
+
+            except Exception as ex:
+                self._logger.exception(f"Error executing entity operation '{operation.operation}' on entity type '{entity_type}': {ex}")
+                
+                # Create failure result
+                failure_details = ph.new_failure_details(ex)
+                result = pb.OperationResult()
+                result.failure.CopyFrom(pb.OperationResultFailure(
+                    failureDetails=failure_details
+                ))
+                results.append(result)
+
+        # Return batch result
+        return pb.EntityBatchResult(
+            results=results,
+            actions=actions,
+            entityState=ph.get_string_value(shared.to_json(current_state)) if current_state is not None else None
+        )
+
+
 def _get_non_determinism_error(
         task_id: int, action_name: str
 ) -> task.NonDeterminismError:

tests/durabletask/test_entities.py

@@ -4,6 +4,7 @@
 import unittest
 from datetime import datetime
 from durabletask import task
+from durabletask import worker as task_worker
 
 
 class TestEntityTypes(unittest.TestCase):
@@ -108,5 +109,81 @@ def test_entity_query_result_creation(self):
         self.assertEqual(result.continuation_token, "next-page-token")
 
 
+class TestEntityWorkerIntegration(unittest.TestCase):
+    
+    def test_worker_entity_registration(self):
+        """Test that entities can be registered with the worker."""
+        worker = task_worker.TaskHubGrpcWorker()
+        
+        def counter_entity(ctx: task.EntityContext, input):
+            if ctx.operation_name == "increment":
+                current_count = ctx.get_state() or 0
+                new_count = current_count + (input or 1)
+                ctx.set_state(new_count)
+                return new_count
+            elif ctx.operation_name == "get":
+                return ctx.get_state() or 0
+            elif ctx.operation_name == "reset":
+                ctx.set_state(0)
+                return 0
+        
+        # Test registration
+        entity_name = worker.add_entity(counter_entity)
+        self.assertEqual(entity_name, "counter_entity")
+        
+        # Test that entity is in registry
+        self.assertIsNotNone(worker._registry.get_entity("counter_entity"))
+        
+        # Test error for duplicate registration
+        with self.assertRaises(ValueError):
+            worker.add_entity(counter_entity)
+    
+    def test_entity_execution(self):
+        """Test entity execution via the EntityExecutor."""
+        from durabletask.worker import _Registry, _EntityExecutor
+        import durabletask.internal.orchestrator_service_pb2 as pb
+        import durabletask.internal.helpers as ph
+        import logging
+        
+        # Create registry and register entity
+        registry = _Registry()
+        
+        def counter_entity(ctx: task.EntityContext, input):
+            if ctx.operation_name == "increment":
+                current_count = ctx.get_state() or 0
+                new_count = current_count + (input or 1)
+                ctx.set_state(new_count)
+                return new_count
+            elif ctx.operation_name == "get":
+                return ctx.get_state() or 0
+        
+        # Register the entity with a specific name 
+        registry.add_named_entity("Counter", counter_entity)
+        
+        # Create executor
+        logger = logging.getLogger("test")
+        executor = _EntityExecutor(registry, logger)
+        
+        # Create test request
+        req = pb.EntityBatchRequest()
+        req.instanceId = "Counter@test-key"  # Instance ID with entity type prefix matching registration
+        req.entityState.CopyFrom(ph.get_string_value("0"))  # Initial state
+        
+        # Add increment operation
+        operation = pb.OperationRequest()
+        operation.operation = "increment"
+        operation.input.CopyFrom(ph.get_string_value("5"))
+        req.operations.append(operation)
+        
+        # Execute
+        result = executor.execute(req)
+        
+        # Verify result
+        self.assertEqual(len(result.results), 1)
+        self.assertTrue(result.results[0].HasField("success"))
+        self.assertEqual(result.results[0].success.result.value, "5")
+        self.assertEqual(result.entityState.value, "5")
+
+
 if __name__ == '__main__':
     unittest.main()