Preserve workflow memo upserts when continue-as-new is yielded

Author: durable-workflow-ops

CHANGELOG.md

@@ -89,6 +89,9 @@ Earlier SDK versions remain historical releases and are not alternate supported
   `NamespaceDescription.deleted`.
 
 ### Fixed
+- Yielded continue-as-new commands now preserve pending memo and search-
+  attribute upserts in the same order as returned continuations, while replayed
+  metadata updates remain consumed exactly once.
 - High-level workers now explicitly advertise memo-upsert and typed-search-
   attribute command support when registering with Server. Worker requests use
   protocol 1.16, while direct `Client.register_worker()` callers retain full

src/durable_workflow/workflow.py

@@ -4168,7 +4168,7 @@ def _terminal_state(value: Any, *, include_pending: bool) -> _ReplayState:
                 pending.extend(cmd)
                 return _state(pending)
             if isinstance(cmd, ContinueAsNew):
-                return _state([cmd])
+                return _terminal_state(cmd, include_pending=True)
             if isinstance(cmd, NexusServiceCall):
                 if result_cursor < len(resolved_results):
                     _assert_next_step_matches(cmd)

tests/fixtures/replay_regressions/yielded-continue-preserves-pending-metadata.json

@@ -0,0 +1,58 @@
+{
+  "$schema": "https://raw.githubusercontent.com/durable-workflow/.github/main/regression-corpus/evidence-schema.json",
+  "fixture_schema": "durable-workflow.replay-regression/v1",
+  "id": "yielded-continue-preserves-pending-metadata",
+  "protocol_version": "1.16",
+  "bindings": [
+    "python"
+  ],
+  "workflow": {
+    "type": "tests.replay.yielded-continue-after-metadata",
+    "input": [],
+    "payload_codec": "avro"
+  },
+  "command_sequence": [
+    {
+      "type": "upsert_search_attributes",
+      "attributes": {
+        "stage": "continued"
+      }
+    },
+    {
+      "type": "upsert_memo",
+      "entries": {
+        "added": "from-upsert",
+        "overwritten": "after"
+      }
+    },
+    {
+      "type": "continue_as_new",
+      "arguments": [
+        "successor"
+      ]
+    }
+  ],
+  "expected": {
+    "command_sequence": [
+      {
+        "type": "upsert_search_attributes",
+        "attributes": {
+          "stage": "continued"
+        }
+      },
+      {
+        "type": "upsert_memo",
+        "entries": {
+          "added": "from-upsert",
+          "overwritten": "after"
+        }
+      },
+      {
+        "type": "continue_as_new",
+        "arguments": [
+          "successor"
+        ]
+      }
+    ]
+  }
+}

tests/integration/test_memo_restart.py

@@ -8,6 +8,7 @@
 import pytest
 
 from durable_workflow import Client, Worker, serializer, workflow
+from durable_workflow.workflow import commands_to_server_commands, replay
 
 MEMO_ENTRIES: dict[str, Any] = {
     "binary": b"same",
@@ -26,6 +27,22 @@
     "CGxvbmcEDgxuZXN0ZWQOBAphbHBoYQQCCGJldGEEBAAIdGV4dAoIc2FtZQA="
 )
 
+CONTINUE_INITIAL_MEMO = {
+    "existing": "retained",
+    "overwritten": "before",
+}
+
+CONTINUE_MEMO_PATCH = {
+    "added": "from-upsert",
+    "overwritten": "after",
+}
+
+CONTINUE_MERGED_MEMO = {
+    "added": "from-upsert",
+    "existing": "retained",
+    "overwritten": "after",
+}
+
 
 @workflow.defn(name="tests.memo-restart-python")
 class MemoRestartWorkflow:
@@ -46,6 +63,15 @@ def run(self, ctx):  # type: ignore[no-untyped-def]
         return "python-replayed-memo"
 
 
+@workflow.defn(name="tests.memo-yielded-continue-python")
+class MemoYieldedContinueWorkflow:
+    def run(self, ctx, generation: int):  # type: ignore[no-untyped-def]
+        if generation == 0:
+            yield ctx.upsert_memo(CONTINUE_MEMO_PATCH)
+            yield ctx.continue_as_new(1)
+        return "python-continued-with-memo"
+
+
 async def _wait_until_waiting_with_memo(handle: Any) -> Any:
     deadline = asyncio.get_running_loop().time() + 30
     last_description = None
