improve context management

timl3136 · timl3136 · commit f6a8ed85ba6b · 2025-09-16T10:02:56.000-07:00
diff --git a/cadence/worker/_base_task_handler.py b/cadence/worker/_base_task_handler.py
@@ -34,26 +34,19 @@ async def handle_task(self, task: T) -> None:
         Handle a single task.
         
         This method provides the base implementation for task handling that includes:
-        - Context propagation
         - Error handling
         - Cleanup
         
         Args:
             task: The task to handle
         """
         try:
-            # Propagate context from task parameters
-            await self._propagate_context(task)
-            
-            # Handle the task
+            # Handle the task implementation
             await self._handle_task_implementation(task)
                 
         except Exception as e:
             logger.exception(f"Error handling task: {e}")
             await self.handle_task_failure(task, e)
-        finally:
-            # Clean up context
-            await self._unset_current_context()
     
     @abstractmethod
     async def _handle_task_implementation(self, task: T) -> None:
@@ -75,20 +68,3 @@ async def handle_task_failure(self, task: T, error: Exception) -> None:
             error: The exception that occurred
         """
         pass
-    
-    async def _propagate_context(self, task: T) -> None:
-        """
-        Propagate context from task parameters.
-        
-        Args:
-            task: The task containing context information
-        """
-        # Default implementation - subclasses should override if needed
-        pass
-    
-    async def _unset_current_context(self) -> None:
-        """
-        Unset the current context after task completion.
-        """
-        # Default implementation - subclasses should override if needed
-        pass
diff --git a/cadence/worker/_decision_task_handler.py b/cadence/worker/_decision_task_handler.py
@@ -11,6 +11,7 @@
 from cadence.client import Client
 from cadence.worker._base_task_handler import BaseTaskHandler
 from cadence._internal.workflow.workflow_engine import WorkflowEngine, DecisionResult
+from cadence._internal.workflow.context import Context
 from cadence.workflow import WorkflowInfo
 from cadence.worker._registry import Registry
 
@@ -52,8 +53,7 @@ async def _handle_task_implementation(self, task: PollForDecisionTaskResponse) -
         
         if not workflow_execution or not workflow_type:
             logger.error("Decision task missing workflow execution or type")
-            await self.handle_task_failure(task, ValueError("Missing workflow execution or type"))
-            return
+            raise ValueError("Missing workflow execution or type")
         
         workflow_id = workflow_execution.workflow_id
         run_id = workflow_execution.run_id
@@ -69,8 +69,7 @@ async def _handle_task_implementation(self, task: PollForDecisionTaskResponse) -
                 workflow_func = self._registry.get_workflow(workflow_type_name)
             except KeyError:
                 logger.error(f"Workflow type '{workflow_type_name}' not found in registry")
-                await self.handle_task_failure(task, KeyError(f"Workflow type '{workflow_type_name}' not found"))
-                return
+                raise KeyError(f"Workflow type '{workflow_type_name}' not found")
             
             # Create workflow info and engine
             workflow_info = WorkflowInfo(
@@ -86,12 +85,22 @@ async def _handle_task_implementation(self, task: PollForDecisionTaskResponse) -
                 workflow_func=workflow_func
             )
         
-        # Process the decision using the workflow engine
+        # Create workflow context and execute with it active
         workflow_engine = self._workflow_engines[engine_key]
-        decision_result = await workflow_engine.process_decision(task)
+        workflow_info = WorkflowInfo(
+            workflow_type=workflow_type_name,
+            workflow_domain=self._client.domain,
+            workflow_id=workflow_id,
+            workflow_run_id=run_id
+        )
         
-        # Respond with the decisions
-        await self._respond_decision_task_completed(task, decision_result)
+        context = Context(client=self._client, info=workflow_info)
+        with context._activate():
+            # Process the decision using the workflow engine
+            decision_result = await workflow_engine.process_decision(task)
+            
+            # Respond with the decisions
+            await self._respond_decision_task_completed(task, decision_result)
         
         logger.info(f"Successfully processed decision task for workflow {workflow_id}")
     
