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<(IPersistentComponentStateScenario Scenario, Action Callback, bool IsRecurring)> _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,82 @@ internal bool TryTakeFromJson(string key, [DynamicallyAccessedMembers(JsonSerial
         }
     }
 
+    /// <summary>
+    /// Registers a callback to be invoked when state is restored for a specific scenario.
+    /// If state for the scenario is already available, the callback is invoked immediately.
+    /// </summary>
+    /// <param name="scenario">The scenario for which to register the callback.</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(
+        IPersistentComponentStateScenario scenario, 
+        Action callback)
+    {
+        ArgumentNullException.ThrowIfNull(scenario);
+        ArgumentNullException.ThrowIfNull(callback);
+
+        var isRecurring = scenario.IsRecurring;
+        _restoringCallbacks.Add((scenario, callback, isRecurring));
+
+        // If we already have a current scenario and it matches, invoke immediately
+        if (CurrentScenario != null && ShouldInvokeCallback(scenario, CurrentScenario))
+        {
+            callback();
+        }
+
+        return new RestoringComponentStateSubscription(_restoringCallbacks, scenario, callback, isRecurring);
+    }
+
+    /// <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 (callbackScenario, callback, isRecurring) = _restoringCallbacks[i];
+
+            if (ShouldInvokeCallback(callbackScenario, scenario))
+            {
+                callback();
+
+                // Remove non-recurring callbacks after invocation
+                if (!isRecurring)
+                {
+                    _restoringCallbacks.RemoveAt(i);
+                }
+            }
+        }
+    }
+
+    private static bool ShouldInvokeCallback(IPersistentComponentStateScenario callbackScenario, IPersistentComponentStateScenario currentScenario)
+    {
+        // For now, match scenarios by type and properties
+        // This can be enhanced with more sophisticated matching logic
+        return callbackScenario.GetType() == currentScenario.GetType() &&
+               callbackScenario.Equals(currentScenario);
+    }
+
     private bool TryTake(string key, out byte[]? value)
     {
         ArgumentNullException.ThrowIfNull(key);

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

@@ -55,9 +55,33 @@ 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 (scenario == null)
+        {
+            // First-time initialization
+            State.InitializeExistingState(data);
+        }
+        else
+        {
+            // Scenario-based update
+            State.UpdateExistingState(data, scenario);
+        }
+
         _servicesRegistry?.Restore(State);
     }
 

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

@@ -21,3 +21,13 @@ 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?
+Microsoft.AspNetCore.Components.PersistentComponentState.RegisterOnRestoring(Microsoft.AspNetCore.Components.IPersistentComponentStateScenario! scenario, 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/RestoringComponentStateSubscription.cs

@@ -0,0 +1,44 @@
+// 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<(IPersistentComponentStateScenario Scenario, Action Callback, bool IsRecurring)>? _callbacks;
+    private readonly IPersistentComponentStateScenario? _scenario;
+    private readonly Action? _callback;
+    private readonly bool _isRecurring;
+
+    internal RestoringComponentStateSubscription(
+        List<(IPersistentComponentStateScenario Scenario, Action Callback, bool IsRecurring)> callbacks,
+        IPersistentComponentStateScenario scenario,
+        Action callback,
+        bool isRecurring)
+    {
+        _callbacks = callbacks;
+        _scenario = scenario;
+        _callback = callback;
+        _isRecurring = isRecurring;
+    }
+
+    /// <inheritdoc />
+    public void Dispose()
+    {
+        if (_callbacks != null && _scenario != null && _callback != null)
+        {
+            for (int i = _callbacks.Count - 1; i >= 0; i--)
+            {
+                var (scenario, callback, isRecurring) = _callbacks[i];
+                if (ReferenceEquals(scenario, _scenario) && ReferenceEquals(callback, _callback) && isRecurring == _isRecurring)
+                {
+                    _callbacks.RemoveAt(i);
+                    break;
+                }
+            }
+        }
+    }
+}

src/Components/Components/src/SupplyParameterFromPersistentComponentStateValueProvider.cs

