UN-3721 [FIX] PG finalization strand — terminal guard at the model layer (backend stale-object clobber)

backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py

@@ -233,6 +233,192 @@ def test_pg_allows_idempotent_completed_rewrite(self):
         assert ex.status == ExecutionStatus.COMPLETED.value
 
 
+class ModelStaleWriterGuardTests(TestCase):
+    """The terminal-one-way guard at the MODEL layer (update_execution) — where the
+    HTTP-endpoint guard can't reach. A stale backend object (created EXECUTING, NULL
+    counters, re-saved after the callback set COMPLETED) must not revert a
+    protected-terminal (COMPLETED/STOPPED) PG execution back to EXECUTING+NULL."""
+
+    def setUp(self):
+        self.wf = Workflow.objects.create(workflow_name="wf-stale")
+
+    def _stale(self, ex, status, **fields):
+        """A separate in-memory instance of ex with stale field values."""
+        stale = WorkflowExecution.objects.get(pk=ex.pk)
+        stale.status = status.value
+        for k, v in fields.items():
+            setattr(stale, k, v)
+        return stale
+
+    def test_pg_refuses_stale_revert_of_completed(self):
+        ex = WorkflowExecution.objects.create(
+            workflow=self.wf,
+            status=ExecutionStatus.COMPLETED,
+            queue_message_id=11,
+            total_files=1,
+            successful_files=1,
+            failed_files=0,
+        )
+        stale = self._stale(
+            ex, ExecutionStatus.EXECUTING, successful_files=None, failed_files=None
+        )
+        stale.update_execution(status=ExecutionStatus.EXECUTING)  # the clobber
+        ex.refresh_from_db()
+        assert ex.status == ExecutionStatus.COMPLETED.value  # not reverted
+        assert ex.successful_files == 1  # counters not nulled
+
+    def test_pg_refuses_stale_revert_of_stopped(self):
+        ex = WorkflowExecution.objects.create(
+            workflow=self.wf,
+            status=ExecutionStatus.STOPPED,
+            queue_message_id=12,
+            total_files=2,
+            successful_files=1,
+            failed_files=1,
+        )
+        self._stale(
+            ex, ExecutionStatus.EXECUTING, successful_files=None, failed_files=None
+        ).update_execution(status=ExecutionStatus.EXECUTING)
+        ex.refresh_from_db()
+        assert ex.status == ExecutionStatus.STOPPED.value
+        assert ex.successful_files == 1 and ex.failed_files == 1  # counters preserved
+
+    def test_pg_same_status_does_not_clobber_counters(self):
+        # Same-status COMPLETED→COMPLETED: the guard is a no-op, so update_fields must
+        # be what stops the stale NULL counters from clobbering the real ones.
+        ex = WorkflowExecution.objects.create(
+            workflow=self.wf,
+            status=ExecutionStatus.COMPLETED,
+            queue_message_id=17,
+            total_files=1,
+            successful_files=1,
+            failed_files=0,
+        )
+        self._stale(
+            ex, ExecutionStatus.COMPLETED, successful_files=None, failed_files=None
+        ).update_execution(status=ExecutionStatus.COMPLETED)
+        ex.refresh_from_db()
+        assert ex.successful_files == 1  # not nulled
+
+    def test_pg_stale_null_marker_still_guarded_and_not_nulled(self):
+        # A stale object whose in-memory queue_message_id is None (snapshotted before
+        # dispatch recorded it) must STILL be guarded (routing reads the persisted
+        # marker) AND the marker must not be nulled by the write.
+        ex = WorkflowExecution.objects.create(
+            workflow=self.wf,
+            status=ExecutionStatus.COMPLETED,
+            queue_message_id=99,
+            successful_files=1,
+        )
+        self._stale(
+            ex, ExecutionStatus.EXECUTING, queue_message_id=None, successful_files=None
+        ).update_execution(status=ExecutionStatus.EXECUTING)
+        ex.refresh_from_db()
+        assert ex.status == ExecutionStatus.COMPLETED.value  # still guarded
+        assert ex.queue_message_id == 99  # marker NOT nulled
+        assert ex.successful_files == 1
+
+    def test_pg_refused_status_still_applies_error_and_attempt(self):
+        # The refused status must NOT silently drop error / increment_attempt.
+        ex = WorkflowExecution.objects.create(
+            workflow=self.wf,
+            status=ExecutionStatus.COMPLETED,
+            queue_message_id=16,
+            attempts=0,
+        )
+        self._stale(ex, ExecutionStatus.EXECUTING).update_execution(
+            status=ExecutionStatus.EXECUTING,
+            error="late failure",
+            increment_attempt=True,
+        )
+        ex.refresh_from_db()
+        assert ex.status == ExecutionStatus.COMPLETED.value  # status refused
+        assert ex.error_message == "late failure"  # error applied
+        assert ex.attempts == 1  # increment applied
+
+    def test_service_helper_update_execution_is_guarded(self):
+        from workflow_manager.workflow_v2.execution import (
+            WorkflowExecutionServiceHelper,
+        )
+
+        ex = WorkflowExecution.objects.create(
+            workflow=self.wf, status=ExecutionStatus.COMPLETED, queue_message_id=18
+        )
+        helper = WorkflowExecutionServiceHelper.__new__(WorkflowExecutionServiceHelper)
+        helper.execution_id = str(ex.id)
+        helper.update_execution(status=ExecutionStatus.EXECUTING)  # delegates → guarded
+        ex.refresh_from_db()
+        assert ex.status == ExecutionStatus.COMPLETED.value
+
+    def test_update_execution_err_is_guarded(self):
+        from workflow_manager.workflow_v2.execution import (
+            WorkflowExecutionServiceHelper,
+        )
+
+        ex = WorkflowExecution.objects.create(
+            workflow=self.wf, status=ExecutionStatus.COMPLETED, queue_message_id=19
+        )
+        WorkflowExecutionServiceHelper.update_execution_err(str(ex.id), "late error")
+        ex.refresh_from_db()
+        assert ex.status == ExecutionStatus.COMPLETED.value  # COMPLETED not reverted
+
+    def test_pg_allows_error_corrected_to_completed(self):
+        # ERROR is not protected — a premature ERROR must stay correctable.
+        ex = WorkflowExecution.objects.create(
+            workflow=self.wf, status=ExecutionStatus.ERROR, queue_message_id=13
+        )
+        WorkflowExecution.objects.get(pk=ex.pk).update_execution(
+            status=ExecutionStatus.COMPLETED
+        )
+        ex.refresh_from_db()
+        assert ex.status == ExecutionStatus.COMPLETED.value
+
+    def test_pg_allows_idempotent_completed_rewrite(self):
+        ex = WorkflowExecution.objects.create(
+            workflow=self.wf, status=ExecutionStatus.COMPLETED, queue_message_id=14
+        )
+        WorkflowExecution.objects.get(pk=ex.pk).update_execution(
+            status=ExecutionStatus.COMPLETED
+        )
+        ex.refresh_from_db()
+        assert ex.status == ExecutionStatus.COMPLETED.value
+
+    def test_celery_row_unaffected(self):
+        # queue_message_id NULL → legacy behavior, the stale write goes through.
+        ex = WorkflowExecution.objects.create(
+            workflow=self.wf,
+            status=ExecutionStatus.COMPLETED,
+            queue_message_id=None,
+            task_id=uuid.uuid4(),
+        )
+        self._stale(ex, ExecutionStatus.EXECUTING).update_execution(
+            status=ExecutionStatus.EXECUTING
+        )
+        ex.refresh_from_db()
+        assert ex.status == ExecutionStatus.EXECUTING.value  # legacy: overwritten
+
+    def test_result_acknowledge_does_not_touch_status_or_counters(self):
+        # A stale object acknowledging a result must not rewrite status/counters
+        # (update_fields scoping).
+        from workflow_manager.workflow_v2.workflow_helper import WorkflowHelper
+
+        ex = WorkflowExecution.objects.create(
+            workflow=self.wf,
+            status=ExecutionStatus.COMPLETED,
+            queue_message_id=15,
+            total_files=1,
+            successful_files=1,
+            failed_files=0,
+            result_acknowledged=False,
+        )
+        stale = self._stale(
+            ex, ExecutionStatus.EXECUTING, successful_files=None, failed_files=None
+        )
+        WorkflowHelper._set_result_acknowledge(stale)
+        ex.refresh_from_db()
+        assert ex.status == ExecutionStatus.COMPLETED.value
+        assert ex.successful_files == 1
+        assert ex.result_acknowledged is True
 class RetrieveNotFoundTests(TestCase):
     """A missing execution must return 404, not 500 (UN-3719). The reaper's
     orphan-claim sweep relies on the deterministic 404 to GC claims for deleted

backend/workflow_manager/workflow_v2/execution.py

@@ -3,7 +3,6 @@
 import uuid
 
 from account_v2.constants import Common
-from django.utils import timezone
 from platform_settings_v2.platform_auth_service import PlatformAuthenticationService
 from tags.models import Tag
 from tool_instance_v2.models import ToolInstance
@@ -174,29 +173,13 @@ def update_execution(
         error: str | None = None,
         increment_attempt: bool = False,
     ) -> None:
+        # Delegate to the model method, which owns the single, robust terminal-one-way
+        # guard (atomic select_for_update, PG-scoped, field-scoped writes). Duplicating
+        # the guard here risked drift + its own TOCTOU/dropped-field bugs.
         execution = WorkflowExecution.objects.get(pk=self.execution_id)
-
-        if status is not None:
-            execution.status = status.value
-
-            if (
-                status
-                in [
-                    ExecutionStatus.COMPLETED,
-                    ExecutionStatus.ERROR,
-                    ExecutionStatus.STOPPED,
-                ]
-                and not execution.execution_time
-            ):
-                execution.execution_time = round(
-                    (timezone.now() - execution.created_at).total_seconds(), 3
-                )
-        if error:
-            execution.error_message = error[:EXECUTION_ERROR_LENGTH]
-        if increment_attempt:
-            execution.attempts += 1
-
-        execution.save()
+        execution.update_execution(
+            status=status, error=error, increment_attempt=increment_attempt
+        )
 
     def has_successful_compilation(self) -> bool:
         return self.compilation_result["success"] is True
@@ -396,9 +379,10 @@ def initiate_tool_execution(
     def update_execution_err(execution_id: str, err_msg: str = "") -> WorkflowExecution:
         try:
             execution = WorkflowExecution.objects.get(pk=execution_id)
-            execution.status = ExecutionStatus.ERROR.value
-            execution.error_message = err_msg[:EXECUTION_ERROR_LENGTH]
-            execution.save()
+            # Route through the guarded model method instead of a direct save, so a
+            # late error handler can't revert a PG execution the callback already
+            # finalized COMPLETED/STOPPED. (ERROR over EXECUTING/PENDING still applies.)
+            execution.update_execution(status=ExecutionStatus.ERROR, error=err_msg)
             return execution
         except WorkflowExecution.DoesNotExist:
             logger.error(f"execution doesn't exist {execution_id}")

backend/workflow_manager/workflow_v2/models/execution.py

@@ -369,8 +369,31 @@
             error (Optional[str], optional): Error message if any. Defaults to None.
             increment_attempt (bool, optional): Whether to increment attempt counter. Defaults to False.
         """
