From fa1049b1c93cbb22d6641d38e191b83cae412a55 Mon Sep 17 00:00:00 2001 From: Ruiming Zhao Date: Fri, 14 Aug 2026 16:48:18 -0700 Subject: [PATCH 1/2] Python: serialize per-step checkpoint saves in functional workflows Concurrent steps (asyncio.gather) could both read the same checkpoint chain head before either wrote it back, creating sibling root checkpoints and leaving part of the history unreachable from the latest checkpoint. Guard the read-save-update with an asyncio.Lock and add a regression test using a storage whose save() yields, which reproduces the interleaving deterministically. --- .../agent_framework/_workflows/_functional.py | 9 ++- .../workflow/test_functional_workflow.py | 68 +++++++++++++++++++ 2 files changed, 76 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 9e4deb340b..5bfdddb776 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -39,6 +39,7 @@ # pyright: reportPrivateUsage=false # Classes in this module (RunContext, StepWrapper, FunctionalWorkflow) form a # cohesive unit and intentionally access each other's underscore-prefixed members. +import asyncio import functools import hashlib import inspect @@ -1023,9 +1024,15 @@ async def _run_core( # Use a mutable list so the closure can update prev_checkpoint_id ckpt_chain: list[str | None] = [prev_checkpoint_id] if storage is not None: + # Concurrent steps (e.g. via asyncio.gather) may complete around + # the same time. Serialize the read-save-update of the chain head + # so each checkpoint links to the previous one instead of creating + # sibling root checkpoints that fork the lineage. + ckpt_chain_lock = asyncio.Lock() async def _on_step_completed() -> None: - ckpt_chain[0] = await self._save_checkpoint(ctx, storage, ckpt_chain[0]) + async with ckpt_chain_lock: + ckpt_chain[0] = await self._save_checkpoint(ctx, storage, ckpt_chain[0]) ctx._on_step_completed = _on_step_completed diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 96b7e6edf6..40cc1c99c1 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -24,6 +24,7 @@ InMemoryCheckpointStorage, RunContext, StepWrapper, + WorkflowCheckpoint, WorkflowEvent, WorkflowRunResult, WorkflowRunState, @@ -68,6 +69,21 @@ def decorate(fn: Callable[..., Awaitable[Any]]) -> FunctionalWorkflow: return decorate(func) if func is not None else decorate +class _YieldingCheckpointStorage(InMemoryCheckpointStorage): + """In-memory checkpoint storage whose ``save()`` yields to the event loop. + + Real backends (files, databases) suspend while persisting, which lets two + concurrent per-step checkpoint saves interleave. The plain in-memory + implementation never suspends, so it cannot exercise the race this file's + parallel-checkpoint regression test guards against. Yielding here + reproduces that interleaving deterministically. + """ + + async def save(self, checkpoint: WorkflowCheckpoint) -> str: + await asyncio.sleep(0) + return await super().save(checkpoint) + + @step async def add_one(x: int) -> int: return x + 1 @@ -797,6 +813,58 @@ async def wf(x: int) -> int: checkpoints = await storage.list_checkpoints(workflow_name="wf") assert len(checkpoints) == 3 # 2 from first run + 1 final from restore + async def test_parallel_steps_keep_single_checkpoint_lineage(self): + """Concurrent step completions must not fork the checkpoint chain. + + When steps run via ``asyncio.gather``, their completion callbacks can + overlap: both read the same ``previous_checkpoint_id`` before either + writes the updated chain head back, creating sibling root checkpoints + that leave part of the history unreachable from the latest checkpoint. + + A storage whose ``save()`` yields to the event loop reproduces the + interleaving a real (file/database) backend would see, deterministically. + """ + storage = _YieldingCheckpointStorage() + + @step + async def left(value: int) -> int: + return value + 1 + + @step + async def right(value: int) -> int: + return value + 2 + + @built_workflow(checkpoint_storage=storage) + async def parallel(value: int) -> list[int]: + return await asyncio.gather(left(value), right(value)) + + result = await parallel.run(1) + assert result.get_outputs() == [[2, 3]] + + checkpoints = await storage.list_checkpoints(workflow_name="parallel") + by_id = {cp.checkpoint_id: cp for cp in checkpoints} + + # Both steps plus the final save. + assert len(checkpoints) == 3 + + # Exactly one root: the chain head is read and updated under a lock. + roots = [cp for cp in checkpoints if cp.previous_checkpoint_id is None] + assert len(roots) == 1 + + # Every checkpoint is reachable from the latest one - no forked lineage. + latest = await storage.get_latest(workflow_name="parallel") + reachable: set[str] = set() + cursor = latest + while cursor is not None: + reachable.add(cursor.checkpoint_id) + cursor = by_id.get(cursor.previous_checkpoint_id) if cursor.previous_checkpoint_id else None + assert reachable == set(by_id) + + +# --------------------------------------------------------------------------- +# Branching / control flow +# --------------------------------------------------------------------------- + # --------------------------------------------------------------------------- # Branching / control flow From f68cb2b706cbfc810f1877c2125555b510fafaf9 Mon Sep 17 00:00:00 2001 From: Ruiming Zhao Date: Fri, 14 Aug 2026 16:52:54 -0700 Subject: [PATCH 2/2] Python: drop duplicated section header in workflow tests --- .../packages/core/tests/workflow/test_functional_workflow.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 40cc1c99c1..7a30291e7b 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -866,11 +866,6 @@ async def parallel(value: int) -> list[int]: # --------------------------------------------------------------------------- -# --------------------------------------------------------------------------- -# Branching / control flow -# --------------------------------------------------------------------------- - - class TestControlFlow: async def test_if_else_branching(self): @dataclass