@@ -42,6 +42,12 @@ public bool CanSupplyValue(in CascadingParameterInfo parameterInfo)
         var componentState = (ComponentState)key!;
         var storageKey = ComputeKey(componentState, parameterInfo.PropertyName);
 
+        // Check if there are scenario filters on the property
+        if (ShouldFilterByScenario(componentState, parameterInfo))
+        {
+            return null; // Don't provide value if scenario filtering rejects it
+        }
+
         return state.TryTakeFromJson(storageKey, parameterInfo.PropertyType, out var value) ? value : null;
     }
 
@@ -281,4 +287,44 @@ private static bool IsSerializableKey(object key)
 
         return result;
     }
+
+    [UnconditionalSuppressMessage("Trimming", "IL2072:'type' argument does not satisfy 'DynamicallyAccessedMemberTypes.PublicProperties' in call to 'Microsoft.AspNetCore.Components.SupplyParameterFromPersistentComponentStateValueProvider.GetPropertyForScenarioFiltering(Type, String)'. The return value of method 'System.Object.GetType()' does not have matching annotations.", Justification = "Properties of rendered components are preserved through other means and won't get trimmed.")]
+    private bool ShouldFilterByScenario(ComponentState componentState, in CascadingParameterInfo parameterInfo)
+    {
+        // If there's no current scenario, don't filter
+        if (state.CurrentScenario == null)
+        {
+            return false;
+        }
+
+        // Get the property info to check for filter attributes
+        var componentType = componentState.Component.GetType();
+        var propertyInfo = GetPropertyForScenarioFiltering(componentType, parameterInfo.PropertyName);
+        if (propertyInfo == null)
+        {
+            return false;
+        }
+
+        // Check for IPersistentStateFilter attributes
+        var filterAttributes = propertyInfo.GetCustomAttributes(typeof(IPersistentStateFilter), inherit: true);
+        if (filterAttributes.Length == 0)
+        {
+            return false; // No filters, allow state
+        }
+
+        // Check if any filter allows the current scenario
+        foreach (IPersistentStateFilter filter in filterAttributes)
+        {
+            if (filter.ShouldRestore(state.CurrentScenario))
+            {
+                return false; // At least one filter allows it
+            }
+        }
+
+        return true; // No filter allows it, so filter it out
+    }
+
+    [UnconditionalSuppressMessage("Trimming", "IL2075:'this' argument does not satisfy 'DynamicallyAccessedMemberTypes.PublicProperties' in call to 'System.Type.GetProperty(String)'. The return value of method 'System.Object.GetType()' does not have matching annotations.", Justification = "Properties of rendered components are preserved through other means and won't get trimmed.")]
+    private static PropertyInfo? GetPropertyForScenarioFiltering([DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicProperties)] Type type, string propertyName)
+        => type.GetProperty(propertyName);
 }

src/Components/Components/test/Microsoft.AspNetCore.Components.Tests.csproj

@@ -7,6 +7,7 @@
 
   <ItemGroup>
     <Reference Include="Microsoft.AspNetCore.Components" />
+    <Reference Include="Microsoft.AspNetCore.Components.Web" />
     <Reference Include="Microsoft.Extensions.DependencyInjection" />
     <Reference Include="Microsoft.Extensions.Diagnostics.Testing" />
   </ItemGroup>

src/Components/Components/test/ScenarioBasedPersistentComponentStateTest.cs

