.NET: defer checkpoint cleanup until workflow restoration succeeds

[Yashvant Mahadev Hange (YashvantHange)]

Motivation & Context

FoundryJsonCheckpointStore.RetrieveCheckpointAsync prunes a checkpoint's ancestry as part of reading it. Reading is not the same as resuming: after that call, CheckpointManagerImpl.LookupCheckpointAsync still has to deserialize the stored JSON, InProcessRunner.RestoreCheckpointCoreAsync still has to check the checkpoint against the workflow being resumed, and the runner still has to import workflow, executor and edge state.

If any of those fail, the ancestors are already gone. An incompatible workflow, corrupt checkpoint data, or a failing state import therefore destroys exactly the state the caller could otherwise have fallen back to, and retrying the restore erodes the ancestry a little further each time. Explicitly restoring an older checkpoint has the same effect.

Description & Review Guide

• What are the major changes?

  Cleanup moves to a point where the restore is known to have succeeded.

  ICheckpointRestorationObserver is a new opt-in interface in Microsoft.Agents.AI.Workflows.Checkpointing that a store implements alongside ICheckpointStore<TStoreObject>. It is a separate interface rather than a member on ICheckpointStore<TStoreObject> because that interface is public and the package multi-targets netstandard2.0 and net472, so a default interface method is not available and a required member would break every external implementer. A store that does not implement it is simply never notified.

  The internal ICheckpointManager gains a matching member, implemented by CheckpointManagerImpl<T>, InMemoryCheckpointManager, and the public CheckpointManager facade as an explicit interface implementation, so no public surface changes there.

  InProcessRunner.RestoreCheckpointCoreAsync raises the notification only after the workflow match, the three state imports and the executor OnCheckpointRestoredAsync callbacks have all completed. FoundryJsonCheckpointStore implements the interface and moves its existing PruneObsoleteCheckpointsAsync call there; RetrieveCheckpointAsync becomes a pure read.

  Two details worth a look. First, housekeeping must not fail a restore that already succeeded, so exceptions from the notification are recorded on the active span and not rethrown. ActivityExtensions.CaptureException is deliberately not used, because it also sets the span status to Error and the restore did not fail. Second, a store that threw may still have deleted checkpoints before throwing, so the runner re-reads the checkpoint index unless the store declined the notification outright. Without that, Checkpoints would keep offering checkpoints that no longer exist.

  The pruning logic itself is untouched: sibling branches and checkpoints committed after the resumed one are retained exactly as before.

• What are the impacts of these changes?

  Reading a checkpoint no longer deletes anything, so a failed restore leaves the ancestry recoverable. Release builds of Microsoft.Agents.AI.Workflows and Microsoft.Agents.AI.Foundry.Hosting are clean across all target frameworks with no warnings and no Package Validation diagnostics. Microsoft.Agents.AI.Workflows.UnitTests is 776/776 and Microsoft.Agents.AI.Foundry.Hosting.UnitTests is 453/453 on net10.0, and dotnet format --verify-no-changes reports no changes.

  Two behavioral notes for the record. ICheckpointRestorationObserver carries [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] even though its neighbour ICheckpointStore<TStoreObject> does not, on the grounds that it is new API in a released package and can graduate later; that does mean a store author implementing both needs an MAAI001 suppression for one and not the other. And an external caller invoking FoundryJsonCheckpointStore.RetrieveCheckpointAsync directly no longer gets the implicit prune. Every in-repo resume path goes through InProcessRunner.RestoreCheckpointCoreAsync, so no first-party consumer loses pruning, and pruning was never part of the read contract.

• What is the test coverage?

  The four cases the issue asks for, each asserting the store was not notified and so had no cue to clean up: an incompatible workflow, malformed checkpoint data, a failure raised from an executor during the restore, and restoring a non-latest checkpoint. Alongside those: a success case, a case where the observer throws and the restore still succeeds, and two cases where the observer really does delete the ancestry, covering the index re-read. On the store side, RetrieveCheckpointAsync is asserted to delete nothing, including across repeated reads, and the existing pruning tests were moved to the new entry point so their coverage is preserved.

  I confirmed the tests fail against the old ordering: moving the notification back to before the validation makes 5 of the 8 new runner tests fail, and defaulting the index re-read flag to false fails the prune-then-throw test.

  One gap, stated rather than implied: no test fails from inside the ImportStateAsync calls themselves. The executor-callback and malformed-data cases are the nearest proxies.

• What do you want reviewers to focus on?

  Whether a separate opt-in interface is the contract shape you want, and the naming. OnRestorationCompletedAsync is deliberately not OnCheckpointRestoredAsync, since Executor.OnCheckpointRestoredAsync already exists and runs during the restore where it can still fail it, whereas this one runs after and cannot. Renaming is cheap if you would rather it read differently.

