Skip to content

[BREAKING] Python: Reject unsafe functional step replay - #9391

Draft
Atharva Vichare (atty57) wants to merge 2 commits into
microsoft:mainfrom
atty57:fix/python-functional-replay-9385
Draft

Atharva Vichare (atty57) wants to merge 2 commits into
microsoft:mainfrom
atty57:fix/python-functional-replay-9385

Conversation

@atty57

@atty57 Atharva Vichare (atty57) commented Oct 14, 2026 •

Copy link
Copy Markdown
Contributor

Motivation & Context

Bypassing a cached outer @step skips its inner call counters. A later nested or top-level call can then silently reuse an earlier invocation's result from an order-only cache. Legacy checkpoints remain unsafe even when current arguments are fingerprintable or callers supply explicit replay keys; opaque arguments could also produce unsafe new checkpoints.

Description & Review Guide

  • What are the major changes? Reject checkpointed calls without a stable replay identity before executing their bodies, reject every unverified legacy order-only cache hit, and prevent opaque results from being retained for response-only replay. Preserve versioned cache lookup precedence, add regression coverage, and clarify the step docstrings.
  • What is the impact of these changes? Opaque arguments require an explicit replay_key for checkpoint/HITL replay. Legacy order-only checkpoints are no longer replayable, including flat sequential checkpoints, and must be regenerated in a new run. Existing versioned checkpoints continue to replay. Sequential execution without replay still supports opaque arguments. With explicit keys, nested crash recovery returns [['A', 'B']] and only reruns b.
  • What do you want reviewers to focus on? Rejecting unverified legacy fallback at every execution depth, including direct calls after a cached outer step; keeping verified versioned cache hits ahead of the legacy guard; and the documented checkpoint compatibility change.

This chooses the fail-loud fix from issue option #3 over automatically scoping counters (option #1), keeping the fix focused and requiring explicit invocation identity for opaque inputs.

Related Issue

Fixes #9385

Maintainer agreement on this direction is pending; no competing open PR was found before opening this PR.

AI Assistance

  • No material AI assistance was used.
  • This is an AI-assisted contribution. I reviewed, understood, and verified all submitted content and accept responsibility for it.

AI assistance details: Issue analysis, implementation, tests, docstrings, and this PR description.

