[cross-repo from server#267] Conformance: replay approve signal does not complete after activity state query (#84)

Author: durable-workflow-ops

CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Fixed
+- Signal and update receivers recorded after an activity result now replay after
+  the workflow consumes that activity result, so receiver-mutated state is not
+  overwritten by deterministic post-activity setup before a `wait_condition`.
 - Python workflow replay now throws terminal `ActivityFailed` history events
   into the generator as a typed `ActivityFailed` exception, including the
   recorded activity and failure metadata. Activity-only saga workflows can

src/durable_workflow/workflow.py

@@ -2052,8 +2052,16 @@ def _state(commands: list[Command]) -> _ReplayState:
     signal_registry: dict[str, str] = getattr(workflow_cls, "__workflow_signals__", {}) or {}
     update_registry: dict[str, str] = getattr(workflow_cls, "__workflow_updates__", {}) or {}
 
-    def _apply_due_receivers() -> None:
-        while pending_receivers and pending_receivers[0][0] <= result_cursor:
+    def _apply_due_receivers(*, before_consuming_result: bool = False) -> None:
+        while pending_receivers:
+            receiver_index = pending_receivers[0][0]
+            due = (
+                receiver_index < result_cursor
+                if before_consuming_result
+                else receiver_index <= result_cursor
+            )
+            if not due:
+                break
             _, kind, name, args = pending_receivers.pop(0)
             if kind == "signal":
                 method_name = signal_registry.get(name)
@@ -2085,19 +2093,30 @@ def _apply_due_receivers() -> None:
     pending: list[Command] = []
     advanced_cmd: Any = None
     wait_yield_count = 0
+
+    def _terminal_state(value: Any, *, include_pending: bool) -> _ReplayState:
+        _apply_due_receivers()
+        commands = list(pending) if include_pending else []
+        if isinstance(value, ContinueAsNew):
+            return _state(commands + [value])
+        return _state(commands + [CompleteWorkflow(result=value)])
+
     try:
         while True:
             # Cursor-0 receivers are start-boundary events. Enter run() once
             # before applying them so workflow initialization observes the
             # same WorkflowStarted-before-Signal/Update ordering the server
             # records in history.
             if not first:
-                _apply_due_receivers()
+                _apply_due_receivers(before_consuming_result=True)
             if advanced_cmd is not None:
                 cmd = advanced_cmd
                 advanced_cmd = None
             else:
-                cmd = gen.send(None) if first else gen.send(next_value)
+                try:
+                    cmd = gen.send(None) if first else gen.send(next_value)
+                except StopIteration as stop:
+                    return _terminal_state(stop.value, include_pending=True)
                 first = False
             _apply_due_receivers()
             if isinstance(cmd, list):
@@ -2111,9 +2130,7 @@ def _apply_due_receivers() -> None:
                             advanced_cmd = gen.throw(failed)
                             continue
                         except StopIteration as stop:
-                            if isinstance(stop.value, ContinueAsNew):
-                                return _state([stop.value])
-                            return _state([CompleteWorkflow(result=stop.value)])
+                            return _terminal_state(stop.value, include_pending=False)
                     next_value = vals
                     continue
                 ctx.logger._set_replaying(False)
@@ -2208,19 +2225,21 @@ def _apply_due_receivers() -> None:
                             advanced_cmd = gen.throw(val)
                             continue
                         except StopIteration as stop:
-                            if isinstance(stop.value, ContinueAsNew):
-                                return _state([stop.value])
-                            return _state([CompleteWorkflow(result=stop.value)])
+                            return _terminal_state(stop.value, include_pending=False)
                     next_value = val
                     continue
                 ctx.logger._set_replaying(False)
                 pending.append(cmd)
                 return _state(pending)
             raise TypeError(f"workflow yielded unsupported command: {cmd!r}")
     except StopIteration as stop:
-        if isinstance(stop.value, ContinueAsNew):
-            return _state(pending + [stop.value])
-        return _state(pending + [CompleteWorkflow(result=stop.value)])
+        try:
+            return _terminal_state(stop.value, include_pending=True)
+        except Exception as exc:
+            return _state([FailWorkflow(
+                message=str(exc),
+                exception_type=type(exc).__name__,
+            )])
     except Exception as exc:
         return _state([FailWorkflow(
             message=str(exc),

tests/test_updates.py

@@ -130,6 +130,16 @@ def _workflow_started_event() -> dict[str, object]:
     }
 
 
+def _activity_completed_event(result: object) -> dict[str, object]:
+    return {
+        "event_type": "ActivityCompleted",
+        "payload": {
+            "result": serializer.encode(result, codec="json"),
+            "payload_codec": "json",
+        },
+    }
+
+
 def _update_accepted_event(
     update_id: str,
     name: str,
@@ -210,6 +220,21 @@ def test_replay_applies_committed_update_events_to_workflow_state(self) -> None:
             payload_codec="json",
         ) == 7
 
+    def test_query_state_applies_update_after_final_activity_result(self) -> None:
+        history = [
+            _signal_received_event("increment", [1]),
+            _activity_completed_event(None),
+            _update_applied_event("upd-final", "increment", [5]),
+        ]
+
+        assert query_state(
+            StatefulUpdateReceiver,
+            history,
+            [],
+            "count",
+            payload_codec="json",
+        ) == 6
+
     def test_apply_update_completes_accepted_update_with_handler_result(self) -> None:
         history = [
             _signal_received_event("increment", [2]),

tests/test_wait_condition.py

@@ -6,6 +6,7 @@
     FailWorkflow,
     WaitCondition,
     WorkflowContext,
+    query_state,
     replay,
 )
 
@@ -60,6 +61,38 @@ def run(self, ctx: WorkflowContext):  # type: ignore[no-untyped-def]
         return "satisfied" if satisfied else "no"
 
 
+@workflow.defn(name="activity-then-wait-for-approval")
+class ActivityThenWaitForApproval:
+    def __init__(self) -> None:
+        self.activity_result: str | None = None
+        self.approved_by: str | None = None
+
+    @workflow.signal("approve")
+    def approve(self, approved_by: str) -> None:
+        self.approved_by = approved_by
+
+    @workflow.query("state")
+    def state(self) -> dict[str, str | None]:
+        return {
+            "activity_result": self.activity_result,
+            "approved_by": self.approved_by,
+        }
+
+    def run(self, ctx: WorkflowContext):  # type: ignore[no-untyped-def]
+        self.activity_result = yield ctx.schedule_activity("load-state", [])
+        self.approved_by = None
+        yield ctx.wait_condition(
+            lambda: self.approved_by == "alice",
+            key="approval",
+            timeout=30,
+        )
+
+        return {
+            "activity_result": self.activity_result,
+            "approved_by": self.approved_by,
+        }
+
+
 class TestCtxWaitCondition:
     def test_wait_condition_returns_dataclass_with_predicate_and_key(self) -> None:
         ctx = WorkflowContext(run_id="x")
@@ -188,6 +221,32 @@ def test_condition_wait_satisfied_event_resolves_wait_even_if_predicate_false(se
         assert isinstance(outcome.commands[0], CompleteWorkflow)
         assert outcome.commands[0].result == "approved"
 
+    def test_signal_after_activity_result_applies_after_result_is_consumed(self) -> None:
+        history = [
+            {
+                "event_type": "ActivityCompleted",
+                "payload": {"result": '"loaded"'},
+            },
+            {
+                "event_type": "ConditionWaitOpened",
+                "payload": {"condition_wait_id": "wait-1", "condition_key": "approval"},
+            },
+            _signal_received_event("approve", ["alice"]),
+        ]
+
+        outcome = replay(ActivityThenWaitForApproval, history, [])
+
+        assert len(outcome.commands) == 1
+        assert isinstance(outcome.commands[0], CompleteWorkflow)
+        assert outcome.commands[0].result == {
+            "activity_result": "loaded",
+            "approved_by": "alice",
+        }
+        assert query_state(ActivityThenWaitForApproval, history, [], "state") == {
+            "activity_result": "loaded",
+            "approved_by": "alice",
+        }
+
     def test_open_with_no_resolution_and_predicate_false_re_emits_wait_condition(self) -> None:
         history = [
             {