diff --git a/tests/cadence/worker/test_base_task_handler.py b/tests/cadence/worker/test_base_task_handler.py
@@ -16,8 +16,6 @@ def __init__(self, client, task_list: str, identity: str, **options):
         super().__init__(client, task_list, identity, **options)
         self._handle_task_implementation_called = False
         self._handle_task_failure_called = False
-        self._propagate_context_called = False
-        self._unset_current_context_called = False
         self._last_task: str = ""
         self._last_error: Exception | None = None
     
@@ -33,15 +31,6 @@ async def handle_task_failure(self, task: str, error: Exception) -> None:
         self._handle_task_failure_called = True
         self._last_task = task
         self._last_error = error
-    
-    async def _propagate_context(self, task: str) -> None:
-        """Test implementation of context propagation."""
-        self._propagate_context_called = True
-        self._last_task = task
-    
-    async def _unset_current_context(self) -> None:
-        """Test implementation of context cleanup."""
-        self._unset_current_context_called = True
 
 
 class TestBaseTaskHandler:
@@ -71,10 +60,8 @@ async def test_handle_task_success(self):
         
         await handler.handle_task("test_task")
         
-        # Verify all methods were called in correct order
-        assert handler._propagate_context_called
+        # Verify implementation was called
         assert handler._handle_task_implementation_called
-        assert handler._unset_current_context_called
         assert not handler._handle_task_failure_called
         assert handler._last_task == "test_task"
         assert handler._last_error is None
@@ -88,72 +75,12 @@ async def test_handle_task_failure(self):
         await handler.handle_task("raise_error")
         
         # Verify error handling was called
-        assert handler._propagate_context_called
         assert handler._handle_task_implementation_called
         assert handler._handle_task_failure_called
-        assert handler._unset_current_context_called
         assert handler._last_task == "raise_error"
         assert isinstance(handler._last_error, ValueError)
         assert str(handler._last_error) == "Test error"
     
-    @pytest.mark.asyncio
-    async def test_handle_task_with_context_propagation_error(self):
-        """Test task handling when context propagation fails."""
-        client = Mock()
-        handler = ConcreteTaskHandler(client, "test_task_list", "test_identity")
-        
-        # Override _propagate_context to raise an error
-        async def failing_propagate_context(task):
-            raise RuntimeError("Context propagation failed")
-        
-        # Use setattr to avoid mypy error about method assignment
-        setattr(handler, '_propagate_context', failing_propagate_context)
-        
-        await handler.handle_task("test_task")
-        
-        # Verify error handling was called
-        assert handler._handle_task_failure_called
-        assert handler._unset_current_context_called
-        assert isinstance(handler._last_error, RuntimeError)
-        assert str(handler._last_error) == "Context propagation failed"
-    
-    @pytest.mark.asyncio
-    async def test_handle_task_with_cleanup_error(self):
-        """Test task handling when cleanup fails."""
-        client = Mock()
-        handler = ConcreteTaskHandler(client, "test_task_list", "test_identity")
-        
-        # Override _unset_current_context to raise an error
-        async def failing_unset_context():
-            raise RuntimeError("Cleanup failed")
-        
-        # Use setattr to avoid mypy error about method assignment
-        setattr(handler, '_unset_current_context', failing_unset_context)
-        
-        # Cleanup errors in finally block will propagate
-        with pytest.raises(RuntimeError, match="Cleanup failed"):
-            await handler.handle_task("test_task")
-    
-    @pytest.mark.asyncio
-    async def test_handle_task_with_implementation_and_cleanup_errors(self):
-        """Test task handling when both implementation and cleanup fail."""
-        client = Mock()
-        handler = ConcreteTaskHandler(client, "test_task_list", "test_identity")
-        
-        # Override _unset_current_context to raise an error
-        async def failing_unset_context():
-            raise RuntimeError("Cleanup failed")
-        
-        # Use setattr to avoid mypy error about method assignment
-        setattr(handler, '_unset_current_context', failing_unset_context)
-        
-        # The implementation error should be handled, but cleanup error will propagate
-        with pytest.raises(RuntimeError, match="Cleanup failed"):
-            await handler.handle_task("raise_error")
-        
-        # Verify the implementation error was handled before cleanup error
-        assert handler._handle_task_failure_called
-        assert isinstance(handler._last_error, ValueError)
     
     @pytest.mark.asyncio
     async def test_abstract_methods_not_implemented(self):
