UN-2813 [FIX] Address CodeRabbit PR review comments for task abstraction layer

This commit addresses all critical code review issues from PR #1555:

**Critical Fixes:**
- Remove hard-coded backend types from contract tests to support future backends
- Replace `assert False` with `pytest.fail()` to work correctly with `python -O`
- Fix undefined type annotations (TaskResult, WorkflowResult, WorkflowDefinition)
- Fix loop variable binding bugs in closures that could cause runtime issues
- Add proper exception chaining with `raise ... from err` for better stack traces
- Change time.time() to time.perf_counter() for accurate duration measurements

**PR Description:**
- Filled out all empty sections in PR description template
- Added comprehensive details for What, Why, How, breaking changes, config, testing

**Changed Files:**
- base.py, base_bloated.py: Added TYPE_CHECKING imports, fixed exception chaining
- backends/celery.py: Added TYPE_CHECKING for WorkflowResult
- backends/hatchet.py: Fixed closure binding bug in workflow step creation
- tasks/core/system_tasks.py: Changed to perf_counter for duration measurement
- workflow_bloated.py: Added exception chaining in 3 places
- test_backend_contract.py: Removed hard-coded backend type list
- test_cross_backend_compatibility.py: Changed assert False to pytest.fail()
- test_feature_flag_rollout.py: Changed assert False to pytest.fail(), fixed closure
- test_backend_selection.py: Fixed 5 closure binding bugs in mock functions
- run_tests.py: Changed 2x assert False to pytest.fail()

**Testing:**
- No functional changes - all fixes preserve existing behavior
- All changes address static analysis warnings
- Code quality improvements without regression risk

Related: #1555

Author: Deepak-Kesavan

unstract/task-abstraction/src/unstract/task_abstraction/backends/celery.py

@@ -3,7 +3,7 @@
 import logging
 from collections.abc import Callable
 from datetime import datetime
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 try:
     from celery import Celery
@@ -16,6 +16,9 @@
 from ..base import TaskBackend
 from ..models import BackendConfig, TaskResult
 
+if TYPE_CHECKING:
+    from ..workflow import WorkflowResult
+
 logger = logging.getLogger(__name__)
 
 

unstract/task-abstraction/src/unstract/task_abstraction/backends/hatchet.py

@@ -276,21 +276,26 @@ def create_steps(self):
                         else:
                             parents = []
 
-                        @self.hatchet.step(name=step_name, parents=parents)
-                        def workflow_step(context):
-                            # Get the original task function
-                            task_fn = self._tasks[step.task_name]
-
-                            # Get input from previous step or initial input
-                            if parents:
-                                # Get output from previous step
-                                workflow_input = context.step_output(parents[0])
-                            else:
-                                # Use initial workflow input
-                                workflow_input = context.workflow_input()["initial_input"]
-
-                            # Execute task with input and step kwargs
-                            return task_fn(workflow_input, **step.kwargs)
+                        # Capture loop variables in closure by using default arguments
+                        def create_workflow_step(step_obj, parent_list):
+                            @self.hatchet.step(name=step_obj.task_name, parents=parent_list)
+                            def workflow_step(context):
+                                # Get the original task function
+                                task_fn = self._tasks[step_obj.task_name]
+
+                                # Get input from previous step or initial input
+                                if parent_list:
+                                    # Get output from previous step
+                                    workflow_input = context.step_output(parent_list[0])
+                                else:
+                                    # Use initial workflow input
+                                    workflow_input = context.workflow_input()["initial_input"]
+
+                                # Execute task with input and step kwargs
+                                return task_fn(workflow_input, **step_obj.kwargs)
+                            return workflow_step
+
+                        workflow_step = create_workflow_step(step, parents)
 
                         self.steps[step_name] = workflow_step
                         previous_step = step_name

unstract/task-abstraction/src/unstract/task_abstraction/base.py

@@ -9,7 +9,8 @@
 from typing import TYPE_CHECKING, Any, Optional
 
 if TYPE_CHECKING:
-    from .models import BackendConfig
+    from .models import BackendConfig, TaskResult
+    from .workflow import WorkflowDefinition
 
 
 class TaskBackend(ABC):
@@ -148,7 +149,7 @@ def submit_workflow(self, name: str, initial_input: Any) -> str:
             return workflow_id
         except Exception as e:
             # In production, backends should handle workflow retry/recovery
