Python: [Bug]: Functional workflow checkpoint replay returns a wrong cached result when a @step calls another @step #9383

Description

A @step that calls another @step works on the first run. After a crash and a resume from a checkpoint, a later call of the inner step returns the cached result of a different call, and its body never runs. The wrong result is returned silently, with no exception. This still reproduces on 1.21.0 when a step argument cannot be fingerprinted (for example a pydantic BaseModel), because MAF falls back to the call-order cache key.

Related: #8292 (closed). The argument fingerprint added in 1.20 fixes the case where all arguments are supported types, but not this one.

Code Sample

import asyncio

import agent_framework as maf
from pydantic import BaseModel

calls: list[str] = []
crash_at = "b"


class Job(BaseModel):
    label: str


@maf.step
async def tick(job: Job) -> str:
    calls.append(job.label)
    if job.label == crash_at:
        raise RuntimeError("killed")
    return job.label.upper()


@maf.step
async def wrapper(job: Job) -> str:
    return await tick(job)  # a step called from a step


@maf.workflow(name="repro")
async def repro(_: str) -> list[str]:
    return [await wrapper(Job(label="a")), await wrapper(Job(label="b"))]


async def main() -> None:
    global crash_at
    storage = maf.InMemoryCheckpointStorage()
    try:
        await repro.build(checkpoint_storage=storage).run("go")
    except RuntimeError as exc:
        print("first run crashed:", exc)

    calls.clear()
    crash_at = ""
    latest = await storage.get_latest(workflow_name="repro")
    result = await repro.build(checkpoint_storage=storage).run(checkpoint_id=latest.checkpoint_id)
    print("outputs:", result.get_outputs(), "step bodies run on resume:", calls)


asyncio.run(main())

Expected: [['A', 'B']], with tick(b) running on the resume.
Actual (1.18.0 and 1.21.0): [['A', 'A']], and no step body runs on the resume.

With str arguments the same workflow resumes correctly on 1.20.0 and 1.21.0 ([['A', 'B']], only b re-runs), and returns [['A', 'A']] on 1.18.0 and 1.19.0.

Cause (1.18.0)

StepWrapper.__call__ allocates the cache key with ctx._get_step_cache_key(name), which reads and increments _step_call_counters[name], so the key is (name, call order). On a cache hit the step returns without running its body, so steps called inside it never allocate their indices.

• First run: tick(a) gets ("tick", 0), wrapper(a) gets ("wrapper", 0), tick(b) gets ("tick", 1) and crashes.
• Resume: wrapper(a) is a hit, so tick's counter is not advanced. wrapper(b) runs, and its inner tick(b) gets ("tick", 0), which is already cached as "A".

Suggestions

1. Scope the counters per calling step (for example key (parent_key, name, idx)), or advance the children's counters on a cache hit.
2. Or raise a clear error when a step is entered while another step body is running, and document that steps must be flat.
3. Or, for arguments that cannot be fingerprinted, fail loudly instead of silently using the order-based key.

The docs do not currently say whether steps may call steps.

Package Versions

agent-framework-core: 1.18.0 and 1.21.0

Python Version

Python 3.12