-        should_release_rate_limit = False
+        # Route on the PERSISTED transport marker, never the (possibly stale) in-memory
+        # self.queue_message_id: a PG row snapshotted before dispatch recorded its
+        # handle would otherwise be mis-routed to the legacy path, whose full save()
+        # would revert the row AND write NULL back into queue_message_id — permanently
+        # disabling this and every sibling guard (all gate on queue_message_id).
+        committed_qmid = (
+            WorkflowExecution.objects.filter(pk=self.pk)
+            .values_list("queue_message_id", flat=True)
+            .first()
+        )
+        if committed_qmid is None:
+            self._update_execution_legacy(status, error, increment_attempt)
+        else:
+            self._update_execution_pg_guarded(status, error, increment_attempt)
 
+    def _update_execution_legacy(
+        self,
+        status: ExecutionStatus | None,
+        error: str | None,
+        increment_attempt: bool,
+    ) -> None:
+        """Celery-path (queue_message_id NULL) status update — unchanged legacy
+        behavior, a full save().
+        """
+        should_release_rate_limit = False
         if status is not None:
             status = ExecutionStatus(status)
             self.status = status.value
@@ -381,13 +404,89 @@
             ]:
                 self.execution_time = CommonUtils.time_since(self.created_at, 3)
                 should_release_rate_limit = True