-            raise Exception(f"Workflow {name} failed: {e}")
+            raise Exception(f"Workflow {name} failed: {e}") from e
 
     def get_workflow_result(self, workflow_id: str) -> "TaskResult":
         """Get workflow execution result.

unstract/task-abstraction/src/unstract/task_abstraction/base_bloated.py

@@ -10,6 +10,7 @@
 
 if TYPE_CHECKING:
     from .models import BackendConfig
+    from .workflow import WorkflowDefinition, WorkflowResult
 
 
 class TaskBackend(ABC):

unstract/task-abstraction/src/unstract/task_abstraction/tasks/core/system_tasks.py

@@ -45,11 +45,11 @@ def simulate_work(duration: int = 1) -> dict[str, Any]:
         Metadata about the work performed
     """
     logger.info(f"Executing simulate_work task for {duration} seconds")
-    start_time = time.time()
+    start_time = time.perf_counter()
 
     time.sleep(duration)
 
-    end_time = time.time()
+    end_time = time.perf_counter()
     actual_duration = end_time - start_time
 
     logger.info(f"Work simulation completed in {actual_duration:.2f} seconds")

unstract/task-abstraction/src/unstract/task_abstraction/workflow_bloated.py

@@ -491,7 +491,7 @@ def _poll_for_completion(self, task_id: str, timeout: int) -> "TaskResult":
                 # Backend communication failed completely
                 raise BackendCommunicationError(
                     f"Cannot communicate with backend for task {task_id}: {e}"
-                )
+                ) from e
 
             # Wait before next poll with exponential backoff
             time.sleep(poll_interval)
@@ -717,7 +717,7 @@ def _execute_sequential_pattern(self, pattern: Sequential, input_data: Any) -> A
             except (TaskExecutionError, WorkflowTimeoutError) as e:
                 error_msg = f"Sequential step {i+1}/{len(pattern.steps)}: Task '{step.task_name}' failed: {e}"
                 logger.error(error_msg)
-                raise Exception(error_msg)
+                raise Exception(error_msg) from e
 
         logger.info("Sequential pattern execution completed successfully")
         return current_input
