[ShiYuan Xiong (hehe764)]

Motivation & Context

A functional workflow can lose a completed step's in-flight checkpoint when a parallel sibling fails. Resuming then repeats the completed step and its side effects, depending on when the storage adapter awaits I/O. Fixes #9615.

Description & Review Guide

• What are the major changes?

  • Run each per-step checkpoint save as a tracked task. Shield the normal and drain awaits, and protect storage-created descendant tasks during framework-owned failure cleanup.
  • Preserve caller cancellation: when the caller cancels the run, cancel in-flight saves before draining workflow tasks. The same distinction applies during human-input pause settlement.
  • Add three sibling-failure/resume regression cases for inline, task-based, and gather-based storage, plus 18 cancellation cases covering direct cancellation and asyncio.wait_for timeouts during step execution, failure cleanup, and pause settlement.

• What is the impact of these changes?

  • Framework-owned cleanup after a sibling failure lets completed steps finish saving. When storage succeeds, a resumed run can replay those results without repeating the completed steps.
  • Caller cancellation can interrupt in-flight saves with cooperative storage adapters; a hung save no longer necessarily holds the caller until it finishes. An explicit caller cancellation may therefore discard an in-flight checkpoint, as before this PR.
  • This is not a durability guarantee across storage errors or process termination. Framework-owned cleanup still waits for saves when the caller has not cancelled; adapters that suppress cancellation remain outside the cancellation guarantee.

• What do you want reviewers to focus on?

  • Whether the separation between caller cancellation and framework-owned failure cleanup is correct, including human-input pause settlement.
  • Task-parent tracking and protection of storage-created descendants, while still cancelling unrelated workflow tasks.
  • The regression tests' coverage of checkpoint preservation and prompt caller cancellation.

Validation

Recorded local results for the caller-cancellation follow-up 35dd2c8, on Windows / Python 3.13.16:

• Full python/packages/core/tests/workflow/ suite: 1,300 passed, 6 skipped, 1 xfailed. The xfail is the existing test_run_request_with_full_history_clears_service_session_id, tracking the executor-layer follow-up to Python: [Bug]: "Duplicate item found" error when using AzureAIClient with tools in workflows #3295; its marker was not added or changed by this PR.
• Of the 18 caller-cancellation regression cases, 12 failed before the cancellation fix and six existing-behavior controls passed; all 18 pass after the fix. The three sibling-failure checkpoint/resume cases also pass.
• Ruff lint/format and strict source Pyright passed. The changed test file passed the repository-configured Pyright, mypy, pyrefly, ty, and zuban checks.

Head d363147 includes an upstream merge. A reviewer independently rechecked that head: a 3-second save no longer delays a 0.2-second caller timeout, and a hung save no longer holds the caller. This is focused reproduction evidence, not a claim that the complete suite was rerun on that merged head.

The complete core suite and live storage backends were not rerun. The held upstream test workflows still require maintainer approval; no fully green hosted CI result is claimed.

Follow-up on the Copilot snapshot-race review, fixed in 1e920c3: checkpoint construction now happens synchronously at step completion (_build_checkpoint), and only storage.save is scheduled, so a ready sibling cannot mutate state before the snapshot is taken. Regression test_step_checkpoint_snapshots_state_before_scheduling_save fails before the fix and passes after. Full tests/workflow/ suite on head b058ca2b: 1,304 passed, 6 skipped, 1 xfailed; Ruff and source Pyright are clean. Copilot's latest summary flagged two moderate residual risks in repeated-cancellation and pause-settlement cleanup without inline details. Added test_double_caller_cancellation_finishes_checkpoint_save and test_second_cancellation_during_pause_settlement_finishes_save; both pass on the existing implementation, so no code change was required.

Related Issue

Fixes #9615.

The issue has an automated reproduction and the proposed implementation approach. Explicit maintainer agreement on the direction, or a ready-for-implementation label, is still pending. The reviewer's cancellation verification is not presented as maintainer implementation approval.

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.

Tool used: OpenAI Codex.

AI assistance covered root-cause analysis, implementation, regression tests, local verification, and drafting the PR description and review replies. The validation section distinguishes recorded local checks, independent reviewer reproductions, and checks that remain unverified.

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 all contribution-process requirements — maintainer issue agreement is pending as noted above.
• 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.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[ShiYuan Xiong (hehe764)]

@microsoft-github-policy-service agree

[Somay kousis (Somay-kousis)]

The fix does what it says for the sibling-failure case. It also makes a run impossible to cancel while a step checkpoint save is in flight.

What I ran (PR head c6414f1, merge base edc2427, Python 3.12, macOS)

