Support persistent component state across enhanced page navigations

src/Components/Components/src/IPersistentComponentStateScenario.cs

@@ -0,0 +1,16 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.AspNetCore.Components;
+
+/// <summary>
+/// Represents a scenario for persistent component state restoration.
+/// </summary>
+public interface IPersistentComponentStateScenario
+{
+    /// <summary>
+    /// Gets a value indicating whether callbacks for this scenario can be invoked multiple times.
+    /// If false, callbacks are automatically unregistered after first invocation.
+    /// </summary>
+    bool IsRecurring { get; }
+}

src/Components/Components/src/IPersistentStateFilter.cs

@@ -0,0 +1,17 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.AspNetCore.Components;
+
+/// <summary>
+/// Defines filtering logic for persistent component state restoration.
+/// </summary>
+public interface IPersistentStateFilter
+{
+    /// <summary>
+    /// Determines whether state should be restored for the given scenario.
+    /// </summary>
+    /// <param name="scenario">The restoration scenario.</param>
+    /// <returns>True if state should be restored; otherwise false.</returns>
+    bool ShouldRestore(IPersistentComponentStateScenario scenario);
+}

src/Components/Components/src/PersistentComponentState.cs

@@ -16,6 +16,7 @@ public class PersistentComponentState
     private readonly IDictionary<string, byte[]> _currentState;
 
     private readonly List<PersistComponentStateRegistration> _registeredCallbacks;