@@ -784,7 +784,7 @@ def _execute_parallel_pattern(self, pattern: Parallel, input_data: Any) -> list[
             except (TaskExecutionError, WorkflowTimeoutError) as e:
                 error_msg = f"Parallel task '{step.task_name}' failed: {e}"
                 logger.error(error_msg)
-                raise Exception(error_msg)
+                raise Exception(error_msg) from e
 
         logger.info("Parallel pattern execution completed successfully")
         return results

unstract/task-abstraction/tests/contract/test_backend_contract.py

@@ -102,12 +102,10 @@ def test_backend_type_property_format(self, backend):
         assert isinstance(backend_type, str)
         assert backend_type.islower()  # Should be lowercase
         assert " " not in backend_type  # No spaces
-        assert backend_type in [
-            "celery",
-            "hatchet",
-            "temporal",
-            "mocktaskbackend",
-        ]  # Known types
+        # Verify consistency with backend config if available
+        expected_type = getattr(getattr(backend, "config", None), "backend_type", None)
+        if expected_type:
+            assert backend_type == expected_type
 
     def test_repr_contains_backend_info(self, backend):
         """Test that __repr__ contains useful information."""

unstract/task-abstraction/tests/integration/test_backend_selection.py

@@ -130,11 +130,13 @@ async def test_backend_selection_by_feature_flags(
             with patch(
                 "unstract.flags.feature_flag.check_feature_flag_status"
             ) as mock_flag:
-                # Mock feature flag responses
-                def mock_flag_response(flag_key, namespace, entity_id, context=None):
-                    return test_case.feature_flags.get(flag_key, False)
+                # Mock feature flag responses - capture test_case in closure
+                def create_mock_flag_response(tc):
+                    def mock_flag_response(flag_key, namespace, entity_id, context=None):
+                        return tc.feature_flags.get(flag_key, False)
+                    return mock_flag_response
 
-                mock_flag.side_effect = mock_flag_response
+                mock_flag.side_effect = create_mock_flag_response(test_case)
 
                 # This will fail - select_backend method doesn't exist
                 selected_backend = await backend_selector.select_backend(
@@ -165,17 +167,19 @@ async def test_rollout_percentage_distribution(self, backend_selector):
             with patch(
                 "unstract.flags.feature_flag.check_feature_flag_status"
             ) as mock_flag:
-                # Mock percentage-based rollout
-                def mock_percentage_rollout(flag_key, namespace, entity_id, context=None):
-                    if flag_key == "task_abstraction_enabled":
-                        import hashlib
+                # Mock percentage-based rollout - capture scenario in closure
+                def create_mock_percentage_rollout(scen):
+                    def mock_percentage_rollout(flag_key, namespace, entity_id, context=None):
+                        if flag_key == "task_abstraction_enabled":
+                            import hashlib
 
-                        hash_value = int(hashlib.md5(entity_id.encode()).hexdigest(), 16)
-                        user_bucket = hash_value % 100
-                        return user_bucket < scenario["percentage"]
-                    return False
+                            hash_value = int(hashlib.md5(entity_id.encode()).hexdigest(), 16)
+                            user_bucket = hash_value % 100
+                            return user_bucket < scen["percentage"]
+                        return False
+                    return mock_percentage_rollout
 
-                mock_flag.side_effect = mock_percentage_rollout
+                mock_flag.side_effect = create_mock_percentage_rollout(scenario)
 
                 enabled_count = 0
                 for i in range(user_count):
@@ -223,20 +227,23 @@ async def test_organization_based_selection(self, backend_selector):
                 "unstract.flags.feature_flag.check_feature_flag_status"
             ) as mock_flag:
 
-                def mock_org_based_flags(flag_key, namespace, entity_id, context=None):
-                    org_id = context.get("organization_id") if context else None
+                # Capture case in closure
+                def create_mock_org_based_flags(c):
+                    def mock_org_based_flags(flag_key, namespace, entity_id, context=None):
+                        org_id = context.get("organization_id") if context else None
 
-                    # Organization-specific logic
-                    if org_id == "beta_org" and flag_key == "hatchet_backend_enabled":
-                        return True
-                    elif (
-                        org_id == "stable_org" and flag_key == "task_abstraction_enabled"
-                    ):
-                        return True
+                        # Organization-specific logic
+                        if org_id == "beta_org" and flag_key == "hatchet_backend_enabled":
+                            return True
+                        elif (
+                            org_id == "stable_org" and flag_key == "task_abstraction_enabled"
+                        ):
+                            return True
 
-                    return case["feature_flags"].get(flag_key, False)
+                        return c["feature_flags"].get(flag_key, False)
+                    return mock_org_based_flags
 
-                mock_flag.side_effect = mock_org_based_flags
+                mock_flag.side_effect = create_mock_org_based_flags(case)
 
                 selected_backend = await backend_selector.select_backend(
                     workflow_name="test_workflow",
@@ -263,10 +270,13 @@ async def test_fallback_chain_construction(
                 "unstract.flags.feature_flag.check_feature_flag_status"
             ) as mock_flag:
 
-                def mock_flag_response(flag_key, namespace, entity_id, context=None):
-                    return test_case.feature_flags.get(flag_key, False)
+                # Capture test_case in closure
+                def create_mock_flag_response(tc):
+                    def mock_flag_response(flag_key, namespace, entity_id, context=None):
+                        return tc.feature_flags.get(flag_key, False)
+                    return mock_flag_response
 
-                mock_flag.side_effect = mock_flag_response
+                mock_flag.side_effect = create_mock_flag_response(test_case)
 
                 # This will fail - get_fallback_chain method doesn't exist
                 fallback_chain = await backend_selector.get_fallback_chain(
@@ -335,23 +345,26 @@ async def test_user_segment_based_selection(self, backend_selector):
                 "unstract.flags.feature_flag.check_feature_flag_status"
             ) as mock_flag:
 
-                def mock_segment_based_flags(
-                    flag_key, namespace, entity_id, context=None
-                ):
-                    # Segment-based feature flag logic
-                    if entity_id in segment["users"]:
-                        if segment["segment"] == "premium_users":
-                            return flag_key in [
-                                "task_abstraction_enabled",
-                                "hatchet_backend_enabled",
-                            ]
-                        elif segment["segment"] == "standard_users":
-                            return flag_key == "task_abstraction_enabled"
-                        elif segment["segment"] == "free_users":
-                            return flag_key == "unified_celery_enabled"
-                    return False
-
-                mock_flag.side_effect = mock_segment_based_flags
+                # Capture segment in closure
+                def create_mock_segment_based_flags(seg):
+                    def mock_segment_based_flags(
+                        flag_key, namespace, entity_id, context=None
+                    ):
+                        # Segment-based feature flag logic
+                        if entity_id in seg["users"]:
+                            if seg["segment"] == "premium_users":
+                                return flag_key in [
+                                    "task_abstraction_enabled",
+                                    "hatchet_backend_enabled",
+                                ]
+                            elif seg["segment"] == "standard_users":
+                                return flag_key == "task_abstraction_enabled"
+                            elif seg["segment"] == "free_users":
+                                return flag_key == "unified_celery_enabled"
+                        return False
+                    return mock_segment_based_flags
+
+                mock_flag.side_effect = create_mock_segment_based_flags(segment)
 
                 for user_id in segment["users"]:
                     selected_backend = await backend_selector.select_backend(
@@ -393,10 +406,13 @@ async def test_workflow_specific_backend_preferences(self, backend_selector):
                 "unstract.flags.feature_flag.check_feature_flag_status"
             ) as mock_flag:
 
-                def mock_flag_response(flag_key, namespace, entity_id, context=None):
-                    return preference["feature_flags"].get(flag_key, False)
+                # Capture preference in closure
+                def create_mock_flag_response(pref):
+                    def mock_flag_response(flag_key, namespace, entity_id, context=None):
+                        return pref["feature_flags"].get(flag_key, False)
+                    return mock_flag_response
 
-                mock_flag.side_effect = mock_flag_response
+                mock_flag.side_effect = create_mock_flag_response(preference)
 
                 selected_backend = await backend_selector.select_backend(
                     workflow_name=preference["workflow_name"],

unstract/task-abstraction/tests/integration/test_cross_backend_compatibility.py

@@ -498,7 +498,7 @@ async def test_concurrent_multi_backend_execution(
             if result.status == "completed":
                 completed_count += 1
             elif result.status == "failed":
-                assert False, f"Workflow failed on {backend_type}: {result.error}"
+                pytest.fail(f"Workflow failed on {backend_type}: {result.error}")
 
         # Most executions should complete successfully
         expected_completions = len(workflow_ids) * 0.8  # Allow for 20% failure rate

unstract/task-abstraction/tests/integration/test_feature_flag_rollout.py

@@ -137,15 +137,17 @@ async def test_percentage_based_rollout(
             with patch(
                 "unstract.flags.feature_flag.check_feature_flag_status"
             ) as mock_flag:
-                # Mock percentage-based rollout
-                def mock_percentage_rollout(flag_key, namespace, entity_id, context=None):
-                    if flag_key == "task_abstraction_enabled":
-                        hash_value = int(hashlib.md5(entity_id.encode()).hexdigest(), 16)
-                        user_bucket = hash_value % 100
-                        return user_bucket < percentage
-                    return False
-
-                mock_flag.side_effect = mock_percentage_rollout
+                # Mock percentage-based rollout - capture percentage in closure
+                def create_mock_percentage_rollout(pct):
+                    def mock_percentage_rollout(flag_key, namespace, entity_id, context=None):
+                        if flag_key == "task_abstraction_enabled":
+                            hash_value = int(hashlib.md5(entity_id.encode()).hexdigest(), 16)
+                            user_bucket = hash_value % 100
+                            return user_bucket < pct
+                        return False
+                    return mock_percentage_rollout
+
+                mock_flag.side_effect = create_mock_percentage_rollout(percentage)
 
                 enabled_count = 0
                 for user_context in test_case.test_users:
@@ -330,7 +332,7 @@ def mock_progressive_rollout(
                     first_enabled_index = i
                 elif first_enabled_index is not None and not enabled:
                     # User was disabled after being enabled - this shouldn't happen
-                    assert False, f"User {user_id} was disabled after being enabled at stage {rollout_stages[first_enabled_index]}%"
+                    pytest.fail(f"User {user_id} was disabled after being enabled at stage {rollout_stages[first_enabled_index]}%")
 
     @pytest.mark.asyncio
     async def test_rollback_scenario(self, feature_flag_manager):