microsoft / microsoft/agent-framework
Python: [Bug]: Functional workflow checkpoint replay can swap cached results between concurrent calls to the same @step
- Dominant language
- Python
- Stars
- 13.6k
- Forks
- 2.3k
- Avg merge
- 2d 45m
- Merged PRs (30d)
- 358
Description
### Description
A Functional Workflow can silently associate a cached result with the wrong logical branch when concurrent branches call the same decorated `@step` and the workflow is restored from a checkpoint.
The cache currently identifies completed invocations by `(step_name, call_index)`. A fresh replay restarts those per-step counters. If an earlier decorated step becomes a cache hit, its original suspension disappears and concurrent branches can reach a shared step in a different order:
```text
Initial execution:
B reaches shared_step first -> cache index 0 -> result:B
A reaches shared_step second -> cache index 1 -> result:A
Checkpoint replay:
A reaches shared_step first -> cache index 0 -> receives result:B
B reaches shared_step second -> cache index 1 -> receives result:A
```
The replay succeeds without an exception but returns the wrong branch-to-result mapping.
Expected behavior: checkpoint replay should preserve the logical association between a completed step invocation and its cached result. `shared_step("A")` should not receive the cached result of `shared_step("B")` solely because replay changed scheduling.
### Code Sample
```markdown
import asyncio
from agent_framework import InMemoryCheckpointStorage, step, workflow
async def main() -> None:
storage = InMemoryCheckpointStorage()
predecessor_started = asyncio.Event()
release_predecessor = asyncio.Event()
branch_a_shared_completed = asyncio.Event()
calls: list[str] = []
@step
async def predecessor(value: str) -> str:
predecessor_started.set()
await release_predecessor.wait()
return value
@step
async def shared_step(value: str) -> str:
calls.append(value)
if value == "B":
release_predecessor.set()
await branch_a_shared_completed.wait()
else:
branch_a_shared_completed.set()
return f"result:{value}"
async def branch_a() -> str:
await predecessor("A")
return await shared_step("A")
async def branch_b() -> str:
await predecessor_started.wait()
return await shared_step("B")
@workflow
async def parallel_workflow(_: str) -> list[str]:
return list(await asyncio.gather(branch_a(), branch_b()))
runnable = parallel_workflow.build(checkpoint_storage=storage)
initial = await runnable.run("input")
checkpoint = await storage.get_latest(workflow_name="parallel_workflow")
assert checkpoint is not None
replayed = await runnable.run(checkpoint_id=checkpoint.checkpoint_id)
print("initial:", initial.get_outputs())
print("replayed:", replayed.get_outputs())
print("executed shared inputs:", calls)
asyncio.run(main())
```
### Error Messages / Stack Traces
```markdown
No exception is raised. The replay succeeds with swapped cached values:
initial: [['result:A', 'result:B']]
replayed: [['result:B', 'result:A']]
executed shared inputs: ['B', 'A']
Expected replay output:
[['result:A', 'result:B']]
```
### Package Versions
agent-framework-core==1.18.0 (affected release: python-1.18.0); reproduced against upstream source e4309e5456c42ce391b982b6c61e60c8197b345c; the relevant Functional Workflow source was unchanged from the release tag and current main at verification time.
### Python Version
Python 3.12.12
### Additional Context
The `asyncio.Event` synchronization makes the initial order deterministic: B invokes `shared_step` first, releases A's predecessor, and then waits until A's shared invocation completes. The reproduction uses no sleeps, network calls, models, random scheduling, or nondeterministic input.
The [Functional Workflow documentation](https://learn.microsoft.com/en-us/agent-framework/concepts/workflows/functional) supports native `asyncio.gather` and explicitly says it also works with functions decorated by `@step`. The in-repository parallel sample further says each decorated branch is independently cached on HITL resume or checkpoint restore.
Functional Workflows were introduced in [PR #4238](https://github.com/microsoft/agent-framework/pull/4238). Its review discussion says the `(step_name, call_index)` cache relies on workflow determinism with respect to step results. This reproduction satisfies that constraint: inputs and completed step results are deterministic; the framework's cache hit removes an await and changes arrival order. The same review also rejected arbitrary argument equality as unreliable, confirmed that `@step` composes with `asyncio.gather`, and specifically suggested testing parallel decorated steps with checkpoint storage because invocation-index cache keys could expose ordering collisions.
This differs from [Issue #7647](https://github.com/microsoft/agent-framework/issues/7647). That issue concerns parallel saves forking checkpoint ancestry while final and restored values remain correct; this reproduction returns values associated with the wrong logical branches.
### Design note
Local fix exploration suggests this is not solved generally by replacing the global call index with another execution-order-derived ordinal. Assigning task identity at first `@step` arrival is also replay-order-dependent, and task-creation lineage can change when native async control flow creates later work based on completion order.
Late duplicate detection is also insufficient: an intermediate checkpoint may contain only one cached invocation, which can already be consumed by the wrong branch and alter control flow before another same-step call appears.
Static call-site identity collides in loops and shared helpers, while arbitrary argument hashing/equality is not suitable for the supported Python call surface. A generally safe fix therefore appears to require a replay-identity contract decision, such as an explicit durable logical invocation key/branch scope or early deterministic rejection of ambiguous replay patterns.
This appears to be ordinary functional correctness, not a security-boundary issue.
Contributor guide
Assessment
This issue has not been assessed yet.