[cross-repo from server#267] Conformance: replay approve signal does not complete after activity state query (#95)

Author: durable-workflow-ops

CHANGELOG.md

@@ -7,10 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Fixed
-- Workflow-task completion transport failures now cause the worker to report
-  the task as failed instead of treating the local replay as completed. This
-  lets the server re-dispatch quickly instead of leaving a workflow task leased
-  until timeout when the completion request never records.
+- Ambiguous workflow-task completion failures no longer get reported back as
+  durable task failures after commands have been produced. Definite server
+  rejections are still treated as failed workflow tasks even when the
+  best-effort failure report cannot be sent, but transport/ownership ambiguity
+  preserves the emitted commands so replay-driven signal completion is not
+  converted into a stuck failed task.
 - Repeated condition-wait openings for the same logical wait now replay through
   every matching signal before deciding whether the wait is still pending, so
   long-running signal/query workflows do not get stuck on the first signal.

src/durable_workflow/worker.py

@@ -45,7 +45,15 @@
     Client,
     WorkflowExecution,
 )
-from .errors import ActivityCancelled, AvroNotInstalledError, NonRetryableError, QueryFailed, ServerError
+from .errors import (
+    ActivityCancelled,
+    AvroNotInstalledError,
+    DurableWorkflowError,
+    InvalidArgument,
+    NonRetryableError,
+    QueryFailed,
+    ServerError,
+)
 from .external_storage import ExternalPayloadCache, ExternalStorageDriver
 from .interceptors import (
     ActivityInterceptorContext,
@@ -74,6 +82,15 @@
     "query_task_not_leased",
     "query_task_timed_out",
 }
+_WORKFLOW_TASK_COMPLETION_AMBIGUOUS_REJECTION_REASONS = {
+    "lease_expired",
+    "lease_owner_mismatch",
+    "run_already_closed",
+    "run_closed",
+    "task_not_found",
+    "task_not_leased",
+    "workflow_task_attempt_mismatch",
+}
 _WORKER_WORKFLOW_FINGERPRINTS: dict[tuple[str, str], str] = {}
 
 
@@ -111,6 +128,19 @@ def _is_final_query_task_rejection(error: BaseException) -> bool:
     )
 
 
