[cross-repo from server#289] Conformance finding: replay state query times out on current artifacts (#104)

Author: durable-workflow-ops

src/durable_workflow/errors.py

@@ -122,6 +122,11 @@ def __init__(self, schedule_id: str) -> None:
 class QueryFailed(DurableWorkflowError):
     """A workflow query was rejected or the workflow raised while handling it."""
 
+    def __init__(self, message: str, *, reason: str | None = None, body: object | None = None) -> None:
+        super().__init__(message)
+        self.reason = reason
+        self.body = body
+
 
 class WorkflowPayloadDecodeError(DurableWorkflowError):
     """A committed workflow history payload could not be decoded during replay."""
@@ -281,9 +286,12 @@ def _raise_for_status(status: int, body: object, *, context: str = "") -> None:
     if status == 401:
         raise Unauthorized(message or "unauthorized")
 
+    def query_failed(default: str) -> QueryFailed:
+        return QueryFailed(message or default, reason=reason if isinstance(reason, str) else None, body=body)
+
     if status == 404:
         if reason in ("query_not_found", "rejected_unknown_query"):
-            raise QueryFailed(message or "query not found")
+            raise query_failed("query not found")
         if reason == "schedule_not_found":
             raise ScheduleNotFound(context)
         if reason in ("instance_not_found", "workflow_not_found") or "workflow" in context.lower():
@@ -297,15 +305,24 @@ def _raise_for_status(status: int, body: object, *, context: str = "") -> None:
             raise ScheduleAlreadyExists(context)
         if reason == "duplicate_not_allowed":
             raise WorkflowAlreadyStarted(context)
-        if reason in ("query_rejected", "query_worker_unavailable"):
-            raise QueryFailed(message or "query rejected")
+        if reason in (
+            "query_rejected",
+            "query_worker_unavailable",
+            "query_worker_incompatible",
+            "query_workflow_state_unavailable",
+        ):
+            raise query_failed("query rejected")
         if reason == "update_rejected":
             raise UpdateRejected(message or "update rejected")
         raise ServerError(status, body)
 
     if status == 504:
-        if reason == "query_worker_timeout":
-            raise QueryFailed(message or "query worker timed out")
+        if reason in (
+            "query_worker_timeout",
+            "query_worker_execution_timeout",
+            "query_task_not_claimed",
+        ):
+            raise query_failed("query worker timed out")
         raise ServerError(status, body)
 
     if status == 422:

src/durable_workflow/worker.py

@@ -1245,11 +1245,17 @@ async def _run_query_task_core(self, task: dict[str, Any], *, client: Client | N
             )
             return "failed"
         except QueryFailed as e:
-            reason = "rejected_unknown_query" if "unknown query" in str(e) else "query_rejected"
+            message = str(e)
+            if "unknown query" in message:
+                reason = "rejected_unknown_query"
+            elif message.startswith("workflow replay failed before query:"):
+                reason = "query_workflow_state_unavailable"
+            else:
+                reason = "query_rejected"
             await self._fail_query_task(
                 query_task_id,
                 attempt,
-                str(e),
+                message,
                 reason=reason,
                 failure_type=type(e).__name__,
                 stack_trace=traceback.format_exc(),

tests/test_client.py

@@ -842,19 +842,91 @@ async def test_worker_routed_query_unavailable_raises_query_failed(self, client:
         resp = _mock_response(409, {"reason": "query_worker_unavailable", "message": "no worker"})
         with (
             patch.object(client._http, "request", new_callable=AsyncMock, return_value=resp),
-            pytest.raises(QueryFailed),
+            pytest.raises(QueryFailed) as excinfo,
+        ):
+            await client.query_workflow("wf-1", "status")
+
+        assert excinfo.value.reason == "query_worker_unavailable"
+
+    @pytest.mark.asyncio
+    async def test_worker_routed_query_incompatible_raises_query_failed(self, client: Client) -> None:
+        resp = _mock_response(
+            409,
+            {
+                "reason": "query_worker_incompatible",
+                "message": "no compatible worker supports this workflow type",
+            },
+        )
+        with (
+            patch.object(client._http, "request", new_callable=AsyncMock, return_value=resp),
+            pytest.raises(QueryFailed) as excinfo,
         ):
             await client.query_workflow("wf-1", "status")
 
+        assert excinfo.value.reason == "query_worker_incompatible"
+
+    @pytest.mark.asyncio
+    async def test_worker_routed_query_state_unavailable_raises_query_failed(self, client: Client) -> None:
+        resp = _mock_response(
+            409,
+            {
+                "reason": "query_workflow_state_unavailable",
+                "message": "workflow state is not queryable yet",
+            },
+        )
+        with (
+            patch.object(client._http, "request", new_callable=AsyncMock, return_value=resp),
+            pytest.raises(QueryFailed) as excinfo,
+        ):
+            await client.query_workflow("wf-1", "status")
+
+        assert excinfo.value.reason == "query_workflow_state_unavailable"
+
     @pytest.mark.asyncio
     async def test_worker_routed_query_timeout_raises_query_failed(self, client: Client) -> None:
         resp = _mock_response(504, {"reason": "query_worker_timeout", "message": "timed out"})
         with (
             patch.object(client._http, "request", new_callable=AsyncMock, return_value=resp),
-            pytest.raises(QueryFailed),
+            pytest.raises(QueryFailed) as excinfo,
+        ):
+            await client.query_workflow("wf-1", "status")
+
+        assert excinfo.value.reason == "query_worker_timeout"
+
+    @pytest.mark.asyncio
+    async def test_worker_routed_query_not_claimed_timeout_raises_query_failed(self, client: Client) -> None:
+        resp = _mock_response(
+            504,
+            {
+                "reason": "query_task_not_claimed",
+                "message": "timed out waiting for a compatible worker to claim the query",
+            },
+        )
+        with (
+            patch.object(client._http, "request", new_callable=AsyncMock, return_value=resp),
+            pytest.raises(QueryFailed) as excinfo,
         ):
             await client.query_workflow("wf-1", "status")
 
+        assert excinfo.value.reason == "query_task_not_claimed"
+
+    @pytest.mark.asyncio
+    async def test_worker_routed_query_execution_timeout_raises_query_failed(self, client: Client) -> None:
+        resp = _mock_response(
+            504,
+            {
+                "reason": "query_worker_execution_timeout",
+                "message": "worker leased the query but did not complete it",
+            },
+        )
+        with (
+            patch.object(client._http, "request", new_callable=AsyncMock, return_value=resp),
+            pytest.raises(QueryFailed) as excinfo,
+        ):
+            await client.query_workflow("wf-1", "status")
+
+        assert excinfo.value.reason == "query_worker_execution_timeout"
+
 
 class TestListWorkflows:
     @pytest.mark.asyncio

tests/test_worker.py

@@ -79,6 +79,16 @@ def run(self, ctx):  # type: ignore[no-untyped-def]
         return self.status
 
 
+@workflow.defn(name="query-state-unavailable-wf")
+class QueryStateUnavailableWorkflow:
+    @workflow.query("status")
+    def status_query(self) -> dict[str, str]:
+        return {"status": "ready"}
+
+    def run(self, ctx):  # type: ignore[no-untyped-def]
+        raise RuntimeError("state not ready")
+
+
 @workflow.defn(name="counter-query-wf")
 class CounterQueryWorkflow:
     def __init__(self) -> None:
@@ -1130,6 +1140,28 @@ async def test_query_task_reports_unknown_query(self, mock_client: AsyncMock) ->
         assert call_kwargs["query_task_attempt"] == 2
         assert call_kwargs["reason"] == "rejected_unknown_query"
 
+    @pytest.mark.asyncio
+    async def test_query_task_reports_state_unavailable_when_replay_fails(self, mock_client: AsyncMock) -> None:
+        worker = Worker(mock_client, task_queue="q1", workflows=[QueryStateUnavailableWorkflow], activities=[])
+        task = {
+            "query_task_id": "qt-state-unavailable",
+            "query_task_attempt": 1,
+            "workflow_type": "query-state-unavailable-wf",
+            "query_name": "status",
+            "history_events": [],
+            "workflow_arguments": serializer.envelope([], codec="json"),
+            "query_arguments": serializer.envelope([], codec="json"),
+            "payload_codec": "json",
+        }
+
+        outcome = await worker._run_query_task(task)
+
+        assert outcome == "failed"
+        mock_client.fail_query_task.assert_awaited_once()
+        call_kwargs = mock_client.fail_query_task.call_args.kwargs
+        assert call_kwargs["query_task_id"] == "qt-state-unavailable"
+        assert call_kwargs["reason"] == "query_workflow_state_unavailable"
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
         "reason",