Contribution Checklist

  • The code builds clean without any errors or warnings
  • All unit tests pass, and I have added new tests where possible
  • The PR follows the Contribution Guidelines
  • This PR links to an agreed issue with no competing open PR, or the Related Issue section documents a trivial-change or repository-automation exception.
  • This is not a breaking change. If it is a breaking change, add the breaking change label (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.

The title flags the intentional compatibility change for legacy order-only checkpoints. All 168 functional workflow tests pass, including the nested and mixed nested/direct legacy replay regressions. Ruff lint/format, source and test Pyright, and test mypy checks pass on the touched files. The full poe check is left to CI; maintainer agreement is pending.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Nested legacy replay can still silently return the wrong cached result when arguments are fingerprintable.

1 open finding
What changed in this PR

Adds fail-loud safeguards for functional workflow replay when step arguments lack stable identities.

Changes:

  • Rejects unsafe checkpoint and HITL replay for opaque arguments.
  • Validates legacy cache hits and replayable results.
  • Adds regression tests and updates replay documentation.
File Description
python/​packages/​core/​agent_framework/​_workflows/​_functional.py Adds replay-safety guards and documentation.
python/​packages/​core/​tests/​workflow/​test_functional_workflow.py Adds opaque-argument replay regression tests.

🧠 Review effort: Balanced


💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.

@atty57 Atharva Vichare (atty57) changed the title Python: Reject unsafe functional step replay [BREAKING] Python: Reject unsafe functional step replay Oct 14, 2026
@agent-framework-automation agent-framework-automation Bot added the breaking change Usage: [PRs], Target: all PRs that introduce changes that are not backward compatible label Oct 14, 2026
@snovik75

Copy link
Copy Markdown

Eduard van Valkenburg (@eavanvalkenburg) thanks for the plan. The nested legacy-replay failure and the opaque-argument example share one root cause, and one small change fixes both without rejecting any legacy cache hit. Findings and a sketch of the change below.

Root cause. StepWrapper.__call__ allocates a legacy name::idx key for every call, including calls that have an automatic or explicit identity, and falls back to it when the versioned entry is missing. The counter behind it (RunContext._step_call_counters) is global per run. A cached parent is skipped, so the calls inside its body never advance it, and a later inner call takes an earlier invocation's entry. The automatic-identity occurrence counter (_step_identity_counters) has the same property.

Change. This is suggestion 1 of the issue: keep the order-based counters, but scope them to the calling step.

  • A ContextVar holds a hash of the cache key of the step body that is running (empty at workflow level).
  • The legacy counter is keyed by (scope, name).
  • An automatic identity is prefixed with the scope when nested.
  • Root-level keys, explicit keys and versioned-hit precedence are unchanged.
  • No identity is added for opaque arguments. They stay on order-based keys, now per calling step, and replay_key remains the way to identify concurrent or reordered calls.

The other variant in suggestion 1, advancing the children's counters on a cache hit, would need the call counts stored with each cached entry. Entries from earlier releases lack them, so it would not help the legacy case. It was not prototyped.

What was run. agent-framework-core 1.22.0 and main at c4de86d. The four later commits on main (up to 63d0998) do not touch the key allocation, but nothing was re-run on them. Three variants of _functional.py:

  • main unchanged.
  • [BREAKING] Python: Reject unsafe functional step replay #9391: its hunks no longer apply (stale base, conflicts), so its three behavioural changes were ported by hand. Any legacy-key hit raises, checkpointing a step without a replay identity raises, and identity_is_replayable is derived from replay_identity. The port passes all 13 of the PR's new tests.
  • The scoped-counter change.

Legacy-format entries (name::idx) in these runs were written by hand, not by an older release.

# Scenario main #9391 (port) scoped counters
1 Issue repro: nested step, pydantic argument, crash and resume [['A','A']], inner body not run ValueError at the first checkpointed step, no checkpoint written [['A','B']], only tick(b) runs
2 Same, HITL pause and resume, no checkpoint storage [['A','A','yes']] ValueError at the pause [['A','B','yes']]
3 Rows 1 and 2 with stream=True as rows 1 and 2 as rows 1 and 2 as rows 1 and 2
4 Nested step, str arguments, legacy-format entries [['A','A']] ValueError [['A','B']]
5 Nested step, explicit replay_key, legacy-format entries [['A','A']] ValueError [['A','B']]
6 Parent restored from cache, then the inner step called directly from the workflow (new checkpoint, pydantic argument) wrong ValueError at start correct
7 Same as 6, legacy-format entries wrong ValueError wrong (see below)
8 Automatic identity, identical arguments from two parents, non-deterministic inner body [[1,1]] [[1,1]] [[1,2]]
9 Flat legacy checkpoint (compute::0, pinned by test_legacy_checkpoint_replays_sequential_step) replays ValueError replays
10 Flat step with a pydantic argument under checkpointing or HITL works ValueError works

Test files.

Compatibility boundary.

  • Flat legacy checkpoints replay as before, and the pinned test passes.
  • Nested keys gain the scope, so inner entries that an earlier release recorded for a nested call are no longer matched. Those inner steps run once more on such a resume: re-execution instead of aliasing. This was run for legacy-format entries. For automatic identities it follows from the key shape and was not run. A read-only fallback to the old nested keys would avoid the re-run, at the cost of keeping the aliasing for those entries.
  • Residual case (row 7): in a checkpoint written before the change, a workflow-level call can still claim a legacy entry that a nested call wrote. Legacy keys carry no record of the caller, so no key scheme can tell them apart. New checkpoints are not affected (row 6).

Not covered yet: concurrent sibling calls under different parents, a nested Workflow.run inside a step, as_agent, and checkpoint storage other than in-memory. No new tests have been written for the change.

Sketch of the change (against c4de86d)
--- a/python/packages/core/agent_framework/_workflows/_functional.py
+++ b/python/packages/core/agent_framework/_workflows/_functional.py
@@ module level, after _active_run_ctx @@
 _active_run_ctx: ContextVar[RunContext | None] = ContextVar("_active_run_ctx", default=None)
+# Scope of the running step body (hash of its cache key); empty at workflow level.
+_active_step_scope: ContextVar[str] = ContextVar("_active_step_scope", default="")
@@ RunContext.__init__ @@
-        self._step_call_counters: dict[str, int] = {}
+        self._step_call_counters: dict[tuple[str, str], int] = {}
@@ RunContext._get_legacy_step_cache_key, _get_automatic_step_cache_key @@
     def _get_legacy_step_cache_key(self, step_name: str) -> str:
-        idx = self._step_call_counters.get(step_name, 0)
-        self._step_call_counters[step_name] = idx + 1
-        return f"{step_name}::{idx}"
+        scope = _active_step_scope.get()
+        idx = self._step_call_counters.get((scope, step_name), 0)
+        self._step_call_counters[(scope, step_name)] = idx + 1
+        return f"{scope}/{step_name}::{idx}" if scope else f"{step_name}::{idx}"

     def _get_automatic_step_cache_key(self, step_name: str, wrapper_identity: str, identity: str) -> str:
+        scope = _active_step_scope.get()
+        if scope:
+            identity = f"{scope}/{identity}"
         counter_key = (step_name, wrapper_identity, identity)
@@ StepWrapper.__call__ @@
             auto_request_info_index_before = ctx._auto_request_info_index
+            scope_token = _active_step_scope.set(hashlib.sha256(cache_key.encode("utf-8")).hexdigest()[:16])
             try:
                 result = await self._func(*call_args, **call_kwargs)
             except Exception as exc:
                 ...
                 raise
+            finally:
+                _active_step_scope.reset(scope_token)
@@ the workflow's _execute @@
         token = _active_run_ctx.set(ctx)
+        scope_reset_token = _active_step_scope.set("")
         release_task_tracking = _track_workflow_tasks()
         ...
         finally:
             release_task_tracking()
+            _active_step_scope.reset(scope_reset_token)
             _active_run_ctx.reset(token)

This branch was previously deployed

1 inactive deployment
github-app-auth — 9bb14b35 Deployed Oct 14, 2026 by atty57 via add_label #25394
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

breaking change Usage: [PRs], Target: all PRs that introduce changes that are not backward compatible python Usage: [Issues, PRs], Target: Python

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

3 participants