Fix suspension persistence to materialize app state before shutdown save

src/ReactiveUI/Interactions/Interaction.cs

@@ -170,13 +170,10 @@ protected Func<IInteractionContext<TInput, TOutput>, IObservable<Unit>>[] GetHan
     protected virtual IOutputContext<TInput, TOutput> GenerateContext(TInput input) => new InteractionContext<TInput, TOutput>(input);
 
     /// <summary>
-    /// Yields once so asynchronous handlers are not invoked inside the current scheduler trampoline.
+    /// Yields through the default task scheduler so asynchronous handlers are not invoked inside the current scheduler trampoline.
     /// </summary>
-    /// <returns>A task that completes after the current context has yielded.</returns>
-    private static async Task YieldToCurrentContext()
-    {
-        await Task.Yield();
-    }
+    /// <returns>A task that completes after the current scheduler trampoline has yielded.</returns>
+    private static Task YieldToCurrentContext() => Task.Run(static () => { });
 
     /// <summary>
     /// Registers a normalized interaction handler.

src/ReactiveUI/Suspension/SuspensionHostExtensions.cs

@@ -200,7 +200,9 @@ public static IDisposable SetupDefaultSuspendResume(this ISuspensionHost item, I
                     .Subscribe(_ => item.Log().Info("Invalidated app state")));
 
         ret.Add(item.ShouldPersistState
-                    .SelectMany(x => _suspensionDriver.SaveState(item.AppState!).Finally(x.Dispose))
+                    .SelectMany(x => EnsureLoadAppStateOnce(item, _suspensionDriver)
+                        .SelectMany(_ => _suspensionDriver!.SaveState(item.AppState!))
+                        .Finally(x.Dispose))
                     .LoggedCatch(item, Observables.Unit, "Tried to persist app state")
                     .Subscribe(_ => item.Log().Info("Persisted application state")));
 
@@ -246,7 +248,9 @@ public static IDisposable SetupDefaultSuspendResume<TAppState>(this ISuspensionH
                     .Subscribe(_ => item.Log().Info("Invalidated app state")));
 
         ret.Add(item.ShouldPersistState
-                    .SelectMany(x => _suspensionDriver.SaveState(item.AppStateValue!, typeInfo).Finally(x.Dispose))
+                    .SelectMany(x => EnsureLoadAppStateOnce(item, _suspensionDriver, typeInfo)
+                        .SelectMany(_ => _suspensionDriver!.SaveState(item.AppStateValue!, typeInfo))
+                        .Finally(x.Dispose))
                     .LoggedCatch(item, Observables.Unit, "Tried to persist app state")
                     .Subscribe(_ => item.Log().Info("Persisted application state")));
 
@@ -301,6 +305,25 @@ private static IObservable<Unit> EnsureLoadAppState(this ISuspensionHost item, I
         return Observable.Return(Unit.Default);
     }
 
+    /// <summary>
+    /// Runs the pending one-time untyped app-state load, or materializes state directly if the pending loader has already been consumed.
+    /// </summary>
+    /// <param name="item">The suspension host.</param>
+    /// <param name="driver">The suspension driver.</param>
+    /// <returns>A completed observable.</returns>
+    [RequiresUnreferencedCode(
+        "This overload may invoke ISuspensionDriver.LoadState(), which is commonly reflection-based. " +
+        "Prefer EnsureLoadAppStateOnce<TAppState>(ISuspensionHost<TAppState>, ISuspensionDriver?, JsonTypeInfo<TAppState>) for trimming/AOT scenarios.")]
+    [RequiresDynamicCode(
+        "This overload may invoke ISuspensionDriver.LoadState(), which is commonly reflection-based. " +
+        "Prefer EnsureLoadAppStateOnce<TAppState>(ISuspensionHost<TAppState>, ISuspensionDriver?, JsonTypeInfo<TAppState>) for trimming/AOT scenarios.")]
+    private static IObservable<Unit> EnsureLoadAppStateOnce(ISuspensionHost item, ISuspensionDriver? driver)
+    {
+        var ensureLoadAppState = Interlocked.Exchange(ref _ensureLoadAppStateFunc, null);
+
+        return ensureLoadAppState?.Invoke() ?? item.EnsureLoadAppState(driver);
+    }
+
     /// <summary>
     /// Ensures a one-time typed app state load from storage using source-generated JSON metadata (trimming/AOT friendly).
     /// </summary>
