microsoft / microsoft/agent-framework

Python: [Bug]: Functional workflow checkpoint replay can swap cached results between concurrent calls to the same @step

Open
#8,292 5 comments 0 reactions 1 assignee Claimed by @eavanvalkenburg View on GitHub
python reproduced workflows
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.