[cross-repo from server#313] Conformance blocker: expand signals and queries coverage beyond current smoke (#115)

Author: durable-workflow-ops

src/durable_workflow/worker.py

@@ -241,10 +241,6 @@ def _query_history_with_export_signal_arguments(
         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:
@@ -265,14 +261,26 @@ def _query_history_with_export_signal_arguments(
             enriched.append(dict(raw_event))
             continue
 
+        payload = dict(raw_payload)
+        payload_changed = False
+        workflow_sequence = signal.get("workflow_sequence")
+        if isinstance(workflow_sequence, int):
+            payload.setdefault("workflow_sequence", workflow_sequence)
+            payload_changed = True
+        elif isinstance(workflow_sequence, str) and workflow_sequence.isdigit():
+            payload.setdefault("workflow_sequence", int(workflow_sequence))
+            payload_changed = True
+
         envelope = _signal_arguments_envelope_from_export(signal, default_codec=signal_default_codec)
-        if envelope is None:
+        if envelope is not None:
+            payload.setdefault("arguments", envelope)
+            payload.setdefault("payload_codec", envelope.get("codec"))
+            payload_changed = True
+
+        if not payload_changed:
             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)

src/durable_workflow/workflow.py

@@ -2055,6 +2055,7 @@ def _replay_state(
     event_types_by_sequence: dict[int, list[str]] = {}
     details_by_sequence: dict[int, dict[str, Any]] = {}
     resolved_sequences: set[int] = set()
+    condition_wait_ids_by_sequence: dict[int, str] = {}
 
     for event in events:
         event_type = _history_event_type(event)
@@ -2069,6 +2070,10 @@ def _replay_state(
         details_by_sequence.setdefault(sequence, {}).update(_recorded_step_details(payload))
         if _is_resolved_step_event(event_type, payload):
             resolved_sequences.add(sequence)
+        if event_type == "ConditionWaitOpened":
+            wait_id = payload.get("condition_wait_id")
+            if isinstance(wait_id, str) and wait_id:
+                condition_wait_ids_by_sequence[sequence] = wait_id
 
     workflow_start_time: datetime | None = None
     for ev in events:
@@ -2307,6 +2312,12 @@ def _close_condition_wait(wait_id: Any) -> None:
         elif etype == "SignalReceived":
             signal_name = payload.get("signal_name")
             if isinstance(signal_name, str) and signal_name:
+                workflow_sequence = _workflow_sequence(payload)
+                condition_wait_id = (
+                    condition_wait_ids_by_sequence.get(workflow_sequence)
+                    if workflow_sequence is not None
+                    else None
+                )
                 pending_receivers.append(_PendingReceiver(
                     result_index=len(resolved_results),
                     kind="signal",
@@ -2321,7 +2332,7 @@ def _close_condition_wait(wait_id: Any) -> None:
                         external_storage=external_storage,
                         external_storage_cache=external_storage_cache,
                     ),
-                    condition_wait_id=_current_condition_wait_id(),
+                    condition_wait_id=condition_wait_id or _current_condition_wait_id(),
                 ))
         elif etype == "UpdateApplied":
             update_name = payload.get("update_name")

tests/test_wait_condition.py

@@ -11,14 +11,18 @@
 )
 
 
-def _signal_received_event(name: str, args: list) -> dict:
+def _signal_received_event(name: str, args: list, workflow_sequence: int | None = None) -> dict:
+    payload = {
+        "signal_name": name,
+        "value": serializer.envelope(args),
+        "payload_codec": serializer.AVRO_CODEC,
+    }
+    if workflow_sequence is not None:
+        payload["workflow_sequence"] = workflow_sequence
+
     return {
         "event_type": "SignalReceived",
-        "payload": {
-            "signal_name": name,
-            "value": serializer.envelope(args),
-            "payload_codec": serializer.AVRO_CODEC,
-        },
+        "payload": payload,
     }
 
 
@@ -414,6 +418,52 @@ def test_repeated_physical_waits_keep_query_state_current_before_finish(self) ->
             ],
         }
 
+    def test_signal_received_before_next_wait_uses_declared_workflow_sequence(self) -> None:
+        history = [
+            {
+                "event_type": "ConditionWaitOpened",
+                "payload": {
+                    "condition_wait_id": "wait-count-3",
+                    "condition_key": "done",
+                    "sequence": 1,
+                },
+            },
+            _signal_received_event("increment", [3], workflow_sequence=1),
+            _signal_received_event("finish", [], workflow_sequence=2),
+            {
+                "event_type": "ConditionWaitOpened",
+                "payload": {
+                    "condition_wait_id": "wait-finish",
+                    "condition_key": "done",
+                    "sequence": 2,
+                },
+            },
+            {
+                "event_type": "ConditionWaitSatisfied",
+                "payload": {
+                    "condition_wait_id": "wait-finish",
+                    "condition_key": "done",
+                    "sequence": 2,
+                },
+            },
+        ]
+
+        expected = {
+            "count": 3,
+            "done": True,
+            "events": [
+                {"signal": "increment", "amount": 3, "count": 3},
+                {"signal": "finish", "count": 3},
+            ],
+        }
+
+        outcome = replay(SignalCounterUntilFinished, history, [])
+
+        assert len(outcome.commands) == 1
+        assert isinstance(outcome.commands[0], CompleteWorkflow)
+        assert outcome.commands[0].result == expected
+        assert query_state(SignalCounterUntilFinished, history, [], "current") == expected
+
     def test_repeated_wait_after_activity_can_be_satisfied_by_later_signal(self) -> None:
         history = [
             {

tests/test_worker.py

@@ -32,7 +32,11 @@
     WorkflowTaskHandler,
     WorkflowTaskInterceptorContext,
 )
-from durable_workflow.worker import Worker, _should_fail_workflow_task_after_completion_error
+from durable_workflow.worker import (
+    Worker,
+    _query_history_with_export_signal_arguments,
+    _should_fail_workflow_task_after_completion_error,
+)
 
 
 @workflow.defn(name="test-wf")
@@ -916,6 +920,39 @@ async def test_update_task_retries_transient_completion_error(
         assert mock_client.complete_workflow_task.await_count == 2
         mock_client.fail_workflow_task.assert_not_called()
 
+    def test_query_history_enrichment_copies_signal_workflow_sequence_from_export(self) -> None:
+        history = [
+            {
+                "event_type": "SignalReceived",
+                "workflow_command_id": "cmd-finish",
+                "payload": {
+                    "signal_id": "sig-finish",
+                    "workflow_command_id": "cmd-finish",
+                    "signal_name": "finish",
+                },
+            },
+        ]
+        export = {
+            "payloads": {"codec": "json"},
+            "signals": [
+                {
+                    "id": "sig-finish",
+                    "command_id": "cmd-finish",
+                    "name": "finish",
+                    "workflow_sequence": 2,
+                    "payload_codec": "json",
+                    "arguments": serializer.encode([], codec="json"),
+                },
+            ],
+        }
+
+        enriched = _query_history_with_export_signal_arguments(history, export, default_codec="json")
+
+        assert isinstance(enriched, list)
+        payload = enriched[0]["payload"]
+        assert payload["workflow_sequence"] == 2
+        assert payload["arguments"]["codec"] == "json"
+
     @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=[])