WIP

src/Components/Components/src/PersistentState/PersistentServicesRegistry.cs

@@ -89,6 +89,7 @@ private static void PersistInstanceState(object instance, Type type, PersistentC
             {
                 state.PersistAsJson(key, value, propertyType);
             }
+            Console.WriteLine($"[Persist] type: {instance.GetType()}, propertyType: {propertyType}, key: {key}, result: {value}");
         }
     }
 
@@ -136,6 +137,7 @@ private static void RestoreInstanceState(object instance, Type type, PersistentC
                 var (setter, getter) = accessors.GetAccessor(key);
                 setter.SetValue(instance, result!);
             }
+            Console.WriteLine($"[Restore] type: {instance.GetType()}, propertyType: {propertyType}, key: {key}, result: {result}");
         }
     }
 
@@ -211,8 +213,19 @@ private static string ComputeKey(Type keyType, string propertyName)
             // This happens once per type and property combo, so allocations are ok.
             var assemblyName = keyType.Assembly.FullName;
             var typeName = keyType.FullName;
-            var input = Encoding.UTF8.GetBytes(string.Join(".", assemblyName, typeName, propertyName));
-            return Convert.ToBase64String(SHA256.HashData(input));
+
+            // Internal classes can be bundled in different assemblies during prerendering and WASM rendering.
+            bool isTypeInternal = (!keyType.IsPublic && !keyType.IsNested) || keyType.IsNestedAssembly;
+            var inputString = isTypeInternal
+                ? string.Join(".", typeName, propertyName)
+                : string.Join(".", assemblyName, typeName, propertyName);
+
+            var input = Encoding.UTF8.GetBytes(inputString);
+            var hash = SHA256.HashData(input);
+            var key = Convert.ToBase64String(hash);
+
+            Console.WriteLine($"[ComputeKey] inputString: {inputString}, key: {key}");
+            return key;
         }
 
         internal static IEnumerable<PropertyInfo> GetCandidateBindableProperties(

src/Components/WebAssembly/Server/src/AuthenticationStateSerializer.cs

@@ -2,51 +2,34 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using Microsoft.AspNetCore.Components.Authorization;
-using Microsoft.AspNetCore.Components.Web;
 using Microsoft.Extensions.Options;
 
 namespace Microsoft.AspNetCore.Components.WebAssembly.Server;
 
-internal sealed class AuthenticationStateSerializer : IHostEnvironmentAuthenticationStateProvider, IDisposable
+internal sealed class AuthenticationStateSerializer : IHostEnvironmentAuthenticationStateProvider
 {
-    // Do not change. This must match all versions of the server-side DeserializedAuthenticationStateProvider.PersistenceKey.
-    internal const string PersistenceKey = $"__internal__{nameof(AuthenticationState)}";
-
-    private readonly PersistentComponentState _state;
     private readonly Func<AuthenticationState, ValueTask<AuthenticationStateData?>> _serializeCallback;
-    private readonly PersistingComponentStateSubscription _subscription;
 
-    private Task<AuthenticationState>? _authenticationStateTask;
+    [SupplyParameterFromPersistentComponentState]
+    public AuthenticationStateData? CurrentAuthenticationState { get; set; }
 
-    public AuthenticationStateSerializer(PersistentComponentState persistentComponentState, IOptions<AuthenticationStateSerializationOptions> options)
+    public AuthenticationStateSerializer(IOptions<AuthenticationStateSerializationOptions> options)
     {
-        _state = persistentComponentState;
         _serializeCallback = options.Value.SerializationCallback;
-        _subscription = persistentComponentState.RegisterOnPersisting(OnPersistingAsync, RenderMode.InteractiveWebAssembly);
-    }
-
-    private async Task OnPersistingAsync()
-    {
-        if (_authenticationStateTask is null)
-        {
-            throw new InvalidOperationException($"{nameof(SetAuthenticationState)} must be called before the {nameof(PersistentComponentState)}.{nameof(PersistentComponentState.RegisterOnPersisting)} callback.");
-        }
-
-        var authenticationStateData = await _serializeCallback(await _authenticationStateTask);
-        if (authenticationStateData is not null)
-        {
-            _state.PersistAsJson(PersistenceKey, authenticationStateData);
-        }
     }
 
     /// <inheritdoc />
     public void SetAuthenticationState(Task<AuthenticationState> authenticationStateTask)
     {
-        _authenticationStateTask = authenticationStateTask ?? throw new ArgumentNullException(nameof(authenticationStateTask));
+        ArgumentNullException.ThrowIfNull(authenticationStateTask, nameof(authenticationStateTask));
+
+        // fire and forget, not good... This method can throw, especially on serialization.
+        _ = SetAuthenticationStateAsync(authenticationStateTask);
     }
 
-    public void Dispose()
+    private async Task SetAuthenticationStateAsync(Task<AuthenticationState> authenticationStateTask)
     {
-        _subscription.Dispose();
+        var authenticationState = await authenticationStateTask;
+        CurrentAuthenticationState = await _serializeCallback(authenticationState);
     }
 }

src/Components/WebAssembly/Server/src/WebAssemblyRazorComponentsBuilderExtensions.cs

@@ -4,7 +4,9 @@
 using Microsoft.AspNetCore.Components;
 using Microsoft.AspNetCore.Components.Authorization;
 using Microsoft.AspNetCore.Components.Endpoints.Infrastructure;
+using Microsoft.AspNetCore.Components.Infrastructure;
 using Microsoft.AspNetCore.Components.WebAssembly.Server;
+using Microsoft.AspNetCore.Components.Web;
 using Microsoft.AspNetCore.Components.WebAssembly.Services;
 using Microsoft.Extensions.DependencyInjection.Extensions;
 
@@ -41,6 +43,8 @@ public static IRazorComponentsBuilder AddInteractiveWebAssemblyComponents(this I
     public static IRazorComponentsBuilder AddAuthenticationStateSerialization(this IRazorComponentsBuilder builder, Action<AuthenticationStateSerializationOptions>? configure = null)
     {
         builder.Services.TryAddEnumerable(ServiceDescriptor.Scoped<IHostEnvironmentAuthenticationStateProvider, AuthenticationStateSerializer>());
+        builder.Services.TryAddScoped(sp => (AuthenticationStateSerializer)sp.GetRequiredService<IHostEnvironmentAuthenticationStateProvider>());
+        RegisterPersistentComponentStateServiceCollectionExtensions.AddPersistentServiceRegistration<AuthenticationStateSerializer>(builder.Services, RenderMode.InteractiveAuto);
         if (configure is not null)
         {
             builder.Services.Configure(configure);

src/Components/WebAssembly/WebAssembly.Authentication/src/Services/DeserializedAuthenticationStateProvider.cs

@@ -1,40 +1,37 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System.Diagnostics.CodeAnalysis;
 using System.Security.Claims;
 using Microsoft.AspNetCore.Components.Authorization;
 using Microsoft.Extensions.Options;
-using static Microsoft.AspNetCore.Internal.LinkerFlags;
 
 namespace Microsoft.AspNetCore.Components.WebAssembly.Authentication;
 
 internal sealed class DeserializedAuthenticationStateProvider : AuthenticationStateProvider
 {
-    // Do not change. This must match all versions of the server-side AuthenticationStateSerializer.PersistenceKey.
-    private const string PersistenceKey = $"__internal__{nameof(AuthenticationState)}";
+    // restoring part is on DeserializedAuthenticationStateProvider but persisting part is on AuthenticationStateSerializer
+    // how can we make the key the same if these are two different classes in different assemblies?
+    // should we merge them and move to the Shared folder? Or should we allow passing a custom key to [SupplyParameterFromPersistentComponentState] attribute?
+    private readonly Func<AuthenticationStateData?, Task<AuthenticationState>> _deserializeCallback;
+
+    [SupplyParameterFromPersistentComponentState]
+    public AuthenticationStateData? CurrentAuthenticationState { get; set; }
 
     private static readonly Task<AuthenticationState> _defaultUnauthenticatedTask =
         Task.FromResult(new AuthenticationState(new ClaimsPrincipal(new ClaimsIdentity())));
 
-    private readonly Task<AuthenticationState> _authenticationStateTask = _defaultUnauthenticatedTask;
+    public DeserializedAuthenticationStateProvider(IOptions<AuthenticationStateDeserializationOptions> options)
+    {
+        _deserializeCallback = options.Value.DeserializationCallback;
+    }
 
-    [UnconditionalSuppressMessage(
-        "Trimming",
-        "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access otherwise can break functionality when trimming application code",
-        Justification = $"{nameof(DeserializedAuthenticationStateProvider)} uses the {nameof(DynamicDependencyAttribute)} to preserve the necessary members.")]
-    [DynamicDependency(JsonSerialized, typeof(AuthenticationStateData))]
-    [DynamicDependency(JsonSerialized, typeof(IList<ClaimData>))]
-    [DynamicDependency(JsonSerialized, typeof(ClaimData))]
-    public DeserializedAuthenticationStateProvider(PersistentComponentState state, IOptions<AuthenticationStateDeserializationOptions> options)
+    public override Task<AuthenticationState> GetAuthenticationStateAsync()
     {
-        if (!state.TryTakeFromJson<AuthenticationStateData?>(PersistenceKey, out var authenticationStateData) || authenticationStateData is null)
+        if (CurrentAuthenticationState is null)
         {
-            return;
+            return _defaultUnauthenticatedTask;
         }
-
-        _authenticationStateTask = options.Value.DeserializationCallback(authenticationStateData);
+        var authenticationState = _deserializeCallback(CurrentAuthenticationState);
+        return authenticationState;
     }
-
-    public override Task<AuthenticationState> GetAuthenticationStateAsync() => _authenticationStateTask;
 }

src/Components/WebAssembly/WebAssembly.Authentication/src/WebAssemblyAuthenticationServiceCollectionExtensions.cs

@@ -5,6 +5,8 @@
 using System.Reflection;
 using Microsoft.AspNetCore.Components;
 using Microsoft.AspNetCore.Components.Authorization;
+using Microsoft.AspNetCore.Components.Infrastructure;
+using Microsoft.AspNetCore.Components.Web;
 using Microsoft.AspNetCore.Components.WebAssembly.Authentication;
 using Microsoft.AspNetCore.Components.WebAssembly.Authentication.Internal;
 using Microsoft.Extensions.DependencyInjection.Extensions;
@@ -29,7 +31,10 @@ public static class WebAssemblyAuthenticationServiceCollectionExtensions
     public static IServiceCollection AddAuthenticationStateDeserialization(this IServiceCollection services, Action<AuthenticationStateDeserializationOptions>? configure = null)
     {
         services.AddOptions();
-        services.TryAddScoped<AuthenticationStateProvider, DeserializedAuthenticationStateProvider>();
+        services.TryAddSingleton<AuthenticationStateProvider, DeserializedAuthenticationStateProvider>();
+        services.TryAddSingleton(sp => (DeserializedAuthenticationStateProvider)sp.GetRequiredService<AuthenticationStateProvider>());
+        RegisterPersistentComponentStateServiceCollectionExtensions.AddPersistentServiceRegistration<DeserializedAuthenticationStateProvider>(services, RenderMode.InteractiveWebAssembly);
+
         if (configure != null)
         {
             services.Configure(configure);