• test_functional_workflow.py on head: 213 passed.
• Same file with only _functional.py reverted to the merge base: 3 failed (all three test_completed_parallel_step_is_checkpointed_when_sibling_fails params), 210 passed. So the new test catches the bug.
• ruff 0.17.0 on both changed files: clean.

Finding: cancelling the run now waits for the save, and a hung save blocks the cancel forever

_functional.py:2182-2184: when the step task is cancelled, it catches the CancelledError and awaits asyncio.shield(save_task) again. _cancel_workflow_tasks (_functional.py:1087) skips save tasks when cancelling and then waits for them with gather (:1099). Neither path tells a cancel that comes from the framework (a sibling failed) apart from a cancel that comes from the caller. So asyncio.wait_for(wf.run(...), timeout) and task cancellation from a server shutting down both get stuck behind storage.save().

Repro: a single-step workflow whose storage save() sleeps before writing. The caller uses asyncio.wait_for(wf.run(0), timeout=0.2).

	save sleeps 3 s	save never returns
merge base	TimeoutError after 0.20 s	TimeoutError after 0.22 s
this PR	TimeoutError after 3.02 s	still blocked after 5 s guard

A save that never returns is a realistic case: a DB or blob call with no client timeout. Before this PR, wait_for or task.cancel() was the way out. Now the run can't be stopped from outside.

Suggestion: keep the drain for cancels the framework makes, but let a cancel from the caller reach the save. One option I tried, in the except BaseException branch of _execute (_functional.py:2365), before ctx._cancel_workflow_task(workflow_task):

for save_task in list(ctx._checkpoint_save_tasks):
    save_task.cancel()

With that change both repros return after about 0.2 s, and the whole test file still passes (213). That's my own attempt, not a reviewed fix. The cost is that a cancel from the caller drops the in-flight checkpoint again, which is how the merge base behaves. If you'd rather keep the save, a bounded wait (for example asyncio.wait_for(asyncio.shield(save_task), timeout) in the drain) would also stop the hang.

A regression test could build asyncio.wait_for(wf.run(0), 0.2) over a storage whose save awaits an asyncio.Event that is never set, and assert that it raises TimeoutError within about 1 s.

What I did not check

• HITL pause settlement (_settle_workflow_tasks) with a save in flight. I only traced it by reading: the save task becomes pause-owned through its parent step, so settlement should wait for it too. Not run.
• Python 3.10/3.11 and other OSes. Live storage backends (Cosmos, Redis, file) were not run.

[ShiYuan Xiong (hehe764)]

Addressed the caller-cancellation feedback in commit 35dd2c8.

• Caller cancellation now cancels in-flight checkpoint saves before draining the workflow task, so direct cancellation and asyncio.wait_for timeouts can terminate a run while a cooperative storage adapter is waiting.
• The same handling applies during human-input pause settlement.
• Framework-owned cleanup after a sibling failure still lets completed steps finish saving, including storage-created child tasks.

Added 18 regression cases covering inline/task/gather storage, direct cancellation/timeouts, and cancellation during step execution, failure cleanup, and pause settlement. Before the fix, 12 cases failed and 6 existing-behavior controls passed; all 18 now pass. The original three sibling-failure checkpoint/resume cases still pass.

Validation on Windows / Python 3.13.16:

• Full packages/core/tests/workflow/ suite: 1300 passed, 6 skipped, 1 xfailed. The xfail is the existing test_run_request_with_full_history_clears_service_session_id, tracking the executor-layer follow-up to Python: [Bug]: "Duplicate item found" error when using AzureAIClient with tools in workflows #3295; this PR did not add or change its marker.
• Ruff check/format and strict source Pyright: passed.
• The changed test file passes the repository-configured Pyright, mypy, pyrefly, ty, and zuban checks.

The full core suite and live storage backends were not rerun. Storage adapters that suppress cancellation remain outside the cancellation guarantee. This supersedes the earlier description's statement that an in-flight save necessarily prevents caller cancellation.

[Copilot]

🟡 Changes recommended

Pause settlement must distinguish caller cancellation from framework-owned child cancellation before cancelling in-flight checkpoint saves.

1 open issue

• Distinguish caller cancellation from child cancellation during settlement · New

_{Review effort: Lite}

[ShiYuan Xiong (hehe764)]

Follow-up on the latest cancellation review, addressed in 1664ae4.

The pause-settlement path previously handled every escaping CancelledError as cancellation of the caller. That conflated two cases: cancellation delivered to the workflow run itself, and a completed pause-owned sibling propagating its own CancelledError. In the latter case, cleanup incorrectly cancelled in-flight checkpoint saves and could lose a completed sibling's checkpoint.

