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..7a30291e7b 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,53 @@ 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