diff --git a/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py b/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py index 0e5183a71b..44a19b888c 100644 --- a/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py +++ b/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 diff --git a/backend/workflow_manager/workflow_v2/execution.py b/backend/workflow_manager/workflow_v2/execution.py index 83a1c24057..e920468038 100644 --- a/backend/workflow_manager/workflow_v2/execution.py +++ b/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}") diff --git a/backend/workflow_manager/workflow_v2/models/execution.py b/backend/workflow_manager/workflow_v2/models/execution.py index 31498a3254..c764b5cd9f 100644 --- a/backend/workflow_manager/workflow_v2/models/execution.py +++ b/backend/workflow_manager/workflow_v2/models/execution.py @@ -369,8 +369,45 @@ def update_execution( error (Optional[str], optional): Error message if any. Defaults to None. increment_attempt (bool, optional): Whether to increment attempt counter. Defaults to False. """ + from django.db import transaction + should_release_rate_limit = False + with transaction.atomic(): + # Lock the row and read the PERSISTED marker + status together, so routing + # AND the write are atomic. Never trust the (possibly stale) in-memory + # self.queue_message_id: a PG row snapshotted before dispatch recorded its + # handle, mis-routed to an unlocked legacy full save(), would write NULL + # back into queue_message_id and permanently disable every guard. Under + # this lock a concurrent _record_dispatch_handle serializes after us. + locked = ( + WorkflowExecution.objects.select_for_update() + .filter(pk=self.pk) + .values("queue_message_id", "status", "attempts") + .first() + ) + if locked is None: + return + if locked["queue_message_id"] is None: + should_release_rate_limit = self._apply_legacy_update( + status, error, increment_attempt + ) + else: + should_release_rate_limit = self._apply_pg_guarded_update( + locked, status, error, increment_attempt + ) + if should_release_rate_limit and self.pipeline_id: + self._release_api_deployment_rate_limit() + def _apply_legacy_update( + self, + status: ExecutionStatus | None, + error: str | None, + increment_attempt: bool, + ) -> bool: + """Celery-path (queue_message_id NULL) — unchanged full save(); runs inside the + caller's locked transaction. Returns whether to release the rate-limit slot. + """ + should_release = False if status is not None: status = ExecutionStatus(status) self.status = status.value @@ -380,18 +417,90 @@ def update_execution( ExecutionStatus.STOPPED, ]: self.execution_time = CommonUtils.time_since(self.created_at, 3) - should_release_rate_limit = True - + should_release = True if error: self.error_message = error[:EXECUTION_ERROR_LENGTH] if increment_attempt: self.attempts += 1 - self.save() + return should_release - # Release rate limit slot for API deployment executions after save - if should_release_rate_limit and self.pipeline_id: - self._release_api_deployment_rate_limit() + def _apply_pg_guarded_update( + self, + locked: dict, + status: ExecutionStatus | None, + error: str | None, + increment_attempt: bool, + ) -> bool: + """PG-path terminal-one-way guarded write, inside the caller's row lock. Refuses + a revert of a protected-terminal (COMPLETED/STOPPED) status — ERROR stays + correctable — while still applying error / increment_attempt. Field-scoped + (``update_fields``) so a stale object can never clobber counters or the marker, + and the post-save cache publishes the status actually persisted. Returns whether + to release the rate-limit slot. + """ + committed = locked["status"] + should_release = False + update_fields: list[str] = [] + if status is not None: + status_fields, should_release = self._apply_guarded_status( + ExecutionStatus(status), committed, error, increment_attempt + ) + update_fields.extend(status_fields) + 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) + return should_release + + def _apply_guarded_status( + self, + status: ExecutionStatus, + committed: str, + error: str | None, + increment_attempt: bool, + ) -> tuple[list[str], bool]: + """Apply the status change under the terminal-one-way guard. + + Returns ``(update_fields, should_release_rate_limit)``. On a refused revert of + a protected-terminal (COMPLETED/STOPPED) row, keeps ``self.status == committed`` + and returns ``([], False)`` — the caller still applies error / increment. + """ + 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, + ) + self.status = committed # keep self + post-save cache consistent with DB + return [], False + + self.status = status.value + fields = ["status"] + should_release = False + if status in [ + ExecutionStatus.COMPLETED, + ExecutionStatus.ERROR, + ExecutionStatus.STOPPED, + ]: + self.execution_time = CommonUtils.time_since(self.created_at, 3) + fields.append("execution_time") + should_release = True + return fields, should_release def _release_api_deployment_rate_limit(self) -> None: """Release rate limit slot for API deployment executions. diff --git a/backend/workflow_manager/workflow_v2/workflow_helper.py b/backend/workflow_manager/workflow_v2/workflow_helper.py index 0e2e28aacd..d22dcc85bb 100644 --- a/backend/workflow_manager/workflow_v2/workflow_helper.py +++ b/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" ) diff --git a/unstract/core/src/unstract/core/data_models.py b/unstract/core/src/unstract/core/data_models.py index 8e938141fe..7b856d70ca 100644 --- a/unstract/core/src/unstract/core/data_models.py +++ b/unstract/core/src/unstract/core/data_models.py @@ -586,6 +586,18 @@ def _terminal_values(cls) -> frozenset[str]: ExecutionStatus.terminal_values = classmethod(_terminal_values) +# Terminal statuses that a stale/late writer must NOT overwrite (the terminal-one-way +# guard). Derived from terminal_values() minus ERROR, so a *new* terminal status is +# protected-by-default; ERROR is deliberately excluded so a premature ERROR (set by +# an upstream/other path before the real callback) stays correctable to COMPLETED. +def _protected_terminal_values(cls) -> frozenset[str]: + """String values of terminal statuses that are one-way (COMPLETED/STOPPED).""" + return cls.terminal_values() - {cls.ERROR.value} + + +ExecutionStatus.protected_terminal_values = classmethod(_protected_terminal_values) + + # Add the is_completed method as a class method def _is_completed(cls, status: str) -> bool: """Check if the execution status represents a completed (terminal) state."""