The settlement code now records when the cancellation originated from a settled child task. Only cancellation delivered to the settlement/run task cancels checkpoint saves; child-propagated cancellation continues through framework cleanup while completed-step saves are drained normally. Caller cancellation and asyncio.wait_for timeouts therefore remain prompt, while sibling cancellation no longer discards a completed checkpoint.

A deterministic regression test blocks a completed sibling's checkpoint save while another sibling requests input and a third raises CancelledError, then verifies that the save is not cancelled and is persisted. The full test_functional_workflow.py file passes (232 tests); Ruff, strict source Pyright, and mypy for the changed test file also pass.

[Copilot]

🟡 Changes recommended

A critical checkpoint snapshot race and a moderate pause-settlement cancellation bug remain unresolved.

1 open issue

• Snapshot checkpoint state synchronously before scheduling save · New

1 resolved since last review

• Distinguish caller cancellation from child cancellation during settlement

_{Review effort: Lite}

[Copilot]

🔵 Needs a closer look

Two moderate issues remain in repeated-cancellation and pause-settlement cleanup paths, where in-flight checkpoint saves may be abandoned.

No open issues 🎉

1 resolved since last review

• Snapshot checkpoint state synchronously before scheduling save

_{Review effort: Lite}

[Eduard van Valkenburg (eavanvalkenburg)]

ShiYuan Xiong (@hehe764) Please take a look at the two inline checkpoint-lifecycle findings.

[ShiYuan Xiong (hehe764)]

Eduard van Valkenburg (@eavanvalkenburg) Both checkpoint-lifecycle findings are addressed in 4695bbb; I have also replied to each inline comment with the implementation details.

• Cancellation throughout teardown: record caller cancellation on the run context, cancel existing saves, and skip per-step saves created by steps finishing during teardown. Shielding the cleanup drain lets it observe a new caller cancellation promptly, including after sibling-failure cleanup has begun, then cancel saves and propagate cancellation. Framework-owned sibling-failure cleanup still protects completed-step checkpoints.
• Nested-state snapshots: reuse _snapshot_workflow_state synchronously when building checkpoints, so nested dictionaries/lists are detached before the storage task is scheduled. The regression now mutates existing nested containers and checks the persisted checkpoints.

Before the fix, the six teardown-cancellation cases and two nested-state cases fail; all eight pass after the fix. Local validation on the pushed change:

• Windows / Python 3.13.16: workflow suite 1,312 passed, 6 skipped, 1 existing xfailed; functional-workflow file 243 passed.
• Windows / Python 3.10: functional-workflow file 224 passed, 19 skipped.
• Ruff lint/format, strict source Pyright, and the changed test file's configured Pyright, mypy, pyrefly, ty and zuban checks passed.

The full core suite and live storage backends were not rerun. The previous head's hosted CI passed; hosted CI for this new commit is a separate run.

[ShiYuan Xiong (hehe764)]

Eduard van Valkenburg (@eavanvalkenburg) I did a further cancellation/resource-lifecycle audit and pushed 99c10d8, covering two additional cleanup gaps:

1. Repeated caller cancellation could abandon background tasks. The caller-cancellation branch awaited the root task through a separate unshielded gather before reaching the context-wide drain. A second cancellation during storage cleanup could interrupt that await and skip draining unrelated workflow tasks. Caller cancellation now goes directly through the existing cancellation-aware context drain, which includes the root task and all tracked background tasks. I strengthened the repeated-cancellation regression to keep storage cleanup in flight and verify that background-task cleanup also finishes. This gap also reproduces with the upstream base implementation; it was not introduced by the preceding commit.
2. Cancellation before a save task starts could leave its coroutine unclosed. The task factory wraps the storage coroutine in a tracking coroutine. If the task is cancelled before the wrapper starts, the wrapped coroutine was never awaited or closed, resulting in a RuntimeWarning during garbage collection. The factory now closes an unstarted coroutine when its owning task is cancelled. A new regression cancels the run just before the save task can start and checks for unawaited-coroutine warnings.

Both focused regressions fail on 4695bbb and pass after this follow-up. Local verification on the pushed change:

• Windows / Python 3.13.16 workflow suite: 1,313 passed, 6 skipped, 1 existing xfailed; zero unexpected failures.
• Python 3.10 functional-workflow file: 225 passed, 19 skipped.
• Ruff lint/format, strict source Pyright, and the changed test file's Pyright, mypy, pyrefly, ty and zuban checks passed.
• Thirteen additional local audit cases passed, covering storage failures, TaskGroup cleanup with inline/task/gather storage, eager task factories, file-checkpoint persistence/resume without repeating a completed step, and cancellation isolation between concurrent workflows.

The complete core suite and remote storage services were not rerun. Hosted CI for this new commit remains separate from these local results.