-
         if error:
             self.error_message = error[:EXECUTION_ERROR_LENGTH]
         if increment_attempt:
             self.attempts += 1
-
         self.save()
+        if should_release_rate_limit and self.pipeline_id:
+            self._release_api_deployment_rate_limit()
+
+    def _update_execution_pg_guarded(
+        self,
+        status: ExecutionStatus | None,
+        error: str | None,
+        increment_attempt: bool,
+    ) -> None:
+        """PG-path status update with the terminal-one-way guard.
+
+        Locks the row (``select_for_update``) so the committed-status read and the
+        write are atomic — a concurrent callback can't commit COMPLETED in between.
+        A protected-terminal (COMPLETED/STOPPED) row refuses only the *status* change
+        (ERROR stays correctable to COMPLETED); ``error`` / ``increment_attempt`` are
+        still applied. Writes are field-scoped (``update_fields``) so a stale object's
+        NULL counters — or the marker — can never be clobbered by a full-row save, and
+        the post-save cache write publishes the status we actually persisted.
+        """
+        from django.db import transaction
+
+        should_release_rate_limit = False
+        update_fields: list[str] = []
+        with transaction.atomic():
+            locked = (
+                WorkflowExecution.objects.select_for_update()
+                .filter(pk=self.pk)
+                .values("status", "attempts")
+                .first()
+            )
+            if locked is None:
+                return
+            committed = locked["status"]
+
+            if status is not None:
+                status = ExecutionStatus(status)
+                if (
+                    committed in ExecutionStatus.protected_terminal_values()
+                    and status.value != committed
+                ):
+                    logger.warning(
+                        "update_execution: refusing to revert protected status '%s' "
+                        "with '%s' for execution %s (stale-writer guard); error=%s "
+                        "increment_attempt=%s still applied",
+                        committed,
+                        status.value,
+                        self.id,
+                        bool(error),
+                        increment_attempt,
+                    )
+                    # Keep self (and the post-save cache) consistent with the DB.
+                    self.status = committed
+                else:
+                    self.status = status.value
+                    update_fields.append("status")
+                    if status in [
+                        ExecutionStatus.COMPLETED,
+                        ExecutionStatus.ERROR,
+                        ExecutionStatus.STOPPED,
+                    ]:
+                        self.execution_time = CommonUtils.time_since(self.created_at, 3)
+                        update_fields.append("execution_time")
+                        should_release_rate_limit = True
+
+            if error:
+                self.error_message = error[:EXECUTION_ERROR_LENGTH]
+                update_fields.append("error_message")
+            if increment_attempt:
+                # Increment from the locked committed value, not the stale in-memory one.
+                self.attempts = locked["attempts"] + 1
+                update_fields.append("attempts")
+
+            if update_fields:
+                update_fields.append("modified_at")
+                self.save(update_fields=update_fields)
+
+        if should_release_rate_limit and self.pipeline_id:
+            self._release_api_deployment_rate_limit()
 
         # Release rate limit slot for API deployment executions after save
         if should_release_rate_limit and self.pipeline_id:

backend/workflow_manager/workflow_v2/workflow_helper.py

@@ -435,8 +435,16 @@ def _set_result_acknowledge(execution: WorkflowExecution) -> None:
             execution (WorkflowExecution): WorkflowExecution instance
         """
         if not execution.result_acknowledged:
+            # A DB-side queryset UPDATE of the single column, NOT execution.save():
+            # a save() (even with update_fields) rewrites the whole row from this
+            # possibly-stale object AND re-runs _handle_execution_cache(), which would
+            # publish the stale in-memory status (e.g. EXECUTING+NULL) to Redis while
+            # Postgres stays COMPLETED. Acknowledging a result must touch ONLY that
+            # column and never the status/counters or the cache.
+            WorkflowExecution.objects.filter(pk=execution.pk).update(
+                result_acknowledged=True
+            )
             execution.result_acknowledged = True
-            execution.save()
             logger.info(
                 f"ExecutionID [{execution.id}] - Task {execution.task_id} acknowledged"
             )