+def _should_fail_workflow_task_after_completion_error(error: BaseException) -> bool:
+    if isinstance(error, InvalidArgument):
+        return True
+
+    if isinstance(error, ServerError):
+        if error.status >= 500 or error.status == 429:
+            return False
+
+        return error.reason() not in _WORKFLOW_TASK_COMPLETION_AMBIGUOUS_REJECTION_REASONS
+
+    return isinstance(error, DurableWorkflowError)
+
+
 def _signal_arguments_envelope_from_export(
     signal: Mapping[str, Any],
     *,
@@ -775,8 +805,9 @@ async def _run_workflow_task_core(self, task: dict[str, Any]) -> list[dict[str,
                 )
             except Exception as e:
                 log.warning("failed to complete workflow update task %s: %s", task_id, e)
-                await self._fail_workflow_task_after_completion_error(task_id, attempt, e)
-                return None
+                if _should_fail_workflow_task_after_completion_error(e):
+                    await self._report_workflow_task_after_completion_error(task_id, attempt, e)
+                    return None
             return [command]
 
         try:
@@ -849,11 +880,12 @@ async def _run_workflow_task_core(self, task: dict[str, Any]) -> list[dict[str,
             )
         except Exception as e:
             log.warning("failed to complete workflow task %s: %s", task_id, e)
-            await self._fail_workflow_task_after_completion_error(task_id, attempt, e)
-            return None
+            if _should_fail_workflow_task_after_completion_error(e):
+                await self._report_workflow_task_after_completion_error(task_id, attempt, e)
+                return None
         return commands
 
-    async def _fail_workflow_task_after_completion_error(
+    async def _report_workflow_task_after_completion_error(
         self,
         task_id: str,
         attempt: int,

tests/test_worker.py

@@ -22,7 +22,7 @@
     Client,
     WorkflowExecution,
 )
-from durable_workflow.errors import ServerError
+from durable_workflow.errors import InvalidArgument, ServerError, Unauthorized, WorkflowNotFound
 from durable_workflow.interceptors import (
     ActivityHandler,
     ActivityInterceptorContext,
@@ -32,7 +32,7 @@
     WorkflowTaskHandler,
     WorkflowTaskInterceptorContext,
 )
-from durable_workflow.worker import Worker
+from durable_workflow.worker import Worker, _should_fail_workflow_task_after_completion_error
 
 
 @workflow.defn(name="test-wf")
@@ -185,6 +185,28 @@ def compatible_cluster_info(**overrides: object) -> dict[str, object]:
     return info
 
 
+class TestWorkflowTaskCompletionErrorClassification:
+    @pytest.mark.parametrize(
+        ("error", "should_fail"),
+        [
+            (TimeoutError("completion timed out"), False),
+            (RuntimeError("connection reset"), False),
+            (ServerError(409, {"reason": "lease_expired"}), False),
+            (ServerError(409, {"reason": "workflow_task_attempt_mismatch"}), False),
+            (ServerError(429, {"reason": "rate_limited"}), False),
+            (ServerError(503, {"reason": "server_busy"}), False),
+            (ServerError(409, {"reason": "invalid_commands"}), True),
+            (InvalidArgument("invalid command payload"), True),
+            (Unauthorized("missing bearer token"), True),
+            (WorkflowNotFound("wf-missing"), True),
+        ],
+    )
+    def test_classifies_definite_and_ambiguous_completion_errors(
+        self, error: BaseException, should_fail: bool
+    ) -> None:
+        assert _should_fail_workflow_task_after_completion_error(error) is should_fail
+
+
 class TestWorkerRegistration:
     @pytest.mark.asyncio
     async def test_register(self, mock_client: AsyncMock) -> None:
@@ -502,7 +524,7 @@ async def test_schedule_activity_on_first_replay(self, mock_client: AsyncMock) -
         assert serializer.decode(commands[0]["arguments"]["blob"], codec="json") == ["hello"]
 
     @pytest.mark.asyncio
-    async def test_workflow_task_completion_error_fails_task_for_fast_redispatch(
+    async def test_workflow_task_ambiguous_completion_error_preserves_commands(
         self, mock_client: AsyncMock
     ) -> None:
         mock_client.complete_workflow_task.side_effect = TimeoutError("completion timed out")
@@ -518,15 +540,96 @@ async def test_workflow_task_completion_error_fails_task_for_fast_redispatch(
 
         result = await worker._run_workflow_task(task)
 
+        assert result is not None
+        assert result[0]["type"] == "schedule_activity"
+        mock_client.complete_workflow_task.assert_awaited_once()
+        mock_client.fail_workflow_task.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_workflow_task_definite_completion_rejection_fails_task(
+        self, mock_client: AsyncMock
+    ) -> None:
+        mock_client.complete_workflow_task.side_effect = ServerError(409, {"reason": "invalid_commands"})
+        worker = Worker(mock_client, task_queue="q1", workflows=[TestWorkflow], activities=[])
+        task = {
+            "task_id": "t-complete-invalid",
+            "workflow_type": "test-wf",
+            "workflow_task_attempt": 2,
+            "history_events": [],
+            "arguments": '["hello"]',
+            "payload_codec": "json",
+        }
+
+        result = await worker._run_workflow_task(task)
+
+        assert result is None
+        mock_client.fail_workflow_task.assert_awaited_once()
+        call_kwargs = mock_client.fail_workflow_task.await_args.kwargs
+        assert call_kwargs["task_id"] == "t-complete-invalid"
+        assert call_kwargs["workflow_task_attempt"] == 2
+        assert call_kwargs["lease_owner"] == worker.worker_id
+        assert call_kwargs["failure_type"] == "ServerError"
+        assert "invalid_commands" in call_kwargs["message"]
+
+    @pytest.mark.parametrize(
+        ("completion_error", "failure_type", "message_fragment"),
+        [
+            (Unauthorized("missing bearer token"), "Unauthorized", "missing bearer token"),
+            (WorkflowNotFound("wf-typed-missing"), "WorkflowNotFound", "wf-typed-missing"),
+        ],
+    )
+    @pytest.mark.asyncio
+    async def test_workflow_task_typed_completion_rejection_fails_task(
+        self,
+        mock_client: AsyncMock,
+        completion_error: Exception,
+        failure_type: str,
+        message_fragment: str,
+    ) -> None:
+        mock_client.complete_workflow_task.side_effect = completion_error
+        worker = Worker(mock_client, task_queue="q1", workflows=[TestWorkflow], activities=[])
+        task = {
+            "task_id": "t-complete-typed-rejection",
+            "workflow_type": "test-wf",
+            "workflow_task_attempt": 2,
+            "history_events": [],
+            "arguments": '["hello"]',
+            "payload_codec": "json",
+        }
+
+        result = await worker._run_workflow_task(task)
+
         assert result is None
         mock_client.complete_workflow_task.assert_awaited_once()
         mock_client.fail_workflow_task.assert_awaited_once()
         call_kwargs = mock_client.fail_workflow_task.await_args.kwargs
-        assert call_kwargs["task_id"] == "t-complete-timeout"
+        assert call_kwargs["task_id"] == "t-complete-typed-rejection"
         assert call_kwargs["workflow_task_attempt"] == 2
         assert call_kwargs["lease_owner"] == worker.worker_id
-        assert call_kwargs["failure_type"] == "TimeoutError"
-        assert "completion timed out" in call_kwargs["message"]
+        assert call_kwargs["failure_type"] == failure_type
+        assert message_fragment in call_kwargs["message"]
+
+    @pytest.mark.asyncio
+    async def test_workflow_task_definite_completion_rejection_stays_failed_when_report_fails(
+        self, mock_client: AsyncMock
+    ) -> None:
+        mock_client.complete_workflow_task.side_effect = ServerError(409, {"reason": "invalid_commands"})
+        mock_client.fail_workflow_task.side_effect = RuntimeError("failure report unavailable")
+        worker = Worker(mock_client, task_queue="q1", workflows=[TestWorkflow], activities=[])
+        task = {
+            "task_id": "t-complete-invalid-report-fails",
+            "workflow_type": "test-wf",
+            "workflow_task_attempt": 2,
+            "history_events": [],
+            "arguments": '["hello"]',
+            "payload_codec": "json",
+        }
+
+        result = await worker._run_workflow_task(task)
+
+        assert result is None
+        mock_client.complete_workflow_task.assert_awaited_once()
+        mock_client.fail_workflow_task.assert_awaited_once()
 
     @pytest.mark.asyncio
     async def test_workflow_command_payload_warning_uses_client_policy(
@@ -701,7 +804,7 @@ async def test_update_backed_workflow_task_completes_update_command(
         mock_client.fail_workflow_task.assert_not_called()
 
     @pytest.mark.asyncio
-    async def test_update_task_completion_error_fails_task_for_fast_redispatch(
+    async def test_update_task_ambiguous_completion_error_preserves_command(
         self, mock_client: AsyncMock
     ) -> None:
         mock_client.complete_workflow_task.side_effect = TimeoutError("update completion timed out")
@@ -729,14 +832,10 @@ async def test_update_task_completion_error_fails_task_for_fast_redispatch(
 
         result = await worker._run_workflow_task(task)
 
-        assert result is None
+        assert result is not None
+        assert result[0]["type"] == "complete_update"
         mock_client.complete_workflow_task.assert_awaited_once()
-        mock_client.fail_workflow_task.assert_awaited_once()
-        call_kwargs = mock_client.fail_workflow_task.await_args.kwargs
-        assert call_kwargs["task_id"] == "t-update-timeout"
-        assert call_kwargs["workflow_task_attempt"] == 3
-        assert call_kwargs["failure_type"] == "TimeoutError"
-        assert "update completion timed out" in call_kwargs["message"]
+        mock_client.fail_workflow_task.assert_not_called()
 
     @pytest.mark.asyncio
     async def test_query_task_executes_registered_query(self, mock_client: AsyncMock) -> None: