[Blazor] Fix Blazor root component state persistence across render modes (dotnet#62370)

* Adds a new virtual method to get the key in ComponentState.
* Ignores the parent component for key computation purposes when crossing a render boundary.
* Uses the ComponentMarkerKey as the key on root components.
* Overrides GetComponentKey on WebAssembly and InteractiveServer to return the ComponentMarkerKey for root components.

Fixes dotnet#62331

Author: Copilot

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

@@ -18,3 +18,4 @@ static Microsoft.AspNetCore.Components.Infrastructure.RegisterPersistentComponen
 static Microsoft.AspNetCore.Components.Infrastructure.ComponentsMetricsServiceCollectionExtensions.AddComponentsMetrics(Microsoft.Extensions.DependencyInjection.IServiceCollection! services) -> Microsoft.Extensions.DependencyInjection.IServiceCollection!
 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?

src/Components/Components/src/Rendering/ComponentState.cs

@@ -339,6 +339,36 @@ internal ValueTask DisposeInBatchAsync(RenderBatchBuilder batchBuilder)
         return DisposeAsync();
     }
 
+    /// <summary>
+    /// Gets the component key for this component instance.
+    /// This is used for state persistence and component identification across render modes.
+    /// </summary>
+    /// <returns>The component key, or null if no key is available.</returns>
+    protected internal virtual object? GetComponentKey()
+    {
+        if (ParentComponentState is not { } parentComponentState)
+        {
+            return null;
+        }
+
+        // Check if the parentComponentState has a `@key` directive applied to the current component.
+        var frames = parentComponentState.CurrentRenderTree.GetFrames();
+        for (var i = 0; i < frames.Count; i++)
+        {
+            ref var currentFrame = ref frames.Array[i];
+            if (currentFrame.FrameType != RenderTreeFrameType.Component ||
+                !ReferenceEquals(Component, currentFrame.Component))
+            {
+                // Skip any frame that is not the current component.
+                continue;
+            }
+
+            return currentFrame.ComponentKey;
+        }
+
+        return null;
+    }
+
     private string GetDebuggerDisplay()
     {
         return $"ComponentId = {ComponentId}, Type = {Component.GetType().Name}, Disposed = {_componentWasDisposed}";

src/Components/Components/src/SupplyParameterFromPersistentComponentStateValueProvider.cs

@@ -45,7 +45,9 @@ public bool CanSupplyValue(in CascadingParameterInfo parameterInfo)
         return state.TryTakeFromJson(storageKey, parameterInfo.PropertyType, out var value) ? value : null;
     }
 