@@ -0,0 +1,96 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Microsoft.AspNetCore.Components.Web;
+using Xunit;
+
+namespace Microsoft.AspNetCore.Components;
+
+public class ScenarioBasedPersistentComponentStateTest
+{
+    [Fact]
+    public void WebPersistenceContext_Properties_SetCorrectly()
+    {
+        // Arrange & Act
+        var enhancedNavContext = new WebPersistenceContext(WebPersistenceReason.EnhancedNavigation);
+        var prerenderingContext = new WebPersistenceContext(WebPersistenceReason.Prerendering);
+        var reconnectionContext = new WebPersistenceContext(WebPersistenceReason.Reconnection);
+
+        // Assert
+        Assert.Equal(WebPersistenceReason.EnhancedNavigation, enhancedNavContext.Reason);
+        Assert.True(enhancedNavContext.IsRecurring);
+
+        Assert.Equal(WebPersistenceReason.Prerendering, prerenderingContext.Reason);
+        Assert.False(prerenderingContext.IsRecurring);
+
+        Assert.Equal(WebPersistenceReason.Reconnection, reconnectionContext.Reason);
+        Assert.False(reconnectionContext.IsRecurring);
+    }
+
+    [Fact]
+    public void WebPersistenceContext_StaticProperties_ReturnCorrectInstances()
+    {
+        // Act
+        var enhancedNav = WebPersistenceContext.EnhancedNavigation;
+        var prerendering = WebPersistenceContext.Prerendering;
+        var reconnection = WebPersistenceContext.Reconnection;
+
+        // Assert
+        Assert.Equal(WebPersistenceReason.EnhancedNavigation, enhancedNav.Reason);
+        Assert.Equal(WebPersistenceReason.Prerendering, prerendering.Reason);
+        Assert.Equal(WebPersistenceReason.Reconnection, reconnection.Reason);
+    }
+
+    [Fact]
+    public void WebPersistenceContext_Equals_WorksCorrectly()
+    {
+        // Arrange
+        var context1 = new WebPersistenceContext(WebPersistenceReason.EnhancedNavigation);
+        var context2 = new WebPersistenceContext(WebPersistenceReason.EnhancedNavigation);
+        var context3 = new WebPersistenceContext(WebPersistenceReason.Prerendering);
+
+        // Act & Assert
+        Assert.True(context1.Equals(context2));
+        Assert.False(context1.Equals(context3));
+        Assert.False(context1.Equals(null));
+    }
+
+    [Fact]
+    public void WebPersistenceContext_GetHashCode_WorksCorrectly()
+    {
+        // Arrange
+        var context1 = new WebPersistenceContext(WebPersistenceReason.EnhancedNavigation);
+        var context2 = new WebPersistenceContext(WebPersistenceReason.EnhancedNavigation);
+        var context3 = new WebPersistenceContext(WebPersistenceReason.Prerendering);
+
+        // Act & Assert
+        Assert.Equal(context1.GetHashCode(), context2.GetHashCode());
+        Assert.NotEqual(context1.GetHashCode(), context3.GetHashCode());
+    }
+
+    [Fact]
+    public void FilterAttributes_ShouldRestore_WorksCorrectly()
+    {
+        // Arrange
+        var enhancedNavContext = new WebPersistenceContext(WebPersistenceReason.EnhancedNavigation, new InteractiveServerRenderMode());
+        var prerenderingContext = new WebPersistenceContext(WebPersistenceReason.Prerendering, new InteractiveServerRenderMode());
+        var reconnectionContext = new WebPersistenceContext(WebPersistenceReason.Reconnection, new InteractiveServerRenderMode());
+
+        var enhancedNavFilter = new UpdateStateOnEnhancedNavigationAttribute();
+        var prerenderingFilter = new RestoreStateOnPrerenderingAttribute();
+        var reconnectionFilter = new RestoreStateOnReconnectionAttribute();
+
+        // Act & Assert
+        Assert.True(enhancedNavFilter.ShouldRestore(enhancedNavContext));
+        Assert.False(enhancedNavFilter.ShouldRestore(prerenderingContext));
+        Assert.False(enhancedNavFilter.ShouldRestore(reconnectionContext));
+
+        Assert.False(prerenderingFilter.ShouldRestore(enhancedNavContext));
+        Assert.True(prerenderingFilter.ShouldRestore(prerenderingContext));
+        Assert.False(prerenderingFilter.ShouldRestore(reconnectionContext));
+
+        Assert.False(reconnectionFilter.ShouldRestore(enhancedNavContext));
+        Assert.False(reconnectionFilter.ShouldRestore(prerenderingContext));
+        Assert.True(reconnectionFilter.ShouldRestore(reconnectionContext));
+    }
+}