+    private readonly List<RestoreComponentStateRegistration> _restoringCallbacks = new();
 
     internal PersistentComponentState(
         IDictionary<string , byte[]> currentState,
@@ -27,6 +28,11 @@ internal PersistentComponentState(
 
     internal bool PersistingState { get; set; }
 
+    /// <summary>
+    /// Gets the current restoration scenario, if any.
+    /// </summary>
+    public IPersistentComponentStateScenario? CurrentScenario { get; internal set; }
+
     internal void InitializeExistingState(IDictionary<string, byte[]> existingState)
     {
         if (_existingState != null)
@@ -155,6 +161,73 @@ internal bool TryTakeFromJson(string key, [DynamicallyAccessedMembers(JsonSerial
         }
     }
 
+    /// <summary>
+    /// Registers a callback to be invoked when state is restored and the filter allows the current scenario.
+    /// </summary>
+    /// <param name="filter">The filter to determine if the callback should be invoked for a scenario.</param>
+    /// <param name="callback">The callback to invoke during restoration.</param>
+    /// <returns>A subscription that can be disposed to unregister the callback.</returns>
+    public RestoringComponentStateSubscription RegisterOnRestoring(
+        IPersistentStateFilter filter,
+        Action callback)
+    {
+        ArgumentNullException.ThrowIfNull(filter);
+        ArgumentNullException.ThrowIfNull(callback);
+
+        var registration = new RestoreComponentStateRegistration(filter, callback);
+        _restoringCallbacks.Add(registration);
+
+        // If we already have a current scenario and the filter matches, invoke immediately
+        if (CurrentScenario != null && filter.ShouldRestore(CurrentScenario))
+        {
+            callback();
+        }
+
+        return new RestoringComponentStateSubscription(_restoringCallbacks, filter, callback);
+    }
+
+    /// <summary>
+    /// Updates the existing state with new state for subsequent restoration calls.
+    /// Only allowed when existing state is empty (fully consumed).
+    /// </summary>
+    /// <param name="newState">New state dictionary to replace existing state.</param>
+    /// <param name="scenario">The restoration scenario context.</param>
+    internal void UpdateExistingState(IDictionary<string, byte[]> newState, IPersistentComponentStateScenario scenario)
+    {
+        ArgumentNullException.ThrowIfNull(newState);
+        ArgumentNullException.ThrowIfNull(scenario);
+
+        if (_existingState != null && _existingState.Count > 0)
+        {
+            throw new InvalidOperationException("Cannot update existing state when state dictionary is not empty. State must be fully consumed before updating.");
+        }
+
+        _existingState = newState;
+        CurrentScenario = scenario;
+
+        // Invoke matching restoration callbacks
+        InvokeRestoringCallbacks(scenario);
+    }
+
+    private void InvokeRestoringCallbacks(IPersistentComponentStateScenario scenario)
+    {
+        for (int i = _restoringCallbacks.Count - 1; i >= 0; i--)
+        {
+            var registration = _restoringCallbacks[i];
+
+            if (registration.Filter.ShouldRestore(scenario))
+            {
+                registration.Callback();
+
+                // Remove callback if scenario is not recurring (one-time scenarios)
+                if (!scenario.IsRecurring)
+                {
+                    _restoringCallbacks.RemoveAt(i);
+                }
+            }
+        }
+    }
+
     private bool TryTake(string key, out byte[]? value)
     {
         ArgumentNullException.ThrowIfNull(key);

src/Components/Components/src/PersistentState/ComponentStatePersistenceManager.cs

@@ -17,6 +17,7 @@ public class ComponentStatePersistenceManager
     private bool _stateIsPersisted;
     private readonly PersistentServicesRegistry? _servicesRegistry;
     private readonly Dictionary<string, byte[]> _currentState = new(StringComparer.Ordinal);
+    private bool _isFirstRestore = true;
 
     /// <summary>
     /// Initializes a new instance of <see cref="ComponentStatePersistenceManager"/>.
@@ -55,9 +56,43 @@ public ComponentStatePersistenceManager(ILogger<ComponentStatePersistenceManager
     /// <param name="store">The <see cref="IPersistentComponentStateStore"/> to restore the application state from.</param>
     /// <returns>A <see cref="Task"/> that will complete when the state has been restored.</returns>
     public async Task RestoreStateAsync(IPersistentComponentStateStore store)
+    {
+        await RestoreStateAsync(store, scenario: null);
+    }
+
+    /// <summary>
+    /// Restores component state from the given store with scenario context.
+    /// </summary>
+    /// <param name="store">The store to restore state from.</param>
+    /// <param name="scenario">The restoration scenario context.</param>
+    /// <returns>A task that completes when state restoration is finished.</returns>
+    public async Task RestoreStateAsync(
+        IPersistentComponentStateStore store, 
+        IPersistentComponentStateScenario? scenario)
     {
         var data = await store.GetPersistedStateAsync();
-        State.InitializeExistingState(data);
+
+        if (_isFirstRestore)
+        {
+            // First-time initialization
+            State.InitializeExistingState(data);
+            _isFirstRestore = false;
+        }
+        else
+        {
+            // Scenario-based update - only if we have a scenario
+            if (scenario != null)
+            {
+                State.UpdateExistingState(data, scenario);
+            }
+            else
+            {
+                // This is a second call without a scenario, which should fail
+                // (maintaining the original behavior for backward compatibility)
+                throw new InvalidOperationException("PersistentComponentState already initialized.");
+            }
+        }
+
         _servicesRegistry?.Restore(State);
     }
 

src/Components/Components/src/PublicAPI.Unshipped.txt

@@ -21,3 +21,14 @@ static Microsoft.AspNetCore.Components.Infrastructure.ComponentsMetricsServiceCo
 static Microsoft.AspNetCore.Components.Infrastructure.ComponentsMetricsServiceCollectionExtensions.AddComponentsTracing(Microsoft.Extensions.DependencyInjection.IServiceCollection! services) -> Microsoft.Extensions.DependencyInjection.IServiceCollection!
 static Microsoft.Extensions.DependencyInjection.SupplyParameterFromPersistentComponentStateProviderServiceCollectionExtensions.AddSupplyValueFromPersistentComponentStateProvider(this Microsoft.Extensions.DependencyInjection.IServiceCollection! services) -> Microsoft.Extensions.DependencyInjection.IServiceCollection!
 virtual Microsoft.AspNetCore.Components.Rendering.ComponentState.GetComponentKey() -> object?
+Microsoft.AspNetCore.Components.IPersistentComponentStateScenario
+Microsoft.AspNetCore.Components.IPersistentComponentStateScenario.IsRecurring.get -> bool
+Microsoft.AspNetCore.Components.IPersistentStateFilter
+Microsoft.AspNetCore.Components.IPersistentStateFilter.ShouldRestore(Microsoft.AspNetCore.Components.IPersistentComponentStateScenario! scenario) -> bool
+Microsoft.AspNetCore.Components.PersistentComponentState.CurrentScenario.get -> Microsoft.AspNetCore.Components.IPersistentComponentStateScenario?
+*REMOVED*Microsoft.AspNetCore.Components.PersistentComponentState.RegisterOnRestoring(Microsoft.AspNetCore.Components.IPersistentComponentStateScenario! scenario, System.Action! callback) -> Microsoft.AspNetCore.Components.RestoringComponentStateSubscription
+Microsoft.AspNetCore.Components.PersistentComponentState.RegisterOnRestoring(Microsoft.AspNetCore.Components.IPersistentStateFilter! filter, System.Action! callback) -> Microsoft.AspNetCore.Components.RestoringComponentStateSubscription
+Microsoft.AspNetCore.Components.RestoringComponentStateSubscription
+Microsoft.AspNetCore.Components.RestoringComponentStateSubscription.RestoringComponentStateSubscription() -> void
+Microsoft.AspNetCore.Components.RestoringComponentStateSubscription.Dispose() -> void
+Microsoft.AspNetCore.Components.Infrastructure.ComponentStatePersistenceManager.RestoreStateAsync(Microsoft.AspNetCore.Components.IPersistentComponentStateStore! store, Microsoft.AspNetCore.Components.IPersistentComponentStateScenario? scenario) -> System.Threading.Tasks.Task!

src/Components/Components/src/Reflection/PropertyGetter.cs

@@ -14,12 +14,16 @@ internal sealed class PropertyGetter
 
     private readonly Func<object, object?> _GetterDelegate;
 
+    public PropertyInfo PropertyInfo { get; }
+
     [UnconditionalSuppressMessage(
         "ReflectionAnalysis",
         "IL2060:MakeGenericMethod",
         Justification = "The referenced methods don't have any DynamicallyAccessedMembers annotations. See https://github.com/mono/linker/issues/1727")]
     public PropertyGetter(Type targetType, PropertyInfo property)
     {
+        PropertyInfo = property;
+
         if (property.GetMethod == null)
         {
             throw new InvalidOperationException("Cannot provide a value for property " +

src/Components/Components/src/RestoreComponentStateRegistration.cs

@@ -0,0 +1,19 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.AspNetCore.Components;
+
+/// <summary>
+/// Represents a registration for state restoration callbacks.
+/// </summary>
+internal readonly struct RestoreComponentStateRegistration
+{
+    public RestoreComponentStateRegistration(IPersistentStateFilter filter, Action callback)
+    {
+        Filter = filter;
+        Callback = callback;
+    }
+
+    public IPersistentStateFilter Filter { get; }
+    public Action Callback { get; }
+}

src/Components/Components/src/RestoringComponentStateSubscription.cs

@@ -0,0 +1,41 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.AspNetCore.Components;
+
+/// <summary>
+/// Represents a subscription to state restoration notifications.
+/// </summary>
+public readonly struct RestoringComponentStateSubscription : IDisposable
+{
+    private readonly List<RestoreComponentStateRegistration>? _callbacks;
+    private readonly IPersistentStateFilter? _filter;
+    private readonly Action? _callback;
+
+    internal RestoringComponentStateSubscription(
+        List<RestoreComponentStateRegistration> callbacks,
+        IPersistentStateFilter filter,
+        Action callback)
+    {
+        _callbacks = callbacks;
+        _filter = filter;
+        _callback = callback;
+    }
+
+    /// <inheritdoc />
+    public void Dispose()
+    {
+        if (_callbacks != null && _filter != null && _callback != null)
+        {
+            for (int i = _callbacks.Count - 1; i >= 0; i--)
+            {
+                var registration = _callbacks[i];
+                if (ReferenceEquals(registration.Filter, _filter) && ReferenceEquals(registration.Callback, _callback))
+                {
+                    _callbacks.RemoveAt(i);
+                    break;
+                }
+            }
+        }
+    }
+}

src/Components/Components/src/SupplyParameterFromPersistentComponentStateValueProvider.cs

@@ -21,6 +21,7 @@ internal sealed class SupplyParameterFromPersistentComponentStateValueProvider(P
     private static readonly ConcurrentDictionary<(Type, string), PropertyGetter> _propertyGetterCache = new();
 
     private readonly Dictionary<ComponentState, PersistingComponentStateSubscription> _subscriptions = [];
+    private readonly Dictionary<(ComponentState, string), object?> _scenarioRestoredValues = [];
 
     public bool IsFixed => false;
     // For testing purposes only
@@ -40,8 +41,15 @@ public bool CanSupplyValue(in CascadingParameterInfo parameterInfo)
     public object? GetCurrentValue(object? key, in CascadingParameterInfo parameterInfo)
     {
         var componentState = (ComponentState)key!;
+
+        // Check if we have a scenario-restored value first
+        var valueKey = (componentState, parameterInfo.PropertyName);
+        if (_scenarioRestoredValues.TryGetValue(valueKey, out var scenarioValue))
+        {
+            return scenarioValue;
+        }
+
         var storageKey = ComputeKey(componentState, parameterInfo.PropertyName);
-
         return state.TryTakeFromJson(storageKey, parameterInfo.PropertyType, out var value) ? value : null;
     }
 
@@ -52,10 +60,12 @@ public void Subscribe(ComponentState subscriber, in CascadingParameterInfo param
     {
         var propertyName = parameterInfo.PropertyName;
         var propertyType = parameterInfo.PropertyType;
+        var propertyGetter = ResolvePropertyGetter(subscriber.Component.GetType(), propertyName);
+
+        // Register persistence callback
         _subscriptions[subscriber] = state.RegisterOnPersisting(() =>
             {
                 var storageKey = ComputeKey(subscriber, propertyName);
-                var propertyGetter = ResolvePropertyGetter(subscriber.Component.GetType(), propertyName);
                 var property = propertyGetter.GetValue(subscriber.Component);
                 if (property == null)
                 {
@@ -64,6 +74,9 @@ public void Subscribe(ComponentState subscriber, in CascadingParameterInfo param
                 state.PersistAsJson(storageKey, property, propertyType);
                 return Task.CompletedTask;
             }, subscriber.Renderer.GetComponentRenderMode(subscriber.Component));
+
+        // Register scenario-based restoration callback using PropertyGetter's PropertyInfo
+        RegisterScenarioRestorationCallback(subscriber, parameterInfo, propertyGetter.PropertyInfo);
     }
 
     private static PropertyGetter ResolvePropertyGetter(Type type, string propertyName)
@@ -281,4 +294,35 @@ private static bool IsSerializableKey(object key)
 
         return result;
     }
+
+    [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access otherwise can break functionality when trimming application code", Justification = "Property types of rendered components are preserved through other means and won't get trimmed.")]
+    [UnconditionalSuppressMessage("Trimming", "IL2072:'type' argument does not satisfy 'DynamicallyAccessedMemberTypes.PublicConstructors', 'DynamicallyAccessedMemberTypes.PublicFields', 'DynamicallyAccessedMemberTypes.PublicProperties' in call to target method. The return value of the source method does not have matching annotations.", Justification = "Property types of rendered components are preserved through other means and won't get trimmed.")]
+    private void RegisterScenarioRestorationCallback(ComponentState subscriber, in CascadingParameterInfo parameterInfo, PropertyInfo propertyInfo)
+    {
+        // Check for IPersistentStateFilter attributes
+        var filterAttributes = propertyInfo.GetCustomAttributes(typeof(IPersistentStateFilter), inherit: true);
+
+        // Register restoration callbacks for each filter
+        foreach (IPersistentStateFilter filter in filterAttributes)
+        {
+            RegisterRestorationCallback(subscriber, parameterInfo, filter);
+        }
+    }
+
+    [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access otherwise can break functionality when trimming application code", Justification = "Property types of rendered components are preserved through other means and won't get trimmed.")]
+    [UnconditionalSuppressMessage("Trimming", "IL2072:'type' argument does not satisfy 'DynamicallyAccessedMemberTypes.PublicConstructors', 'DynamicallyAccessedMemberTypes.PublicFields', 'DynamicallyAccessedMemberTypes.PublicProperties' in call to target method. The return value of the source method does not have matching annotations.", Justification = "Property types of rendered components are preserved through other means and won't get trimmed.")]
+    private void RegisterRestorationCallback(ComponentState subscriber, in CascadingParameterInfo parameterInfo, IPersistentStateFilter filter)
+    {
+        var storageKey = ComputeKey(subscriber, parameterInfo.PropertyName);
+        var propertyType = parameterInfo.PropertyType;
+        var valueKey = (subscriber, parameterInfo.PropertyName);
+
+        state.RegisterOnRestoring(filter, () =>
+        {
+            if (state.TryTakeFromJson(storageKey, propertyType, out var value))
+            {
+                _scenarioRestoredValues[valueKey] = value;
+            }
+        });
+    }
 }