Fix TempData and SupplyParameterFromSession persistence for streaming SSR case

src/Components/Endpoints/src/RazorComponentEndpointInvoker.cs

@@ -173,6 +173,14 @@ await _renderer.InitializeStandardComponentServicesAsync(
             _renderer.EmitInitializersIfNecessary(context, bufferWriter);
         }
 
+        // Persist TempData and Session values after all components (including streaming)
+        // have finished rendering, so that values modified during async rendering are captured.
+        if (context.RequestServices.GetService<SessionCascadingValueSupplier>() is { } sessionSupplier)
+        {
+            await sessionSupplier.PersistAllValues();
+        }
+        TempDataProviderServiceCollectionExtensions.PersistTempData(context);
+
         // Emit comment containing state.
         if (!isErrorHandlerOrReExecuted)
         {

src/Components/Endpoints/src/SessionCascadingValueSupplier.cs

@@ -18,7 +18,6 @@ internal partial class SessionCascadingValueSupplier
     private static readonly ConcurrentDictionary<(Type, string), PropertyGetter> _propertyGetterCache = new();
     private static readonly JsonSerializerOptions _jsonOptions = new(JsonSerializerDefaults.Web);
     private HttpContext? _httpContext;
-    private bool _onStartingRegistered;
     private readonly Dictionary<string, Func<object?>> _valueCallbacks = new(StringComparer.OrdinalIgnoreCase);
     private readonly ILogger<SessionCascadingValueSupplier> _logger;
 
@@ -37,12 +36,6 @@ internal CascadingParameterSubscription CreateSubscription(
         SupplyParameterFromSessionAttribute attribute,
         CascadingParameterInfo parameterInfo)
     {
-        if (!_onStartingRegistered && _httpContext is not null)
-        {
-            _onStartingRegistered = true;
-            _httpContext.Response.OnStarting(PersistAllValues);
-        }
-
         var sessionKey = attribute.Name ?? parameterInfo.PropertyName;
         var componentType = componentState.Component.GetType();
         var getter = _propertyGetterCache.GetOrAdd((componentType, parameterInfo.PropertyName), PropertyGetterFactory);

src/Components/Endpoints/src/TempData/TempDataProviderServiceCollectionExtensions.cs

@@ -57,12 +57,16 @@ internal static ITempData GetOrCreateTempData(HttpContext httpContext)
         var tempDataService = httpContext.RequestServices.GetRequiredService<TempDataService>();
         var tempDataInstance = tempDataService.CreateEmpty(httpContext);
         httpContext.Items[HttpContextItemKey] = tempDataInstance;
-        httpContext.Response.OnStarting(() =>
-        {
-            tempDataService.Save(httpContext, tempDataInstance);
-            return Task.CompletedTask;
-        });
 
         return tempDataInstance;
     }
+
+    internal static void PersistTempData(HttpContext httpContext)
+    {
+        if (httpContext.Items.TryGetValue(HttpContextItemKey, out var tempDataObj) && tempDataObj is TempData tempData)
+        {
+            var tempDataService = httpContext.RequestServices.GetRequiredService<TempDataService>();
+            tempDataService.Save(httpContext, tempData);
+        }
+    }
 }

src/Components/Endpoints/test/Session/SessionCascadingValueSupplierTest.cs

@@ -160,17 +160,17 @@ public async Task DeleteCallbacks_RemovesCallbacksForKey()
     }
 
     [Fact]
-    public async Task SetRequestContext_DoesNotRegisterOnStarting_UntilSubscriptionCreated()
+    public async Task SetRequestContext_DoesNotPersist_UntilExplicitlyCalled()
     {
         _supplier.RegisterValueCallback("key", () => "value");
 
         var httpContext = CreateHttpContextWithSession(out var responseFeature);
         _supplier.SetRequestContext(httpContext);
 
-        // OnStarting has not been registered yet because no subscription was created
-        await responseFeature.FireOnStartingAsync();
-
         Assert.Null(httpContext.Session.GetString("key"));
+
+        await _supplier.PersistAllValues();
+        Assert.Equal("\"value\"", httpContext.Session.GetString("key"));
     }
 
     internal static DefaultHttpContext CreateHttpContextWithSession()

src/Components/Endpoints/test/Session/SessionSubscriptionTest.cs

@@ -127,7 +127,7 @@ public void GetCurrentValue_ReturnsComponentValue_OnSubsequentCalls()
     [Fact]
     public async Task CreateSubscription_RegistersValueCallbackAndReturnsSubscription()
     {
-        var httpContext = CreateHttpContextWithSession(out var responseFeature);
+        var httpContext = CreateHttpContextWithSession();
         httpContext.Session.SetString(nameof(TestComponent.Value).ToLowerInvariant(), "\"from-session\"");
         _supplier.SetRequestContext(httpContext);
 
@@ -142,7 +142,7 @@ public async Task CreateSubscription_RegistersValueCallbackAndReturnsSubscriptio
         Assert.Equal("from-session", subscription.GetCurrentValue());
 
         _component.Value = "updated";
-        await responseFeature.FireOnStartingAsync();
+        await _supplier.PersistAllValues();
         Assert.Equal("\"updated\"", httpContext.Session.GetString(nameof(TestComponent.Value).ToLowerInvariant()));
     }