Conformance: CLI signal arguments are missing during SDK query replay (#81)

Author: durable-workflow-ops

src/durable_workflow/worker.py

@@ -25,7 +25,7 @@
 import time
 import traceback
 import uuid
-from collections.abc import Awaitable, Callable, Iterable
+from collections.abc import Awaitable, Callable, Iterable, Mapping
 from types import FunctionType
 from typing import Any
 
@@ -111,6 +111,136 @@ def _is_final_query_task_rejection(error: BaseException) -> bool:
     )
 
 
+def _signal_arguments_envelope_from_export(
+    signal: Mapping[str, Any],
+    *,
+    default_codec: str | None,
+) -> dict[str, Any] | None:
+    raw_arguments = signal.get("arguments")
+    if raw_arguments is None:
+        return None
+
+    codec = signal.get("payload_codec")
+    if not isinstance(codec, str) or codec == "":
+        codec = default_codec or serializer.AVRO_CODEC
+
+    if isinstance(raw_arguments, str):
+        if raw_arguments == "":
+            return None
+        return {"codec": codec, "blob": raw_arguments}
+
+    if not isinstance(raw_arguments, Mapping):
+        return None
+
+    envelope = dict(raw_arguments)
+    if "blob" not in envelope and "external_storage" not in envelope:
+        return None
+    envelope.setdefault("codec", codec)
+    return envelope
+
+
+def _query_history_with_export_signal_arguments(
+    history: Any,
+    history_export: Any,
+    *,
+    default_codec: str | None,
+) -> Any:
+    # Query-task history can carry compact SignalReceived rows; the full
+    # signal payload bytes are still present in the accompanying export.
+    if not isinstance(history, list) or not isinstance(history_export, Mapping):
+        return history
+
+    raw_signals = history_export.get("signals")
+    if not isinstance(raw_signals, list):
+        return history
+
+    export_payloads = history_export.get("payloads")
+    export_codec = (
+        export_payloads.get("codec")
+        if isinstance(export_payloads, Mapping)
+        else None
+    )
+    signal_default_codec = default_codec
+    if signal_default_codec is None and isinstance(export_codec, str) and export_codec:
+        signal_default_codec = export_codec
+
+    signals_by_id: dict[str, Mapping[str, Any]] = {}
+    signals_by_command_id: dict[str, Mapping[str, Any]] = {}
+    signals_by_name: dict[str, list[Mapping[str, Any]]] = {}
+    for raw_signal in raw_signals:
+        if not isinstance(raw_signal, Mapping):
+            continue
+        envelope = _signal_arguments_envelope_from_export(raw_signal, default_codec=signal_default_codec)
+        if envelope is None:
+            continue
+        signal_id = raw_signal.get("id")
+        if isinstance(signal_id, str) and signal_id:
+            signals_by_id[signal_id] = raw_signal
+        command_id = raw_signal.get("command_id")
+        if isinstance(command_id, str) and command_id:
+            signals_by_command_id[command_id] = raw_signal
+        name = raw_signal.get("name")
+        if isinstance(name, str) and name:
+            signals_by_name.setdefault(name, []).append(raw_signal)
+
+    if not signals_by_id and not signals_by_command_id and not signals_by_name:
+        return history
+
+    name_offsets: dict[str, int] = {}
+    enriched: list[Any] = []
+    changed = False
+    for raw_event in history:
+        if not isinstance(raw_event, Mapping):
+            enriched.append(raw_event)
+            continue
+        event_type = raw_event.get("event_type") or raw_event.get("type")
+        if event_type != "SignalReceived":
+            enriched.append(raw_event)
+            continue
+        raw_payload = raw_event.get("payload")
+        if not isinstance(raw_payload, Mapping):
+            enriched.append(dict(raw_event))
+            continue
+        if any(raw_payload.get(key) is not None for key in ("value", "input", "arguments")):
+            enriched.append(dict(raw_event))
+            continue
+
+        signal: Mapping[str, Any] | None = None
+        signal_id = raw_payload.get("signal_id")
+        if isinstance(signal_id, str) and signal_id:
+            signal = signals_by_id.get(signal_id)
+        if signal is None:
+            command_id = raw_payload.get("workflow_command_id") or raw_event.get("workflow_command_id")
+            if isinstance(command_id, str) and command_id:
+                signal = signals_by_command_id.get(command_id)
+        if signal is None:
+            signal_name = raw_payload.get("signal_name")
+            if isinstance(signal_name, str) and signal_name:
+                candidates = signals_by_name.get(signal_name, [])
+                offset = name_offsets.get(signal_name, 0)
+                if offset < len(candidates):
+                    signal = candidates[offset]
+                    name_offsets[signal_name] = offset + 1
+        if signal is None:
+            enriched.append(dict(raw_event))
+            continue
+
+        envelope = _signal_arguments_envelope_from_export(signal, default_codec=signal_default_codec)
+        if envelope is None:
+            enriched.append(dict(raw_event))
+            continue
+
+        payload = dict(raw_payload)
+        payload["arguments"] = envelope
+        payload.setdefault("payload_codec", envelope.get("codec"))
+        event = dict(raw_event)
+        event["payload"] = payload
+        enriched.append(event)
+        changed = True
+
+    return enriched if changed else history
+
+
 def _callable_fingerprint_payload(value: object) -> str:
     if isinstance(value, staticmethod | classmethod):
         value = value.__func__