@@ -341,4 +364,20 @@ private static IObservable<Unit> EnsureLoadAppState<TAppState>(this ISuspensionH
 
         return Observable.Return(Unit.Default);
     }
+
+    /// <summary>
+    /// Runs the pending one-time typed app-state load, or materializes state directly if the pending loader has already been consumed.
+    /// </summary>
+    /// <typeparam name="TAppState">The application state type.</typeparam>
+    /// <param name="item">The typed suspension host.</param>
+    /// <param name="driver">The suspension driver.</param>
+    /// <param name="typeInfo">Source-generated metadata for <typeparamref name="TAppState"/>.</param>
+    /// <returns>A completed observable.</returns>
+    private static IObservable<Unit> EnsureLoadAppStateOnce<TAppState>(ISuspensionHost<TAppState> item, ISuspensionDriver? driver, JsonTypeInfo<TAppState> typeInfo)
+        where TAppState : class
+    {
+        var ensureLoadAppState = Interlocked.Exchange(ref _ensureLoadAppStateFunc, null);
+
+        return ensureLoadAppState?.Invoke() ?? item.EnsureLoadAppState(driver, typeInfo);
+    }
 }

src/tests/ReactiveUI.Tests/InteractionsTest.cs

@@ -225,6 +225,130 @@ public async Task ObservableHandlersShouldNotBlockNestedInteractionsBeforeReturn
         await Assert.That(nestedHandledBeforeParentReturned).IsTrue();
     }
 
+    /// <summary>
+    ///     Tests that task handler exceptions are propagated to the interaction observer.
+    /// </summary>
+    /// <returns>A <see cref="Task" /> representing the asynchronous operation.</returns>
+    [Test]
+    public async Task TaskHandlerExceptionsShouldPropagate()
+    {
+        var interaction = new Interaction<Unit, string>();
+        var expected = new InvalidOperationException("task handler failed");
+
+        interaction.RegisterHandler(_ => Task.FromException(expected));
+
+        var ex = await Assert.ThrowsAsync<InvalidOperationException>(() => interaction.Handle(Unit.Default).ToTask());
+        await Assert.That(ex).IsSameReferenceAs(expected);
+    }
+
+    /// <summary>
+    ///     Tests that task handlers can complete without handling the interaction.
+    /// </summary>
+    /// <returns>A <see cref="Task" /> representing the asynchronous operation.</returns>
+    [Test]
+    public async Task TaskHandlersThatCompleteWithoutOutputShouldFallThroughToNextHandler()
+    {
+        var interaction = new Interaction<Unit, string>();
+
+        interaction.RegisterHandler(static context => context.SetOutput("fallback"));
+        interaction.RegisterHandler(static _ => Task.CompletedTask);
+
+        var result = await interaction.Handle(Unit.Default);
+
+        await Assert.That(result).IsEqualTo("fallback");
+    }
+
+    /// <summary>
+    ///     Tests that task handlers which do not set output still surface the unhandled interaction.
+    /// </summary>
+    /// <returns>A <see cref="Task" /> representing the asynchronous operation.</returns>
+    [Test]
+    public async Task TaskHandlersThatCompleteWithoutOutputShouldCauseUnhandledInteractionException()
+    {
+        var interaction = new Interaction<string, Unit>();
+
+        interaction.RegisterHandler(static _ => Task.CompletedTask);
+
+        var ex = await Assert.ThrowsAsync<UnhandledInteractionException<string, Unit>>(() =>
+            interaction.Handle("task").ToTask());
+
+        using (Assert.Multiple())
+        {
+            await Assert.That(ex!.Interaction).IsSameReferenceAs(interaction);
+            await Assert.That(ex.Input).IsEqualTo("task");
+        }
+    }
+
+    /// <summary>
+    ///     Tests that exceptions thrown while creating observable handlers are propagated.
+    /// </summary>
+    /// <returns>A <see cref="Task" /> representing the asynchronous operation.</returns>
+    [Test]
+    public async Task ObservableHandlerFactoryExceptionsShouldPropagate()
+    {
+        var interaction = new Interaction<Unit, string>();
+        var expected = new InvalidOperationException("observable handler factory failed");
+
+        interaction.RegisterHandler<Unit>(_ => throw expected);
+
+        var ex = await Assert.ThrowsAsync<InvalidOperationException>(() => interaction.Handle(Unit.Default).ToTask());
+        await Assert.That(ex).IsSameReferenceAs(expected);
+    }
+
+    /// <summary>
+    ///     Tests that errors produced by observable handlers are propagated.
+    /// </summary>
+    /// <returns>A <see cref="Task" /> representing the asynchronous operation.</returns>
+    [Test]
+    public async Task ObservableHandlerErrorsShouldPropagate()
+    {
+        var interaction = new Interaction<Unit, string>();
+        var expected = new InvalidOperationException("observable handler failed");
+
+        interaction.RegisterHandler(_ => Observable.Throw<Unit>(expected));
+
+        var ex = await Assert.ThrowsAsync<InvalidOperationException>(() => interaction.Handle(Unit.Default).ToTask());
+        await Assert.That(ex).IsSameReferenceAs(expected);
+    }
+
+    /// <summary>
+    ///     Tests that observable handlers can complete without handling the interaction.
+    /// </summary>
+    /// <returns>A <see cref="Task" /> representing the asynchronous operation.</returns>
+    [Test]
+    public async Task ObservableHandlersThatCompleteWithoutOutputShouldFallThroughToNextHandler()
+    {
+        var interaction = new Interaction<Unit, string>();
+
+        interaction.RegisterHandler(static context => context.SetOutput("fallback"));
+        interaction.RegisterHandler(static _ => Observable.Empty<Unit>());
+
+        var result = await interaction.Handle(Unit.Default);
+
+        await Assert.That(result).IsEqualTo("fallback");
+    }
+
+    /// <summary>
+    ///     Tests that observable handlers which do not set output still surface the unhandled interaction.
+    /// </summary>
+    /// <returns>A <see cref="Task" /> representing the asynchronous operation.</returns>
+    [Test]
+    public async Task ObservableHandlersThatCompleteWithoutOutputShouldCauseUnhandledInteractionException()
+    {
+        var interaction = new Interaction<string, Unit>();
+
+        interaction.RegisterHandler(static _ => Observable.Empty<Unit>());
+
+        var ex = await Assert.ThrowsAsync<UnhandledInteractionException<string, Unit>>(() =>
+            interaction.Handle("observable").ToTask());
+
+        using (Assert.Multiple())
+        {
+            await Assert.That(ex!.Interaction).IsSameReferenceAs(interaction);
+            await Assert.That(ex.Input).IsEqualTo("observable");
+        }
+    }
+
     /// <summary>
     ///     Tests that handlers can opt not to handle the interaction.
     /// </summary>