@@ -94,6 +120,7 @@ async def _seed_worker_contract(client: Client, worker: Worker) -> None:
         workflow_definition_fingerprints=worker.workflow_definition_fingerprints,
         workflow_command_contracts=worker.workflow_command_contracts,
         supported_activity_types=[],
+        capabilities=["memo_upserts"],
     )
 
 
@@ -175,3 +202,105 @@ async def test_fresh_python_worker_replays_persisted_server_memo_history(
         final_memo_events = _memo_events(final_history)
         assert len(final_memo_events) == 1
         _assert_typed_memo_event(final_memo_events[0])
+
+
+@pytest.mark.asyncio
+async def test_yielded_continue_inherits_merged_memo_after_worker_restart(
+    server_url: str,
+    server_token: str,
+) -> None:
+    suffix = uuid.uuid4().hex[:8]
+    task_queue = f"memo-continue-python-{suffix}"
+    workflow_id = f"memo-continue-python-{suffix}"
+
+    async with Client(server_url, token=server_token, namespace="default") as client:
+        first_worker = Worker(
+            client,
+            task_queue=task_queue,
+            workflows=[MemoYieldedContinueWorkflow],
+            activities=[],
+            worker_id=f"memo-continue-before-{suffix}",
+        )
+        await _seed_worker_contract(client, first_worker)
+
+        handle = await client.start_workflow(
+            workflow_type="tests.memo-yielded-continue-python",
+            task_queue=task_queue,
+            workflow_id=workflow_id,
+            input=[0],
+            memo=CONTINUE_INITIAL_MEMO,
+        )
+        assert handle.run_id is not None
+        first_run_id = handle.run_id
+
+        try:
+            first_task = await client.poll_workflow_task(
+                worker_id=first_worker.worker_id,
+                task_queue=task_queue,
+                timeout=10.0,
+            )
+            assert first_task is not None, "expected the initial workflow task"
+
+            payload_codec = first_task.get("payload_codec") or serializer.AVRO_CODEC
+            decoded = serializer.decode_envelope(
+                first_task.get("arguments"),
+                codec=payload_codec,
+            )
+            start_input = decoded if isinstance(decoded, list) else [decoded]
+            outcome = replay(
+                MemoYieldedContinueWorkflow,
+                first_task.get("history_events", []),
+                start_input,
+                run_id=first_task.get("run_id", ""),
+            )
+            commands = commands_to_server_commands(
+                outcome.commands,
+                task_queue,
+                payload_codec=payload_codec,
+            )
+            assert [command["type"] for command in commands] == [
+                "upsert_memo",
+                "continue_as_new",
+            ]
+
+            await client.complete_workflow_task(
+                task_id=first_task["task_id"],
+                lease_owner=first_worker.worker_id,
+                workflow_task_attempt=first_task.get("workflow_task_attempt", 1),
+                commands=commands,
+            )
+        finally:
+            await client.deregister_worker_registration(first_worker.worker_id)
+
+        replacement_worker = Worker(
+            client,
+            task_queue=task_queue,
+            workflows=[MemoYieldedContinueWorkflow],
+            activities=[],
+            worker_id=f"memo-continue-after-{suffix}",
+            poll_timeout=1.0,
+            shutdown_timeout=5.0,
+        )
+        completed = await replacement_worker.run_until(
+            workflow_id=workflow_id,
+            timeout=30.0,
+            poll_interval=0.1,
+        )
+
+        assert (completed.status or "").lower() == "completed"
+        assert completed.output == "python-continued-with-memo"
+        assert completed.run_id != first_run_id
+        assert completed.memo == CONTINUE_MERGED_MEMO
+
+        runs = await handle.list_runs()
+        assert runs.run_count == 2
+        assert [run.run_id for run in runs.runs] == [first_run_id, completed.run_id]
+        assert runs.runs[1].memo == CONTINUE_MERGED_MEMO
+
+        first_history = await client.get_history(workflow_id, first_run_id)
+        successor_history = await client.get_history(workflow_id, completed.run_id or "")
+        first_memo_events = _memo_events(first_history)
+        assert len(first_memo_events) == 1
+        assert serializer.decode_envelope(first_memo_events[0]["payload"]["entries"]) == CONTINUE_MEMO_PATCH
+        assert serializer.decode_envelope(first_memo_events[0]["payload"]["merged"]) == CONTINUE_MERGED_MEMO
+        assert _memo_events(successor_history) == []

tests/test_replay.py

@@ -1003,6 +1003,22 @@ def run(self, ctx: WorkflowContext, counter: int):  # type: ignore[no-untyped-de
         return ctx.continue_as_new(counter - 1)
 
 
+@workflow.defn(name="continue-as-new-after-metadata-return-wf")
+class ContinueAsNewAfterMetadataReturnWorkflow:
+    def run(self, ctx: WorkflowContext):  # type: ignore[no-untyped-def]
+        yield ctx.upsert_search_attributes({"stage": "continued"})
+        yield ctx.upsert_memo({"added": "from-upsert", "overwritten": "after"})
+        return ctx.continue_as_new("successor")
+
+
+@workflow.defn(name="continue-as-new-after-metadata-yield-wf")
+class ContinueAsNewAfterMetadataYieldWorkflow:
+    def run(self, ctx: WorkflowContext):  # type: ignore[no-untyped-def]
+        yield ctx.upsert_search_attributes({"stage": "continued"})
+        yield ctx.upsert_memo({"added": "from-upsert", "overwritten": "after"})
+        yield ctx.continue_as_new("successor")
+
+
 @workflow.defn(name="side-effect-wf")
 class SideEffectWorkflow:
     def run(self, ctx: WorkflowContext):  # type: ignore[no-untyped-def]
@@ -1042,6 +1058,40 @@ def test_generator_return_continue(self) -> None:
         assert isinstance(cmd, ContinueAsNew)
         assert cmd.arguments == [4]
 
+    @pytest.mark.parametrize(
+        "workflow_cls",
+        [
+            ContinueAsNewAfterMetadataReturnWorkflow,
+            ContinueAsNewAfterMetadataYieldWorkflow,
+        ],
+    )
+    def test_pending_metadata_precedes_every_continue_authoring_form(self, workflow_cls: type) -> None:
+        first = replay(workflow_cls, [], [])
+
+        assert [command.to_server_command("q")["type"] for command in first.commands] == [
+            "upsert_search_attributes",
+            "upsert_memo",
+            "continue_as_new",
+        ]
+
+        entries = _avro({"added": "from-upsert", "overwritten": "after"})
+        merged = _avro({"added": "from-upsert", "overwritten": "after"})
+        replayed = replay(
+            workflow_cls,
+            [
+                {"event_type": "SearchAttributesUpserted", "payload": {"sequence": 1}},
+                {
+                    "event_type": "MemoUpserted",
+                    "payload": {"sequence": 2, "entries": entries, "merged": merged},
+                },
+            ],
+            [],
+        )
+
+        assert len(replayed.commands) == 1
+        assert isinstance(replayed.commands[0], ContinueAsNew)
+        assert replayed.commands[0].arguments == ["successor"]
+
     def test_server_command_shape(self) -> None:
         cmd = ContinueAsNew(workflow_type="other", arguments=[1, 2], task_queue="q2")
         sc = cmd.to_server_command("default-q")

tests/test_replay_regression_corpus.py

@@ -101,6 +101,14 @@ def run(self, ctx: WorkflowContext):  # type: ignore[no-untyped-def]
         return "memo-replayed"
 
 
+@workflow.defn(name="tests.replay.yielded-continue-after-metadata")
+class YieldedContinueAfterMetadataWorkflow:
+    def run(self, ctx: WorkflowContext):  # type: ignore[no-untyped-def]
+        yield ctx.upsert_search_attributes({"stage": "continued"})
+        yield ctx.upsert_memo({"added": "from-upsert", "overwritten": "after"})
+        yield ctx.continue_as_new("successor")
+
+
 WORKFLOWS = [
     GoldenSagaCompensationWorkflow,
     GoldenSignalWaitWorkflow,
@@ -114,6 +122,7 @@ def run(self, ctx: WorkflowContext):  # type: ignore[no-untyped-def]
     UpdateSignalConditionTimerWorkflow,
     WorkflowStreamAuthorWorkflow,
     WorkflowMemoAuthorWorkflow,
+    YieldedContinueAfterMetadataWorkflow,
 ]
 WORKFLOW_TYPES = {str(getattr(workflow, "__workflow_name__", workflow.__name__)): workflow for workflow in WORKFLOWS}