Fix TempData and SupplyParameterFromSession persistence for streaming SSR case

src/Components/Endpoints/src/RazorComponentEndpointInvoker.cs

@@ -173,6 +173,16 @@ 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();
+            sessionSupplier.CleanupNullValues();
+        }
+        TempDataProviderServiceCollectionExtensions.PersistTempData(context);
+        TempDataProviderServiceCollectionExtensions.CleanupTempDataIfEmpty(context);
+
         // Emit comment containing state.
         if (!isErrorHandlerOrReExecuted)
         {

src/Components/Endpoints/src/SessionCascadingValueSupplier.cs

@@ -98,7 +98,11 @@ internal Task PersistAllValues()
                 }
                 else
                 {
-                    session.Remove(sessionKey);
+                    // Write an empty value (read side treats empty strings as null) so that
+                    // session.Set is called even when all current values are null. This ensures
+                    // the session cookie is issued during the OnStarting callback, before
+                    // streaming chunks flush the response headers and prevent further cookies.
+                    session.SetString(sessionKey, string.Empty);
                 }
             }
             catch (Exception ex)
@@ -109,6 +113,35 @@ internal Task PersistAllValues()
         return Task.CompletedTask;
     }
 
+    internal void CleanupNullValues()
+    {
+        if (_valueCallbacks.Count == 0)
+        {
+            return;
+        }
+
+        var session = GetSession();
+        if (session is null)
+        {
+            return;
+        }
+
+        foreach (var (key, valueGetter) in _valueCallbacks)
+        {
+            try
+            {
+                if (valueGetter() is null)
+                {
+                    session.Remove(key.ToLowerInvariant());
+                }
+            }
+            catch (Exception ex)
+            {
+                Log.SessionPersistFail(_logger, ex);
+            }
+        }
+    }
+
     internal void DeleteValueCallback(string sessionKey)
     {
         _valueCallbacks.Remove(sessionKey);

src/Components/Endpoints/src/TempData/SessionStorageTempDataProvider.cs

@@ -31,6 +31,11 @@ public SessionStorageTempDataProvider(
 
             if (session.TryGetValue(TempDataSessionStateKey, out var value))
             {
+                if (value is null || value.Length == 0)
+                {
+                    return new Dictionary<string, object?>();
+                }
+
                 var dataFromSession = JsonSerializer.Deserialize<Dictionary<string, JsonElement>>(value);
                 if (dataFromSession is null)
                 {
@@ -66,7 +71,12 @@ public void SaveTempData(HttpContext context, IDictionary<string, object?> value
         var session = context.Session;
         if (values.Count == 0)
         {
-            session.Remove(TempDataSessionStateKey);
+            // Write an empty payload (rather than Remove) so that the session is
+            // established and the session cookie is issued. This mirrors the
+            // SessionCascadingValueSupplier behavior and is necessary for streaming SSR
+            // where the response headers flush before async work has populated TempData.
+            // The Load path treats empty bytes as no-TempData.
+            session.Set(TempDataSessionStateKey, []);
             return;
         }
 
@@ -75,6 +85,24 @@ public void SaveTempData(HttpContext context, IDictionary<string, object?> value
         Log.TempDataSessionSaveSuccess(_logger);
     }
 
+    public void CleanupIfEmpty(HttpContext context)
+    {
+        ArgumentNullException.ThrowIfNull(context);
+
+        try
+        {
+            var session = context.Session;
+            if (session.TryGetValue(TempDataSessionStateKey, out var bytes) && bytes.Length == 0)
+            {
+                session.Remove(TempDataSessionStateKey);
+            }
+        }
+        catch (Exception ex)
+        {
+            Log.TempDataSessionLoadFailure(_logger, ex);
+        }
+    }
+
     private static partial class Log
     {
         [LoggerMessage(1, LogLevel.Debug, "TempData was not found in session.", EventName = "TempDataSessionNotFound")]

src/Components/Endpoints/src/TempData/TempDataCascadingValueSupplier.cs

@@ -15,6 +15,7 @@ internal partial class TempDataCascadingValueSupplier
 {
     private static readonly ConcurrentDictionary<(Type, string), PropertyGetter> _propertyGetterCache = new();
     private HttpContext? _httpContext;
+    private bool _onStartingRegistered;
     private readonly Dictionary<string, Func<object?>> _valueCallbacks = new(StringComparer.OrdinalIgnoreCase);
     private readonly ILogger<TempDataCascadingValueSupplier> _logger;
 
@@ -33,6 +34,16 @@ internal CascadingParameterSubscription CreateSubscription(
         SupplyParameterFromTempDataAttribute attribute,
         CascadingParameterInfo parameterInfo)
     {
+        if (!_onStartingRegistered && _httpContext is not null)
+        {
+            _onStartingRegistered = true;
+            _httpContext.Response.OnStarting(() =>
+            {
+                TempDataProviderServiceCollectionExtensions.PersistTempData(_httpContext);
+                return Task.CompletedTask;
+            });
+        }
+
         var tempDataKey = attribute.Name ?? parameterInfo.PropertyName;
         var componentType = componentState.Component.GetType();
         var getter = _propertyGetterCache.GetOrAdd((componentType, parameterInfo.PropertyName), PropertyGetterFactory);

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

@@ -59,10 +59,32 @@ internal static ITempData GetOrCreateTempData(HttpContext httpContext)
         httpContext.Items[HttpContextItemKey] = tempDataInstance;
         httpContext.Response.OnStarting(() =>
         {
+            // The provider's SaveTempData establishes the underlying storage (e.g. issues
+            // the session cookie for SessionStorage) so that the post-streaming
+            // PersistTempData call can write the final state after async rendering.
+            _ = tempDataInstance.ContainsKey(string.Empty);
             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);
+        }
+    }
+
+    internal static void CleanupTempDataIfEmpty(HttpContext httpContext)
+    {
+        var provider = httpContext.RequestServices.GetService<ITempDataProvider>();
+        if (provider is SessionStorageTempDataProvider sessionProvider)
+        {
+            sessionProvider.CleanupIfEmpty(httpContext);
+        }
+    }
 }

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

@@ -59,18 +59,42 @@ public async Task PersistAllValues_SetsValueInSession()
     }
 
     [Fact]
-    public async Task PersistAllValues_RemovesKey_WhenCallbackReturnsNull()
+    public async Task PersistAllValues_WritesEmptyPlaceholder_WhenCallbackReturnsNull()
     {
         var httpContext = CreateHttpContextWithSession();
-        httpContext.Session.SetString("key", "\"existing\"");
 
         _supplier.RegisterValueCallback("key", () => null);
         _supplier.SetRequestContext(httpContext);
         await _supplier.PersistAllValues();
+        Assert.Equal(string.Empty, httpContext.Session.GetString("key"));
+    }
+
+    [Fact]
+    public async Task CleanupNullValues_RemovesKey_WhenCallbackReturnsNull()
+    {
+        var httpContext = CreateHttpContextWithSession();
+
+        _supplier.RegisterValueCallback("key", () => null);
+        _supplier.SetRequestContext(httpContext);
+        await _supplier.PersistAllValues();
+        _supplier.CleanupNullValues();
 
         Assert.Null(httpContext.Session.GetString("key"));
     }
 
+    [Fact]
+    public async Task CleanupNullValues_KeepsKey_WhenCallbackReturnsValue()
+    {
+        var httpContext = CreateHttpContextWithSession();
+
+        _supplier.RegisterValueCallback("key", () => "value");
+        _supplier.SetRequestContext(httpContext);
+        await _supplier.PersistAllValues();
+        _supplier.CleanupNullValues();
+
+        Assert.Equal("\"value\"", httpContext.Session.GetString("key"));
+    }
+
     [Fact]
     public async Task PersistAllValues_HandlesMultipleKeys()
     {
@@ -160,17 +184,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);
+        var httpContext = CreateHttpContextWithSession();
         _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()));
     }
 

src/Components/Endpoints/test/TempData/SessionStorageTempDataProviderTest.cs

@@ -34,18 +34,44 @@ public void Load_ReturnsEmptyTempData_WhenNoSessionDataExists()
     }
 
     [Fact]
-    public void Save_RemovesSessionKey_WhenNoDataToSave()
+    public void Save_WritesEmptyPayload_WhenNoDataToSave()
     {
         var httpContext = CreateHttpContext();
-        var session = (TestSession)httpContext.Session;
-        session.Set(SessionStorageTempDataProvider.TempDataSessionStateKey, new byte[] { 1, 2, 3 });
-
         var tempData = CreateTempData();
+
         _sessionStateTempDataProvider.SaveTempData(httpContext, tempData.Save());
 
+        var session = (TestSession)httpContext.Session;
+        Assert.Contains(SessionStorageTempDataProvider.TempDataSessionStateKey, session.Keys);
+        Assert.True(session.TryGetValue(SessionStorageTempDataProvider.TempDataSessionStateKey, out var bytes));
+        Assert.Empty(bytes);
+        Assert.Empty(_sessionStateTempDataProvider.LoadTempData(httpContext));
+    }
+
+    [Fact]
+    public void CleanupIfEmpty_RemovesSessionKey_WhenPayloadIsEmpty()
+    {
+        var httpContext = CreateHttpContext();
+        var session = (TestSession)httpContext.Session;
+        session.Set(SessionStorageTempDataProvider.TempDataSessionStateKey, []);
+
+        _sessionStateTempDataProvider.CleanupIfEmpty(httpContext);
+
         Assert.DoesNotContain(SessionStorageTempDataProvider.TempDataSessionStateKey, session.Keys);
     }
 
+    [Fact]
+    public void CleanupIfEmpty_DoesNotRemoveSessionKey_WhenPayloadHasData()
+    {
+        var httpContext = CreateHttpContext();
+        var session = (TestSession)httpContext.Session;
+        session.Set(SessionStorageTempDataProvider.TempDataSessionStateKey, new byte[] { 1, 2, 3 });
+
+        _sessionStateTempDataProvider.CleanupIfEmpty(httpContext);
+
+        Assert.Contains(SessionStorageTempDataProvider.TempDataSessionStateKey, session.Keys);
+    }
+
     [Fact]
     public void Save_SetsSessionData_WhenDataExists()
     {

src/Components/test/E2ETest/ServerRenderingTests/StreamingSessionPersistenceTest.cs

@@ -0,0 +1,74 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Components.TestServer.RazorComponents;
+using Microsoft.AspNetCore.Components.E2ETest.Infrastructure;
+using Microsoft.AspNetCore.Components.E2ETest.Infrastructure.ServerFixtures;
+using Microsoft.AspNetCore.E2ETesting;
+using OpenQA.Selenium;
+using TestServer;
+using Xunit.Abstractions;
+
+namespace Microsoft.AspNetCore.Components.E2ETests;
+
+public class StreamingSessionPersistenceTest : ServerTestBase<BasicTestAppServerSiteFixture<RazorComponentEndpointsNoInteractivityStartup<App>>>
+{
+    private const string SessionCookieName = ".AspNetCore.Session";
+
+    public StreamingSessionPersistenceTest(
+        BrowserFixture browserFixture,
+        BasicTestAppServerSiteFixture<RazorComponentEndpointsNoInteractivityStartup<App>> serverFixture,
+        ITestOutputHelper output)
+        : base(browserFixture, serverFixture, output)
+    {
+    }
+
+    protected override void InitializeAsyncCore()
+    {
+        _serverFixture.AdditionalArguments.Add("--UseSessionStorageTempDataProvider=true");
+        _serverFixture.AdditionalArguments.Add("--UseSession=true");
+        Browser.Manage().Cookies.DeleteCookieNamed(SessionCookieName);
+        base.InitializeAsyncCore();
+    }
+
+    [Fact]
+    public void StreamingSSR_PersistsSupplyParameterFromSession_AfterAsyncRendering()
+    {
+        Navigate($"{ServerPathBase}/streaming-session-persistence");
+
+        // Wait for streaming to complete — values are set AFTER an await in OnInitializedAsync
+        Browser.Exists(By.Id("streaming-complete"));
+        Browser.Equal("set-during-streaming", () => Browser.FindElement(By.Id("session-value")).Text);
+
+        // Navigate to the read page with a full page load to verify the value survived the request boundary
+        Navigate($"{ServerPathBase}/streaming-session-persistence/read");
+        Browser.Equal("set-during-streaming", () => Browser.FindElement(By.Id("read-session-value")).Text);
+    }
+
+    [Fact]
+    public void StreamingSSR_PersistsSupplyParameterFromTempData_AfterAsyncRendering()
+    {
+        Navigate($"{ServerPathBase}/streaming-session-persistence");
+
+        // Wait for streaming to complete
+        Browser.Exists(By.Id("streaming-complete"));
+        Browser.Equal("tempdata-set-during-streaming", () => Browser.FindElement(By.Id("tempdata-value")).Text);
+
+        // Navigate to the read page to verify TempData was persisted via session storage
+        Navigate($"{ServerPathBase}/streaming-session-persistence/read");
+        Browser.Equal("tempdata-set-during-streaming", () => Browser.FindElement(By.Id("read-tempdata-supply-value")).Text);
+    }
+
+    [Fact]
+    public void StreamingSSR_PersistsTempDataCascadingParameter_AfterAsyncRendering()
+    {
+        Navigate($"{ServerPathBase}/streaming-session-persistence");
+
+        // Wait for streaming to complete
+        Browser.Exists(By.Id("streaming-complete"));
+
+        // Navigate to the read page to verify ITempData values set during streaming were persisted
+        Navigate($"{ServerPathBase}/streaming-session-persistence/read");
+        Browser.Equal("streaming-tempdata-message", () => Browser.FindElement(By.Id("read-tempdata-message")).Text);
+    }
+}