@@ -175,17 +102,6 @@ async def handle_task_failure(self, task: str, error: Exception) -> None:
         with pytest.raises(NotImplementedError):
             await handler.handle_task_failure("test", Exception("test"))
     
-    @pytest.mark.asyncio
-    async def test_default_context_methods(self):
-        """Test default implementations of context methods."""
-        client = Mock()
-        handler = ConcreteTaskHandler(client, "test_task_list", "test_identity")
-        
-        # Test default _propagate_context (should not raise)
-        await handler._propagate_context("test_task")
-        
-        # Test default _unset_current_context (should not raise)
-        await handler._unset_current_context()
     
     @pytest.mark.asyncio
     async def test_generic_type_parameter(self):
diff --git a/tests/cadence/worker/test_decision_task_handler.py b/tests/cadence/worker/test_decision_task_handler.py
@@ -114,14 +114,8 @@ async def test_handle_task_implementation_missing_workflow_execution(self, handl
         task.workflow_type = Mock()
         task.workflow_type.name = "TestWorkflow"
         
-        with patch.object(handler, 'handle_task_failure', new_callable=AsyncMock) as mock_handle_failure:
+        with pytest.raises(ValueError, match="Missing workflow execution or type"):
             await handler._handle_task_implementation(task)
-            
-            mock_handle_failure.assert_called_once()
-            args = mock_handle_failure.call_args[0]
-            assert args[0] == task
-            assert isinstance(args[1], ValueError)
-            assert "Missing workflow execution or type" in str(args[1])
     
     @pytest.mark.asyncio
     async def test_handle_task_implementation_missing_workflow_type(self, handler):
@@ -133,28 +127,16 @@ async def test_handle_task_implementation_missing_workflow_type(self, handler):
         task.workflow_execution.run_id = "test_run_id"
         task.workflow_type = None
         
-        with patch.object(handler, 'handle_task_failure', new_callable=AsyncMock) as mock_handle_failure:
+        with pytest.raises(ValueError, match="Missing workflow execution or type"):
             await handler._handle_task_implementation(task)
-            
-            mock_handle_failure.assert_called_once()
-            args = mock_handle_failure.call_args[0]
-            assert args[0] == task
-            assert isinstance(args[1], ValueError)
-            assert "Missing workflow execution or type" in str(args[1])
     
     @pytest.mark.asyncio
     async def test_handle_task_implementation_workflow_not_found(self, handler, sample_decision_task, mock_registry):
         """Test decision task handling when workflow is not found in registry."""
         mock_registry.get_workflow.side_effect = KeyError("Workflow not found")
         
-        with patch.object(handler, 'handle_task_failure', new_callable=AsyncMock) as mock_handle_failure:
+        with pytest.raises(KeyError, match="Workflow type 'TestWorkflow' not found"):
             await handler._handle_task_implementation(sample_decision_task)
-            
-            mock_handle_failure.assert_called_once()
-            args = mock_handle_failure.call_args[0]
-            assert args[0] == sample_decision_task
-            assert isinstance(args[1], KeyError)
-            assert "Workflow type 'TestWorkflow' not found" in str(args[1])
     
     @pytest.mark.asyncio
     async def test_handle_task_implementation_reuses_existing_engine(self, handler, sample_decision_task, mock_registry):
@@ -340,13 +322,15 @@ async def test_workflow_engine_creation_with_workflow_info(self, handler, sample
             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
-                mock_workflow_info_class.assert_called_once_with(
-                    workflow_type="TestWorkflow",
-                    workflow_domain="test_domain",
-                    workflow_id="test_workflow_id",
-                    workflow_run_id="test_run_id"
-                )
+                # Verify WorkflowInfo was created with correct parameters (called twice - once for engine, once for context)
+                assert mock_workflow_info_class.call_count == 2
+                for call in mock_workflow_info_class.call_args_list:
+                    assert call[1] == {
+                        'workflow_type': "TestWorkflow",
+                        'workflow_domain': "test_domain",
+                        '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()
diff --git a/tests/cadence/worker/test_task_handler_integration.py b/tests/cadence/worker/test_task_handler_integration.py
@@ -4,6 +4,7 @@
 """
 
 import pytest
+from contextlib import contextmanager
 from unittest.mock import Mock, AsyncMock, patch, PropertyMock
 
 from cadence.api.v1.service_worker_pb2 import PollForDecisionTaskResponse
@@ -100,8 +101,8 @@ async def test_full_task_handling_flow_with_error(self, handler, sample_decision
         assert call_args.identity == handler._identity
     
     @pytest.mark.asyncio
-    async def test_context_propagation_integration(self, handler, sample_decision_task, mock_registry):
-        """Test that context propagation works correctly in the integration."""
+    async def test_context_activation_integration(self, handler, sample_decision_task, mock_registry):
+        """Test that context activation works correctly in the integration."""
         # Mock workflow function
         mock_workflow_func = Mock()
         mock_registry.get_workflow.return_value = mock_workflow_func
@@ -114,27 +115,23 @@ async def test_context_propagation_integration(self, handler, sample_decision_ta
         mock_decision_result.query_results = {}
         mock_engine.process_decision = AsyncMock(return_value=mock_decision_result)
         
-        # Track if context methods are called
-        context_propagated = False
-        context_unset = False
+        # Track if context is activated
+        context_activated = False
         
-        async def track_propagate_context(task):
-            nonlocal context_propagated
-            context_propagated = True
-        
-        async def track_unset_current_context():
-            nonlocal context_unset
-            context_unset = True
-        
-        handler._propagate_context = track_propagate_context
-        handler._unset_current_context = track_unset_current_context
+        def track_context_activation():
+            nonlocal context_activated
+            context_activated = True
         
         with patch('cadence.worker._decision_task_handler.WorkflowEngine', return_value=mock_engine):
-            await handler.handle_task(sample_decision_task)
-        
-        # Verify context methods were called
-        assert context_propagated
-        assert context_unset
+            with patch('cadence.worker._decision_task_handler.Context') as mock_context_class:
+                mock_context = Mock()
+                mock_context._activate = Mock(return_value=contextmanager(lambda: track_context_activation())())
+                mock_context_class.return_value = mock_context
+                
+                await handler.handle_task(sample_decision_task)
+        
+        # Verify context was activated
+        assert context_activated
     
     @pytest.mark.asyncio
     async def test_multiple_workflow_executions(self, handler, mock_registry):
@@ -235,19 +232,22 @@ async def test_error_handling_with_context_cleanup(self, handler, sample_decisio
         mock_engine.process_decision = AsyncMock(side_effect=RuntimeError("Workflow processing failed"))
         
         # Track context cleanup
-        context_unset = False
-        
-        async def track_unset_current_context():
-            nonlocal context_unset
-            context_unset = True
+        context_cleaned_up = False
         
-        handler._unset_current_context = track_unset_current_context
+        def track_context_cleanup():
+            nonlocal context_cleaned_up
+            context_cleaned_up = True
         
         with patch('cadence.worker._decision_task_handler.WorkflowEngine', return_value=mock_engine):
-            await handler.handle_task(sample_decision_task)
+            with patch('cadence.worker._decision_task_handler.Context') as mock_context_class:
+                mock_context = Mock()
+                mock_context._activate = Mock(return_value=contextmanager(lambda: track_context_cleanup())())
+                mock_context_class.return_value = mock_context
+                
+                await handler.handle_task(sample_decision_task)
         
         # Verify context was cleaned up even after error
-        assert context_unset
+        assert context_cleaned_up
         
         # Verify error was handled
         handler._client.worker_stub.RespondDecisionTaskFailed.assert_called_once()
