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

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:

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

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

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 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. 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. 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.