src/tests/ReactiveUI.Tests/Suspension/SuspensionHostExtensionsAotTests.cs

@@ -254,6 +254,135 @@ public async Task SetupDefaultSuspendResume_Typed_ShouldPersistState_CallsDriver
         await Assert.That(driver.LastSavedState).IsSameReferenceAs(appState);
     }
 
+    [Test]
+    public async Task SetupDefaultSuspendResume_Typed_ShouldPersistCreatedState_WhenNoPersistedStateAndPersistOccursBeforeGetAppState()
+    {
+        var createdState = new TestAppState { Value = 321 };
+        var createNewAppStateCallCount = 0;
+        var persistTokenDisposed = false;
+        using var host = new SuspensionHost<TestAppState>
+        {
+            CreateNewAppStateTyped = () =>
+            {
+                createNewAppStateCallCount++;
+                return createdState;
+            },
+            IsLaunchingNew = Observable.Never<Unit>(),
+            IsResuming = Observable.Never<Unit>(),
+            ShouldInvalidateState = Observable.Never<Unit>()
+        };
+
+        var driver = new TestSuspensionDriver<TestAppState>();
+        var persistSubject = new Subject<IDisposable>();
+        host.ShouldPersistState = persistSubject.ObserveOn(ImmediateScheduler.Instance);
+
+        using var disposable = host.SetupDefaultSuspendResume(TestAppStateContext.Default.TestAppState, driver);
+        var persistToken = Disposable.Create(() => persistTokenDisposed = true);
+
+        persistSubject.OnNext(persistToken);
+
+        await Assert.That(driver.LoadStateCallCount).IsEqualTo(1);
+        await Assert.That(createNewAppStateCallCount).IsEqualTo(1);
+        await Assert.That(host.AppStateValue).IsSameReferenceAs(createdState);
+        await Assert.That(driver.SaveStateCallCount).IsEqualTo(1);
+        await Assert.That(driver.LastSavedState).IsSameReferenceAs(createdState);
+        await Assert.That(persistTokenDisposed).IsTrue();
+    }
+
+    [Test]
+    public async Task SetupDefaultSuspendResume_Typed_ShouldPersistCreatedState_WhenLaunchSignalWasRaisedBeforeSetup()
+    {
+        var createdState = new TestAppState { Value = 654 };
+        var createNewAppStateCallCount = 0;
+        using var host = new SuspensionHost<TestAppState>
+        {
+            CreateNewAppStateTyped = () =>
+            {
+                createNewAppStateCallCount++;
+                return createdState;
+            },
+            ShouldInvalidateState = Observable.Never<Unit>()
+        };
+
+        var launchSubject = new Subject<Unit>();
+        var resumeSubject = new Subject<Unit>();
+        var persistSubject = new Subject<IDisposable>();
+        host.IsLaunchingNew = launchSubject.ObserveOn(ImmediateScheduler.Instance);
+        host.IsResuming = resumeSubject.ObserveOn(ImmediateScheduler.Instance);
+        host.ShouldPersistState = persistSubject.ObserveOn(ImmediateScheduler.Instance);
+
+        launchSubject.OnNext(Unit.Default);
+
+        var driver = new TestSuspensionDriver<TestAppState>();
+        using var disposable = host.SetupDefaultSuspendResume(TestAppStateContext.Default.TestAppState, driver);
+
+        persistSubject.OnNext(Disposable.Empty);
+
+        await Assert.That(driver.LoadStateCallCount).IsEqualTo(1);
+        await Assert.That(createNewAppStateCallCount).IsEqualTo(1);
+        await Assert.That(host.AppStateValue).IsSameReferenceAs(createdState);
+        await Assert.That(driver.SaveStateCallCount).IsEqualTo(1);
+        await Assert.That(driver.LastSavedState).IsSameReferenceAs(createdState);
+    }
+
+    [Test]
+    public async Task SetupDefaultSuspendResume_Typed_ShouldPersistLoadedState_WhenPersistOccursBeforeGetAppState()
+    {
+        var loadedState = new TestAppState { Value = 987 };
+        var createNewAppStateCallCount = 0;
+        using var host = new SuspensionHost<TestAppState>
+        {
+            CreateNewAppStateTyped = () =>
+            {
+                createNewAppStateCallCount++;
+                return new TestAppState();
+            },
+            IsLaunchingNew = Observable.Never<Unit>(),
+            IsResuming = Observable.Never<Unit>(),
+            ShouldInvalidateState = Observable.Never<Unit>()
+        };
+
+        var driver = new TestSuspensionDriver<TestAppState> { StateToLoad = loadedState };
+        var persistSubject = new Subject<IDisposable>();
+        host.ShouldPersistState = persistSubject.ObserveOn(ImmediateScheduler.Instance);
+
+        using var disposable = host.SetupDefaultSuspendResume(TestAppStateContext.Default.TestAppState, driver);
+
+        persistSubject.OnNext(Disposable.Empty);
+
+        await Assert.That(driver.LoadStateCallCount).IsEqualTo(1);
+        await Assert.That(createNewAppStateCallCount).IsEqualTo(0);
+        await Assert.That(host.AppStateValue).IsSameReferenceAs(loadedState);
+        await Assert.That(driver.SaveStateCallCount).IsEqualTo(1);
+        await Assert.That(driver.LastSavedState).IsSameReferenceAs(loadedState);
+    }
+
+    [Test]
+    public async Task SetupDefaultSuspendResume_Typed_ShouldDisposePersistTokenAfterSave()
+    {
+        var appState = new TestAppState { Value = 111 };
+        var persistTokenDisposed = false;
+        using var host = new SuspensionHost<TestAppState>
+        {
+            AppStateValue = appState,
+            IsLaunchingNew = Observable.Never<Unit>(),
+            IsResuming = Observable.Never<Unit>(),
+            ShouldInvalidateState = Observable.Never<Unit>()
+        };
+
+        var driver = new TestSuspensionDriver<TestAppState>();
+        var persistSubject = new Subject<IDisposable>();
+        host.ShouldPersistState = persistSubject.ObserveOn(ImmediateScheduler.Instance);
+
+        using var disposable = host.SetupDefaultSuspendResume(TestAppStateContext.Default.TestAppState, driver);
+        var persistToken = Disposable.Create(() => persistTokenDisposed = true);
+
+        persistSubject.OnNext(persistToken);
+
+        await Assert.That(driver.SaveStateCallCount).IsEqualTo(1);
+        await Assert.That(persistTokenDisposed).IsTrue();
+    }
+
     [Test]
     public async Task SetupDefaultSuspendResume_Typed_ShouldInvalidateState_CallsDriverInvalidateState()
     {