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

Author: durable-workflow-ops

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Fixed
+- Workflow-task completion now retries transient transport failures and server
+  throttling/5xx rejections before preserving emitted commands or reporting a
+  definite task failure, reducing stuck waiting runs when a signal-satisfied
+  wait completes immediately after replay-driven query activity.
 - 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

src/durable_workflow/worker.py

@@ -91,6 +91,8 @@
     "task_not_leased",
     "workflow_task_attempt_mismatch",
 }
+_WORKFLOW_TASK_COMPLETION_MAX_ATTEMPTS = 3
+_WORKFLOW_TASK_COMPLETION_RETRY_DELAYS = (0.05, 0.2)
 _WORKER_WORKFLOW_FINGERPRINTS: dict[tuple[str, str], str] = {}
 
 
@@ -141,6 +143,13 @@ def _should_fail_workflow_task_after_completion_error(error: BaseException) -> b
     return isinstance(error, DurableWorkflowError)
 
 
+def _should_retry_workflow_task_completion_error(error: BaseException) -> bool:
+    if isinstance(error, ServerError):
+        return error.status >= 500 or error.status == 429
+
+    return not isinstance(error, DurableWorkflowError)
+
+
 def _signal_arguments_envelope_from_export(
     signal: Mapping[str, Any],
     *,
@@ -797,10 +806,9 @@ async def _run_workflow_task_core(self, task: dict[str, Any]) -> list[dict[str,
                 command["type"],
             )
             try:
-                await self.client.complete_workflow_task(
+                await self._complete_workflow_task_with_retry(
                     task_id=task_id,
-                    lease_owner=self.worker_id,
-                    workflow_task_attempt=attempt,
+                    attempt=attempt,
                     commands=[command],
                 )
             except Exception as e:
@@ -872,10 +880,9 @@ async def _run_workflow_task_core(self, task: dict[str, Any]) -> list[dict[str,
             [c["type"] for c in commands],
         )
         try:
-            await self.client.complete_workflow_task(
+            await self._complete_workflow_task_with_retry(
                 task_id=task_id,
-                lease_owner=self.worker_id,
-                workflow_task_attempt=attempt,
+                attempt=attempt,
                 commands=commands,
             )
         except Exception as e:
@@ -885,6 +892,42 @@ async def _run_workflow_task_core(self, task: dict[str, Any]) -> list[dict[str,
                 return None
         return commands
 
+    async def _complete_workflow_task_with_retry(
+        self,
+        *,
+        task_id: str,
+        attempt: int,
+        commands: list[dict[str, Any]],
+    ) -> None:
+        for completion_attempt in range(1, _WORKFLOW_TASK_COMPLETION_MAX_ATTEMPTS + 1):
+            try:
+                await self.client.complete_workflow_task(
+                    task_id=task_id,
+                    lease_owner=self.worker_id,
+                    workflow_task_attempt=attempt,
+                    commands=commands,
+                )
+                return
+            except Exception as error:
+                if (
+                    completion_attempt >= _WORKFLOW_TASK_COMPLETION_MAX_ATTEMPTS
+                    or not _should_retry_workflow_task_completion_error(error)
+                ):
+                    raise
+
+                delay_index = min(
+                    completion_attempt - 1,
+                    len(_WORKFLOW_TASK_COMPLETION_RETRY_DELAYS) - 1,
+                )
+                log.warning(
+                    "retrying workflow task %s completion after attempt %d/%d failed: %s",
+                    task_id,
+                    completion_attempt,
+                    _WORKFLOW_TASK_COMPLETION_MAX_ATTEMPTS,
+                    error,
+                )
+                await asyncio.sleep(_WORKFLOW_TASK_COMPLETION_RETRY_DELAYS[delay_index])
+
     async def _report_workflow_task_after_completion_error(
         self,
         task_id: str,

tests/test_worker.py

@@ -527,10 +527,10 @@ async def test_schedule_activity_on_first_replay(self, mock_client: AsyncMock) -
     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")
+        mock_client.complete_workflow_task.side_effect = ServerError(409, {"reason": "task_not_leased"})
         worker = Worker(mock_client, task_queue="q1", workflows=[TestWorkflow], activities=[])
         task = {
-            "task_id": "t-complete-timeout",
+            "task_id": "t-complete-not-leased",
             "workflow_type": "test-wf",
             "workflow_task_attempt": 2,
             "history_events": [],
@@ -545,6 +545,31 @@ async def test_workflow_task_ambiguous_completion_error_preserves_commands(
         mock_client.complete_workflow_task.assert_awaited_once()
         mock_client.fail_workflow_task.assert_not_called()
 
+    @pytest.mark.asyncio
+    async def test_workflow_task_retries_transient_completion_error(
+        self, mock_client: AsyncMock
+    ) -> None:
+        mock_client.complete_workflow_task.side_effect = [
+            TimeoutError("completion timed out"),
+            {"outcome": "completed"},
+        ]
+        worker = Worker(mock_client, task_queue="q1", workflows=[TestWorkflow], activities=[])
+        task = {
+            "task_id": "t-complete-retry",
+            "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 not None
+        assert result[0]["type"] == "schedule_activity"
+        assert mock_client.complete_workflow_task.await_count == 2
+        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
@@ -807,10 +832,10 @@ async def test_update_backed_workflow_task_completes_update_command(
     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")
+        mock_client.complete_workflow_task.side_effect = ServerError(409, {"reason": "task_not_leased"})
         worker = Worker(mock_client, task_queue="q1", workflows=[UpdateWorkflow], activities=[])
         task = {
-            "task_id": "t-update-timeout",
+            "task_id": "t-update-not-leased",
             "workflow_type": "update-wf",
             "workflow_task_attempt": 3,
             "workflow_update_id": "upd-worker-1",
@@ -837,6 +862,43 @@ async def test_update_task_ambiguous_completion_error_preserves_command(
         mock_client.complete_workflow_task.assert_awaited_once()
         mock_client.fail_workflow_task.assert_not_called()
 
+    @pytest.mark.asyncio
+    async def test_update_task_retries_transient_completion_error(
+        self, mock_client: AsyncMock
+    ) -> None:
+        mock_client.complete_workflow_task.side_effect = [
+            ServerError(503, {"reason": "server_busy"}),
+            {"outcome": "completed"},
+        ]
+        worker = Worker(mock_client, task_queue="q1", workflows=[UpdateWorkflow], activities=[])
+        task = {
+            "task_id": "t-update-retry",
+            "workflow_type": "update-wf",
+            "workflow_task_attempt": 3,
+            "workflow_update_id": "upd-worker-1",
+            "workflow_wait_kind": "update",
+            "history_events": [
+                {
+                    "event_type": "UpdateAccepted",
+                    "payload": {
+                        "update_id": "upd-worker-1",
+                        "update_name": "increment",
+                        "arguments": serializer.encode([6], codec="json"),
+                        "payload_codec": "json",
+                    },
+                },
+            ],
+            "arguments": "[]",
+            "payload_codec": "json",
+        }
+
+        result = await worker._run_workflow_task(task)
+
+        assert result is not None
+        assert result[0]["type"] == "complete_update"
+        assert mock_client.complete_workflow_task.await_count == 2
+        mock_client.fail_workflow_task.assert_not_called()
+
     @pytest.mark.asyncio
     async def test_query_task_executes_registered_query(self, mock_client: AsyncMock) -> None:
         worker = Worker(mock_client, task_queue="q1", workflows=[QueryWorkflow], activities=[])