@@ -948,7 +1078,11 @@ async def _run_query_task_core(self, task: dict[str, Any], *, client: Client | N
             return "failed"
 
         result_codec = _command_payload_codec(codec)
-        history = task.get("history_events", [])
+        history = _query_history_with_export_signal_arguments(
+            task.get("history_events", []),
+            task.get("history_export"),
+            default_codec=codec,
+        )
 
         cls = self.workflows.get(wf_type)
         if cls is None:

tests/test_worker.py

@@ -79,6 +79,23 @@ def run(self, ctx):  # type: ignore[no-untyped-def]
         return self.status
 
 
+@workflow.defn(name="counter-query-wf")
+class CounterQueryWorkflow:
+    def __init__(self) -> None:
+        self.count = 0
+
+    @workflow.signal("increment")
+    def increment(self, amount: int) -> None:
+        self.count += amount
+
+    @workflow.query("current")
+    def current(self) -> int:
+        return self.count
+
+    def run(self, ctx):  # type: ignore[no-untyped-def]
+        yield ctx.wait_condition(lambda: False)
+
+
 @workflow.defn(name="async-query-wf")
 class AsyncQueryWorkflow:
     @workflow.query("current")
@@ -685,6 +702,62 @@ 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_replays_signal_arguments_from_history_export(
+        self, mock_client: AsyncMock
+    ) -> None:
+        worker = Worker(mock_client, task_queue="q1", workflows=[CounterQueryWorkflow], activities=[])
+        signal_arguments = serializer.encode([3], codec="json")
+        task = {
+            "query_task_id": "qt-signal-export",
+            "query_task_attempt": 1,
+            "workflow_type": "counter-query-wf",
+            "workflow_id": "wf-counter",
+            "run_id": "run-counter",
+            "query_name": "current",
+            "history_events": [
+                {
+                    "event_type": "SignalReceived",
+                    "workflow_command_id": "cmd-increment",
+                    "payload": {
+                        "signal_id": "sig-increment",
+                        "workflow_command_id": "cmd-increment",
+                        "signal_name": "increment",
+                    },
+                },
+            ],
+            "history_export": {
+                "payloads": {"codec": "json"},
+                "signals": [
+                    {
+                        "id": "sig-increment",
+                        "command_id": "cmd-increment",
+                        "name": "increment",
+                        "payload_codec": "json",
+                        "arguments": signal_arguments,
+                    },
+                ],
+            },
+            "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-signal-export",
+            lease_owner=worker.worker_id,
+            query_task_attempt=1,
+            result=3,
+            codec="json",
+            workflow_id="wf-counter",
+            run_id="run-counter",
+            query_name="current",
+        )
+        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=[])