[cross-repo from server#244] Conformance: run signals/queries against current published artifacts (#77)

Author: durable-workflow-ops

src/durable_workflow/worker.py

@@ -979,6 +979,8 @@ async def _run_query_task_core(self, task: dict[str, Any]) -> str:
                 external_storage=self.external_storage,
                 external_storage_cache=self.external_storage_cache,
             )
+            if inspect.isawaitable(result):
+                result = await result
         except AvroNotInstalledError as e:
             await self._fail_query_task(
                 query_task_id,
@@ -1034,7 +1036,49 @@ async def _run_query_task_core(self, task: dict[str, Any]) -> str:
                     e.reason(),
                 )
                 return "expired"
-            raise
+            server_reason = e.reason()
+            await self._fail_query_task(
+                query_task_id,
+                attempt,
+                str(e) or "query result completion was rejected by the server",
+                reason=server_reason if server_reason else "query_result_completion_failed",
+                failure_type=type(e).__name__,
+                stack_trace=traceback.format_exc(),
+            )
+            return "failed"
+        except AvroNotInstalledError as e:
+            await self._fail_query_task(
+                query_task_id,
+                attempt,
+                (
+                    "cannot encode query result with codec 'avro': "
+                    f"{e}. Reinstall durable-workflow with its runtime dependencies."
+                ),
+                reason="query_result_encode_failed",
+                failure_type=type(e).__name__,
+                stack_trace=traceback.format_exc(),
+            )
+            return "failed"
+        except (TypeError, ValueError) as e:
+            await self._fail_query_task(
+                query_task_id,
+                attempt,
+                f"cannot encode query result with codec {result_codec!r}: {e}",
+                reason="query_result_encode_failed",
+                failure_type=type(e).__name__,
+                stack_trace=traceback.format_exc(),
+            )
+            return "failed"
+        except Exception as e:
+            await self._fail_query_task(
+                query_task_id,
+                attempt,
+                str(e) or "query result completion failed",
+                reason="query_result_completion_failed",
+                failure_type=type(e).__name__,
+                stack_trace=traceback.format_exc(),
+            )
+            return "failed"
 
         return "completed"
 

tests/test_worker.py

@@ -78,6 +78,17 @@ def run(self, ctx):  # type: ignore[no-untyped-def]
         return self.status
 
 
+@workflow.defn(name="async-query-wf")
+class AsyncQueryWorkflow:
+    @workflow.query("current")
+    async def current(self) -> int:
+        await asyncio.sleep(0)
+        return 0
+
+    def run(self, ctx):  # type: ignore[no-untyped-def]
+        yield ctx.wait_condition(lambda: False)
+
+
 @activity.defn(name="test-act")
 def echo_activity(val: str) -> str:
     return f"result-{val}"
@@ -673,6 +684,35 @@ async def test_query_task_executes_registered_query(self, mock_client: AsyncMock
         )
         mock_client.fail_query_task.assert_not_called()
 
+    @pytest.mark.asyncio
+    async def test_query_task_awaits_async_query_result(self, mock_client: AsyncMock) -> None:
+        worker = Worker(mock_client, task_queue="q1", workflows=[AsyncQueryWorkflow], activities=[])
+        task = {
+            "query_task_id": "qt-async",
+            "query_task_attempt": 1,
+            "workflow_type": "async-query-wf",
+            "query_name": "current",
+            "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 == "completed"
+        mock_client.complete_query_task.assert_awaited_once_with(
+            query_task_id="qt-async",
+            lease_owner=worker.worker_id,
+            query_task_attempt=1,
+            result=0,
+            codec="json",
+            workflow_id=None,
+            run_id=None,
+            query_name="current",
+        )
+        mock_client.fail_query_task.assert_not_called()
+
     @pytest.mark.asyncio
     async def test_query_task_reports_unknown_query(self, mock_client: AsyncMock) -> None:
         worker = Worker(mock_client, task_queue="q1", workflows=[QueryWorkflow], activities=[])
@@ -723,6 +763,33 @@ async def test_query_task_completion_rejection_after_server_timeout_is_handled(
         mock_client.complete_query_task.assert_awaited_once()
         mock_client.fail_query_task.assert_not_called()
 
+    @pytest.mark.asyncio
+    async def test_query_task_reports_query_result_completion_failure(
+        self, mock_client: AsyncMock
+    ) -> None:
+        worker = Worker(mock_client, task_queue="q1", workflows=[QueryWorkflow], activities=[])
+        mock_client.complete_query_task.side_effect = TypeError("Object is not payload safe")
+        task = {
+            "query_task_id": "qt-result-encode-failure",
+            "query_task_attempt": 1,
+            "workflow_type": "query-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-result-encode-failure"
+        assert call_kwargs["query_task_attempt"] == 1
+        assert call_kwargs["reason"] == "query_result_encode_failed"
+        assert call_kwargs["failure_type"] == "TypeError"
+
 
 class TestActivityTaskExecution:
     @pytest.mark.asyncio