-    [UnconditionalSuppressMessage("Trimming", "IL2075:'this' argument does not satisfy 'DynamicallyAccessedMembersAttribute' in call to target method. The return value of the source method does not have matching annotations.", Justification = "OpenComponent already has the right set of attributes")] [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access otherwise can break functionality when trimming application code", Justification = "OpenComponent already has the right set of attributes")] [UnconditionalSuppressMessage("Trimming", "IL2072:Target parameter argument does not satisfy 'DynamicallyAccessedMembersAttribute' in call to target method. The return value of the source method does not have matching annotations.", Justification = "OpenComponent already has the right set of attributes")]
+    [UnconditionalSuppressMessage("Trimming", "IL2075:'this' argument does not satisfy 'DynamicallyAccessedMembersAttribute' in call to target method. The return value of the source method does not have matching annotations.", Justification = "OpenComponent already has the right set of attributes")]
+    [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access otherwise can break functionality when trimming application code", Justification = "OpenComponent already has the right set of attributes")]
+    [UnconditionalSuppressMessage("Trimming", "IL2072:Target parameter argument does not satisfy 'DynamicallyAccessedMembersAttribute' in call to target method. The return value of the source method does not have matching annotations.", Justification = "OpenComponent already has the right set of attributes")]
     public void Subscribe(ComponentState subscriber, in CascadingParameterInfo parameterInfo)
     {
         var propertyName = parameterInfo.PropertyName;
@@ -221,35 +223,46 @@ private static void GrowBuffer(ref byte[]? pool, ref Span<byte> keyBuffer, int?
 
     private static object? GetSerializableKey(ComponentState componentState)
     {
-        if (componentState.ParentComponentState is not { } parentComponentState)
+        var componentKey = componentState.GetComponentKey();
+        if (componentKey != null && IsSerializableKey(componentKey))
         {
-            return null;
+            return componentKey;
         }
 
-        // Check if the parentComponentState has a `@key` directive applied to the current component.
-        var frames = parentComponentState.CurrentRenderTree.GetFrames();
-        for (var i = 0; i < frames.Count; i++)
+        return null;
+    }
+
+    private static string GetComponentType(ComponentState componentState) => componentState.Component.GetType().FullName!;
+
+    private static string GetParentComponentType(ComponentState componentState)
+    {
+        if (componentState.ParentComponentState == null)
+        {
+            return "";
+        }
+        if (componentState.ParentComponentState.Component == null)
         {
-            ref var currentFrame = ref frames.Array[i];
-            if (currentFrame.FrameType != RenderTree.RenderTreeFrameType.Component ||
-                !ReferenceEquals(componentState.Component, currentFrame.Component))
+            return "";
+        }
+
+        if (componentState.ParentComponentState.ParentComponentState != null)
+        {
+            var renderer = componentState.Renderer;
+            var parentRenderMode = renderer.GetComponentRenderMode(componentState.ParentComponentState.Component);
+            var grandParentRenderMode = renderer.GetComponentRenderMode(componentState.ParentComponentState.ParentComponentState.Component);
+            if (parentRenderMode != grandParentRenderMode)
             {
-                // Skip any frame that is not the current component.
-                continue;
+                // This is the case when EndpointHtmlRenderer introduces an SSRRenderBoundary component.
+                // We want to return "" because the SSRRenderBoundary component is not a real component
+                // and won't appear on the component tree in the WebAssemblyRenderer and RemoteRenderer
+                // interactive scenarios.
+                return "";
             }
-
-            var componentKey = currentFrame.ComponentKey;
-            return !IsSerializableKey(componentKey) ? null : componentKey;
         }
 
-        return null;
+        return GetComponentType(componentState.ParentComponentState);
     }
 
-    private static string GetComponentType(ComponentState componentState) => componentState.Component.GetType().FullName!;
-
-    private static string GetParentComponentType(ComponentState componentState) =>
-        componentState.ParentComponentState == null ? "" : GetComponentType(componentState.ParentComponentState);
-
     private static byte[] KeyFactory((string parentComponentType, string componentType, string propertyName) parts) =>
         SHA256.HashData(Encoding.UTF8.GetBytes(string.Join(".", parts.parentComponentType, parts.componentType, parts.propertyName)));
 

src/Components/Endpoints/src/Rendering/EndpointComponentState.cs

@@ -15,10 +15,12 @@ namespace Microsoft.AspNetCore.Components.Endpoints;
 internal sealed class EndpointComponentState : ComponentState
 {
     private static readonly ConcurrentDictionary<Type, StreamRenderingAttribute?> _streamRenderingAttributeByComponentType = new();
-
+    private readonly EndpointHtmlRenderer _renderer;
     public EndpointComponentState(Renderer renderer, int componentId, IComponent component, ComponentState? parentComponentState)
         : base(renderer, componentId, component, parentComponentState)
     {
+        _renderer = (EndpointHtmlRenderer)renderer;
+
         var streamRenderingAttribute = _streamRenderingAttributeByComponentType.GetOrAdd(component.GetType(),
             type => type.GetCustomAttribute<StreamRenderingAttribute>());
 
@@ -35,6 +37,22 @@ public EndpointComponentState(Renderer renderer, int componentId, IComponent com
 
     public bool StreamRendering { get; }
 
+    protected override object? GetComponentKey()
+    {
+        if (ParentComponentState != null && ParentComponentState.Component is SSRRenderModeBoundary boundary)
+        {
+            var (sequence, key) = _renderer.GetSequenceAndKey(ParentComponentState);
+            var marker = boundary.GetComponentMarkerKey(sequence, key);
+            if (!marker.Equals(default))
+            {
+                return marker.Serialized();
+            }
+        }
+
+        // Fall back to the default implementation
+        return base.GetComponentKey();
+    }
+
     /// <summary>
     /// MetadataUpdateHandler event. This is invoked by the hot reload host via reflection.
     /// </summary>

src/Components/Endpoints/src/Rendering/EndpointHtmlRenderer.Prerendering.cs

@@ -5,6 +5,7 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Text.Encodings.Web;
 using Microsoft.AspNetCore.Components.Rendering;
+using Microsoft.AspNetCore.Components.RenderTree;
 using Microsoft.AspNetCore.Components.Web.HtmlRendering;
 using Microsoft.AspNetCore.Html;
 using Microsoft.AspNetCore.Http;
@@ -298,6 +299,42 @@ internal static ServerComponentInvocationSequence GetOrCreateInvocationId(HttpCo
         return (ServerComponentInvocationSequence)result!;
     }
 
+    internal (int sequence, object? key) GetSequenceAndKey(ComponentState boundaryComponentState)
+    {
+        if (boundaryComponentState is null || boundaryComponentState.Component is not SSRRenderModeBoundary boundary)
+        {
+            throw new InvalidOperationException(
+                "The parent component state must be an SSRRenderModeBoundary to get the sequence and key.");
+        }
+
+        // The boundary is at the root (not supported, but we handle it gracefully)
+        if (boundaryComponentState.ParentComponentState is null)
+        {
+            return (0, null);
+        }
+
+        // Grab the parent of the boundary component. We need to find the SSRRenderModeBoundary component marker frame
+        // within it. As when we do `@rendermode="InteractiveServer" @key="some-key" the sequence we are interested in
+        // is the one on the SSRRenderModeBoundary component marker frame, not the one on the nested component frame.
+        // Same for the key.
+        var targetState = boundaryComponentState.ParentComponentState;
+        var frames = GetCurrentRenderTreeFrames(targetState.ComponentId);
+        for (var i = 0; i < frames.Count; i++)
+        {
+            ref var frame = ref frames.Array[i];
+            if (frame.FrameType == RenderTreeFrameType.Component &&
+                frame.Component is SSRRenderModeBoundary candidate &&
+                ReferenceEquals(candidate, boundary))
+            {
+                // This is the component marker frame, so we can use its sequence and key
+                return (frame.Sequence, frame.ComponentKey);
+            }
+        }
+
+        throw new InvalidOperationException(
+            "The parent component state does not have a valid SSRRenderModeBoundary component marker frame.");
+    }
+
     // An implementation of IHtmlContent that holds a reference to a component until we're ready to emit it as HTML to the response.
     // We don't construct the actual HTML until we receive the call to WriteTo.
     public class PrerenderedComponentHtmlContent : IHtmlAsyncContent

src/Components/Endpoints/src/Rendering/SSRRenderModeBoundary.cs

@@ -220,4 +220,14 @@ private ComponentMarkerKey GenerateMarkerKey(int sequence, object? componentKey)
             FormattedComponentKey = formattedComponentKey,
         };
     }
+
+    /// <summary>
+    /// Gets the ComponentMarkerKey for this boundary if it has been computed.
+    /// This is used for state persistence across render modes.
+    /// </summary>
+    /// <returns>The ComponentMarkerKey if available, null otherwise.</returns>
+    internal ComponentMarkerKey GetComponentMarkerKey(int sequence, object? componentKey)
+    {
+        return _markerKey ??= GenerateMarkerKey(sequence, componentKey);
+    }
 }

src/Components/Server/src/Circuits/RemoteComponentState.cs

@@ -0,0 +1,38 @@
+// 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.Rendering;
+
+namespace Microsoft.AspNetCore.Components.Server.Circuits;
+
+/// <summary>
+/// Specialized ComponentState for Server/Remote rendering that supports ComponentMarkerKey for state persistence.
+/// </summary>
+internal sealed class RemoteComponentState : ComponentState
+{
+    private readonly RemoteRenderer _renderer;
+
+    public RemoteComponentState(
+        RemoteRenderer renderer,
+        int componentId,
+        IComponent component,
+        ComponentState? parentComponentState)
+        : base(renderer, componentId, component, parentComponentState)
+    {
+        _renderer = renderer;
+    }
+
+    protected override object? GetComponentKey()
+    {
+        var markerKey = _renderer.GetMarkerKey(this);
+
+        // If we have a ComponentMarkerKey, return it for state persistence consistency
+        if (markerKey != default)
+        {
+            return markerKey.Serialized();
+        }
+
+        // Fall back to the default implementation
+        return base.GetComponentKey();
+    }
+}

src/Components/Server/src/Circuits/RemoteRenderer.cs

@@ -4,6 +4,7 @@
 using System.Collections.Concurrent;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
+using Microsoft.AspNetCore.Components.Rendering;
 using Microsoft.AspNetCore.Components.RenderTree;
 using Microsoft.AspNetCore.Components.Web;
 using Microsoft.AspNetCore.SignalR;
@@ -313,6 +314,18 @@ protected override IComponent ResolveComponentForRenderMode([DynamicallyAccessed
             _ => throw new NotSupportedException($"Cannot create a component of type '{componentType}' because its render mode '{renderMode}' is not supported by interactive server-side rendering."),
         };
 
+    protected override ComponentState CreateComponentState(int componentId, IComponent component, ComponentState? parentComponentState)
+    {
+        return new RemoteComponentState(this, componentId, component, parentComponentState);
+    }
+
+    internal ComponentMarkerKey GetMarkerKey(RemoteComponentState remoteComponentState)
+    {
+        return remoteComponentState.ParentComponentState != null ?
+            default :
+            _webRootComponentManager!.GetRootComponentKey(remoteComponentState.ComponentId);
+    }
+
     private void ProcessPendingBatch(string? errorMessageOrNull, UnacknowledgedRenderBatch entry)
     {
         var elapsedTime = entry.ValueStopwatch.GetElapsedTime();

src/Components/Shared/src/WebRootComponentManager.cs

@@ -86,12 +86,26 @@ private WebRootComponent GetRequiredWebRootComponent(int ssrComponentId)
 #if COMPONENTS_SERVER
         internal IEnumerable<(int id, ComponentMarkerKey key, (Type componentType, ParameterView parameters))> GetRootComponents()
         {
-            foreach (var (id, (key, type, parameters)) in _webRootComponents)
+            foreach (var (id, (_, key, type, parameters)) in _webRootComponents)
             {
                 yield return (id, key, (type, parameters));
             }
         }
+
 #endif
+        internal ComponentMarkerKey GetRootComponentKey(int componentId)
+        {
+            foreach (var (_, candidate) in _webRootComponents)
+            {
+                var(id, key, _, _) = candidate;
+                if (id == componentId)
+                {
+                    return key;
+                }
+            }
+
+            return default;
+        }
 
         private sealed class WebRootComponent
         {
@@ -135,17 +149,17 @@ private WebRootComponent(
                 _latestParameters = initialParameters;
             }
 
-#if COMPONENTS_SERVER
             public void Deconstruct(
+                out int interactiveComponentId,
                 out ComponentMarkerKey key,
                 out Type componentType,
                 out ParameterView parameters)
             {
+                interactiveComponentId = _interactiveComponentId;
                 key = _key;
                 componentType = _componentType;
                 parameters = _latestParameters.Parameters;
             }
-#endif
 
             public Task UpdateAsync(
                 Renderer renderer,

src/Components/WebAssembly/WebAssembly/src/Rendering/WebAssemblyComponentState.cs

@@ -0,0 +1,38 @@
+// 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.Rendering;
+
+namespace Microsoft.AspNetCore.Components.WebAssembly.Rendering;
+
+/// <summary>
+/// Specialized ComponentState for WebAssembly rendering that supports ComponentMarkerKey for state persistence.
+/// </summary>
+internal sealed class WebAssemblyComponentState : ComponentState
+{
+    private readonly WebAssemblyRenderer _renderer;
+
+    public WebAssemblyComponentState(
+        WebAssemblyRenderer renderer,
+        int componentId,
+        IComponent component,
+        ComponentState? parentComponentState)
+        : base(renderer, componentId, component, parentComponentState)
+    {
+        _renderer = renderer;
+    }
+
+    protected override object? GetComponentKey()
+    {
+        var markerKey = _renderer.GetMarkerKey(this);
+
+        // If we have a ComponentMarkerKey, return it for state persistence consistency
+        if (markerKey != default)
+        {
+            return markerKey.Serialized();
+        }
+
+        // Fall back to the default implementation
+        return base.GetComponentKey();
+    }
+}