Related Issue

Fixes #7796

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 is linked to an issue and there is no other open PR for this issue (see Related Issue above).
• 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.

[Copilot]

Pull request overview

Defers checkpoint pruning until workflow restoration succeeds, preserving recoverable ancestry after failed restores.

Changes:

• Adds an opt-in restoration observer contract.
• Notifies stores after state restoration and refreshes checkpoint indexes.
• Moves Foundry pruning from retrieval to post-restoration, with expanded tests.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
CheckpointRestorationObserverTests.cs	Tests restoration notification behavior.
FoundryJsonCheckpointStoreTests.cs	Tests deferred pruning and retention.
InProcessRunner.cs	Sends post-restore notifications and refreshes indexes.
CheckpointManager.cs	Forwards restoration notifications.
InMemoryCheckpointManager.cs	Implements a no-op notification.
ICheckpointRestorationObserver.cs	Defines the opt-in observer API.
ICheckpointManager.cs	Adds the internal notification contract.
CheckpointManagerImpl.cs	Dispatches notifications to supporting stores.
FoundryJsonCheckpointStore.cs	Moves pruning to successful restoration.

Suppressed comments (1)

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CheckpointRestorationObserverTests.cs:185

• Two consecutive <summary> elements document the same method, producing duplicate summary nodes in generated XML documentation. Consolidate them into a single summary.

    /// <summary>
    /// Builds the workflow the tests checkpoint and resume. The echo executor starts the run, so it is instantiated
    /// and therefore takes part in the restore, which is what lets a test fail the restore from inside an executor.
    /// </summary>
    /// <summary>

---

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

[Yashvant Mahadev Hange (YashvantHange)]

Pushed 833c73e addressing all three review points, including the suppressed one.

• The AddRange finding was a genuine bug and is fixed, with a regression test that reproduces it when the fix is reverted.
• The state-import gap I had disclosed in the description is now closed by a test that fails inside RunContext.ImportStateAsync rather than by proxy.
• The duplicate <summary> elements on BuildWorkflow are merged into one.

While rechecking I also found and fixed a portability break in the new test: string.Replace(string, string, StringComparison) does not exist on net472, so the test project failed to compile there. It now uses the two-argument overload, which is ordinal anyway.

[Manjunath Janardhan (manjunathshiva)]

This matches the option-A shape discussed on the issue (post-restore notification, prune moves out of the read path) — glad to see it picked up.

The separate opt-in interface is the right call over a virtual member on JsonCheckpointStore: given the netstandard2.0/net472 multi-targeting, a default interface member isn't available, and the external-implementer breakage concern is real. That constraint didn't come up in the issue discussion, so nice catch.

One thing that stood out while reading: the "observer threw but already deleted some entries" case forcing an index re-read is a subtle one — good that there's a dedicated test for it.

[Yashvant Mahadev Hange (YashvantHange)]

Pushed beb764540, bringing this onto current main. Three things main changed underneath it, none of which conflicted, so the merge was clean and the branch still would not have built.

Public API baseline. #7935 added the public API analyzers, so ICheckpointRestorationObserver now has to be declared in the Microsoft.Agents.AI.Workflows baseline or the build fails RS0016 once per target framework. Added to PublicAPI.Unshipped.txt for all five. Only the interface needs an entry: NotifyCheckpointRestoredAsync is on internal types or is an explicit implementation of an internal interface, and Foundry.Hosting is not a released package.

FluentAssertions. #7938 removed it from the tree over licensing and rewrote the other tests in this project. CheckpointRestorationObserverTests was the last file in the repo still importing it, so it no longer compiled. Rewritten in xUnit, matching the style that commit established. The reason strings survive as Assert.True's message argument where they were carrying real information.

A silence this branch introduced. OnRestorationCompletedAsync acquired the Foundry state store outside the try, so a store the binding could not hand over threw past the catch that reports a refused prune and into the runner's blanket suppression of everything this observer throws. That contradicted both the method's own docs and the interface contract saying an implementation reports its own failures, and it was new: the same failure used to surface loudly from RetrieveCheckpointAsync. The binding is now acquired inside the try, with a regression test that fails the binding rather than the index write.

One thing worth a decision rather than my assumption: the new interface is [Experimental(MAAI001)] while its sibling ICheckpointStore<T> is not, so a store implementing both suppresses the diagnostic for one half of the pair. That follows dotnet/AGENTS.md for new public API, but say the word if you would rather it match the sibling.