Conformance: Python signal/query workflow stays leased before durable wait (#89)

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 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.
 - 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

@@ -775,6 +775,8 @@ 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
             return [command]
 
         try:
@@ -847,8 +849,32 @@ 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
         return commands
 
+    async def _fail_workflow_task_after_completion_error(
+        self,
+        task_id: str,
+        attempt: int,
+        error: Exception,
+    ) -> None:
+        try:
+            await self.client.fail_workflow_task(
+                task_id=task_id,
+                lease_owner=self.worker_id,
+                workflow_task_attempt=attempt,
+                message=f"workflow task completion failed after commands were produced: {error}",
+                failure_type=type(error).__name__,
+                stack_trace=traceback.format_exc(),
+            )
+        except Exception as fail_error:
+            log.warning(
+                "failed to report workflow task %s completion failure: %s",
+                task_id,
+                fail_error,
+            )
+
     async def _run_activity_task(self, task: dict[str, Any]) -> str:
         task_id: str = task["task_id"]
         attempt_id: str = task.get("activity_attempt_id") or task.get("attempt_id", "")

tests/test_worker.py

@@ -501,6 +501,33 @@ async def test_schedule_activity_on_first_replay(self, mock_client: AsyncMock) -
         assert commands[0]["arguments"]["codec"] == "json"
         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(
+        self, mock_client: AsyncMock
+    ) -> None:
+        mock_client.complete_workflow_task.side_effect = TimeoutError("completion timed out")
+        worker = Worker(mock_client, task_queue="q1", workflows=[TestWorkflow], activities=[])
+        task = {
+            "task_id": "t-complete-timeout",
+            "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["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"]
+
     @pytest.mark.asyncio
     async def test_workflow_command_payload_warning_uses_client_policy(
         self, mock_client: AsyncMock, caplog: pytest.LogCaptureFixture
@@ -673,6 +700,44 @@ 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(
+        self, mock_client: AsyncMock
+    ) -> None:
+        mock_client.complete_workflow_task.side_effect = TimeoutError("update completion timed out")
+        worker = Worker(mock_client, task_queue="q1", workflows=[UpdateWorkflow], activities=[])
+        task = {
+            "task_id": "t-update-timeout",
+            "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 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-update-timeout"
+        assert call_kwargs["workflow_task_attempt"] == 3
+        assert call_kwargs["failure_type"] == "TimeoutError"
+        assert "update completion timed out" in call_kwargs["message"]
+
     @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=[])