Fix TempData and SupplyParameterFromSession persistence for streaming SSR case

src/Components/Endpoints/src/DependencyInjection/TempDataService.cs

@@ -25,17 +25,17 @@ public TempData CreateEmpty(HttpContext httpContext)
         return _tempDataProvider.LoadTempData(httpContext);
     }
 
-    public void Save(HttpContext httpContext, TempData tempData)
+    public void Save(HttpContext httpContext, ITempData tempData)
     {
         if (httpContext.RequestServices.GetService<TempDataCascadingValueSupplier>() is { } supplier)
         {
             supplier.PersistValues(tempData);
         }
 
-        if (!tempData.WasLoaded)
+        if (tempData is not TempData concrete || !concrete.WasLoaded)
         {
             return;
         }
-        _tempDataProvider.SaveTempData(httpContext, tempData.Save());
+        _tempDataProvider.SaveTempData(httpContext, concrete.Save());
     }
 }

src/Components/Endpoints/src/RazorComponentEndpointInvoker.cs

@@ -22,6 +22,7 @@ namespace Microsoft.AspNetCore.Components.Endpoints;
 
 internal partial class RazorComponentEndpointInvoker : IRazorComponentEndpointInvoker
 {
+    private const string SessionEstablishmentKey = "__AspNetCore.Components.Endpoints.SessionEstablishment";
     private readonly EndpointHtmlRenderer _renderer;
     private readonly ILogger<RazorComponentEndpointInvoker> _logger;
     private readonly ComponentsActivityLinkStore _activityLinkStore;
@@ -33,6 +34,34 @@ public RazorComponentEndpointInvoker(EndpointHtmlRenderer renderer, ILogger<Razo
         _logger = logger;
     }
 
+    // Issues the session cookie before the first response chunk is flushed, so that
+    // session-backed values written during or after streaming SSR (TempData,
+    // SupplyParameterFromSession) can still be committed once headers are frozen.
+    private static void TryRegisterSessionEstablishment(HttpContext context)
+    {
+        var session = context.Features.Get<ISessionFeature>()?.Session;
+        if (session is null)
+        {
+            return;
+        }
+
+        context.Response.OnStarting(static state =>
+        {
+            var session = (ISession)state!;
+            if (session.TryGetValue(SessionEstablishmentKey, out _))
+            {
+                return Task.CompletedTask;
+            }
+
+            // Flip the session middleware's establishment gate with a no-op Set+Remove.
+            // Set is the only operation that triggers SessionEstablisher.TryEstablishSession;
+            // Remove leaves no entry behind in the in-memory store and doesn't reset the gate.
+            session.Set(SessionEstablishmentKey, []);
+            session.Remove(SessionEstablishmentKey);
+            return Task.CompletedTask;
+        }, session);
+    }
+
     public Task Render(HttpContext context)
     {
         return _renderer.Dispatcher.InvokeAsync(() => RenderComponentCore(context));
@@ -89,6 +118,12 @@ await _renderer.InitializeStandardComponentServicesAsync(
             handler: result.HandlerName,
             form: result.HandlerName != null && context.Request.HasFormContentType ? await context.Request.ReadFormAsync() : null);
 
+        // Issue the session cookie before any streaming chunks flush the response headers,
+        // so that TempData / [SupplyParameterFromSession] values written during async
+        // rendering can still be persisted after Response.HasStarted. No-op when session
+        // middleware is not registered.
+        TryRegisterSessionEstablishment(context);
+
         // Matches MVC's MemoryPoolHttpResponseStreamWriterFactory.DefaultBufferSize
         var defaultBufferSize = 16 * 1024;
         await using var writer = new HttpResponseStreamWriter(context.Response.Body, Encoding.UTF8, defaultBufferSize, ArrayPool<byte>.Shared, ArrayPool<char>.Shared);
@@ -173,6 +208,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

@@ -10,7 +10,7 @@
 
 namespace Microsoft.AspNetCore.Components.Endpoints;
 
-internal static class TempDataProviderServiceCollectionExtensions
+internal static partial class TempDataProviderServiceCollectionExtensions
 {
     internal static readonly object HttpContextItemKey = new object();
 
@@ -57,12 +57,32 @@ internal static ITempData GetOrCreateTempData(HttpContext httpContext)
         var tempDataService = httpContext.RequestServices.GetRequiredService<TempDataService>();
         var tempDataInstance = tempDataService.CreateEmpty(httpContext);
         httpContext.Items[HttpContextItemKey] = tempDataInstance;
-        httpContext.Response.OnStarting(() =>
+        return tempDataInstance;
+    }
+
+    internal static void PersistTempData(HttpContext httpContext)
+    {
+        if (httpContext.Items.TryGetValue(HttpContextItemKey, out var tempDataObj) && tempDataObj is ITempData tempData)
         {
-            tempDataService.Save(httpContext, tempDataInstance);
-            return Task.CompletedTask;
-        });
+            var provider = httpContext.RequestServices.GetService<ITempDataProvider>();
+            if (provider is CookieTempDataProvider && httpContext.Response.HasStarted)
+            {
+                var logger = httpContext.RequestServices.GetRequiredService<ILoggerFactory>().CreateLogger(typeof(TempDataProviderServiceCollectionExtensions).FullName!);
+                Log.CookieTempDataNotPersistedAfterResponseStarted(logger);
+                return;
+            }
 
-        return tempDataInstance;
+            var tempDataService = httpContext.RequestServices.GetRequiredService<TempDataService>();
+            tempDataService.Save(httpContext, tempData);
+        }
+    }
+
+    private static partial class Log
+    {
+        [LoggerMessage(1, LogLevel.Warning,
+            "TempData values written during or after streaming SSR cannot be persisted by the cookie TempData provider because the response has already started. " +
+            "Switch to the session-storage TempData provider (RazorComponentsServiceOptions.TempDataProviderType = TempDataProviderType.SessionStorage) to enable persistence in streaming SSR scenarios.",
+            EventName = "CookieTempDataNotPersistedAfterResponseStarted")]
+        public static partial void CookieTempDataNotPersistedAfterResponseStarted(ILogger logger);
     }
 }

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

@@ -62,7 +62,7 @@ public async Task PersistAllValues_SetsValueInSession()
     public async Task PersistAllValues_RemovesKey_WhenCallbackReturnsNull()
     {
         var httpContext = CreateHttpContextWithSession();
-        httpContext.Session.SetString("key", "\"existing\"");
+        httpContext.Session.SetString("key", "previous");
 
         _supplier.RegisterValueCallback("key", () => null);
         _supplier.SetRequestContext(httpContext);
@@ -71,6 +71,18 @@ public async Task PersistAllValues_RemovesKey_WhenCallbackReturnsNull()
         Assert.Null(httpContext.Session.GetString("key"));
     }
 
+    [Fact]
+    public async Task PersistAllValues_KeepsKey_WhenCallbackReturnsValue()
+    {
+        var httpContext = CreateHttpContextWithSession();
+
+        _supplier.RegisterValueCallback("key", () => "value");
+        _supplier.SetRequestContext(httpContext);
+        await _supplier.PersistAllValues();
+
+        Assert.Equal("\"value\"", httpContext.Session.GetString("key"));
+    }
+
     [Fact]
     public async Task PersistAllValues_HandlesMultipleKeys()
     {
@@ -160,17 +172,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,7 +34,7 @@ public void Load_ReturnsEmptyTempData_WhenNoSessionDataExists()
     }
 
     [Fact]
-    public void Save_RemovesSessionKey_WhenNoDataToSave()
+    public void Save_RemovesSessionEntry_WhenNoDataToSave()
     {
         var httpContext = CreateHttpContext();
         var session = (TestSession)httpContext.Session;
@@ -44,6 +44,7 @@ public void Save_RemovesSessionKey_WhenNoDataToSave()
         _sessionStateTempDataProvider.SaveTempData(httpContext, tempData.Save());
 
         Assert.DoesNotContain(SessionStorageTempDataProvider.TempDataSessionStateKey, session.Keys);
+        Assert.Empty(_sessionStateTempDataProvider.LoadTempData(httpContext));
     }
 
     [Fact]

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

@@ -0,0 +1,63 @@
+// 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.E2ETest.Tests;
+
+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");
+
+        Browser.Exists(By.Id("streaming-complete"));
+        Navigate($"{ServerPathBase}/supply-parameter-from-session");
+        Browser.Equal("set-during-streaming", () => Browser.FindElement(By.Id("text-email")).Text);
+    }
+
+    [Fact]
+    public void StreamingSSR_PersistsSupplyParameterFromTempData_AfterAsyncRendering()
+    {
+        Navigate($"{ServerPathBase}/streaming-session-persistence");
+
+        Browser.Exists(By.Id("streaming-complete"));
+        Navigate($"{ServerPathBase}/tempdata");
+        Browser.Equal("tempdata-set-during-streaming", () => Browser.FindElement(By.Id("supply-parameter-from-tempdata")).Text);
+    }
+
+    [Fact]
+    public void StreamingSSR_PersistsTempDataCascadingParameter_AfterAsyncRendering()
+    {
+        Navigate($"{ServerPathBase}/streaming-session-persistence");
+
+        Browser.Exists(By.Id("streaming-complete"));
+        Navigate($"{ServerPathBase}/tempdata");
+        Browser.Equal("streaming-tempdata-message", () => Browser.FindElement(By.Id("message")).Text);
+    }
+}

src/Components/test/E2ETest/Tests/TempDataCookieTest.cs

@@ -156,4 +156,14 @@ public void SupplyParameterFromTempDataReadsAndSavesValues()
         Browser.FindElement(By.Id("set-supply-from-tempdata")).Click();
         Browser.Equal("Supplied from TempData", () => Browser.FindElement(By.Id("supply-parameter-from-tempdata")).Text);
     }
+
+    [Fact]
+    public void StreamingSSR_CookieTempData_DoesNotPersistValuesWrittenAfterFirstFlush()
+    {
+        Navigate($"{ServerPathBase}/streaming-session-persistence");
+        Browser.Exists(By.Id("streaming-complete"));
+
+        Navigate($"{ServerPathBase}/tempdata");
+        Browser.Equal("No message", () => Browser.FindElement(By.Id("message")).Text);
+    }
 }

src/Components/test/testassets/Components.TestServer/RazorComponents/Pages/StreamingRendering/StreamingSessionPersistence.razor

@@ -0,0 +1,42 @@
+@page "/streaming-session-persistence"
+@attribute [StreamRendering]
+
+<h1>Streaming Session Persistence Test</h1>
+
+@if (_done)
+{
+    <p id="streaming-complete">Done</p>
+}
+else
+{
+    <p id="streaming-status">Waiting for streaming...</p>
+}
+
+@code {
+    private bool _done;
+
+    [SupplyParameterFromSession]
+    public string? Email { get; set; }
+
+    [SupplyParameterFromTempData]
+    private string SupplyParameterFromTempDataValue { get; set; } = string.Empty;
+
+    [CascadingParameter]
+    public ITempData? TempData { get; set; }
+
+    protected override async Task OnInitializedAsync()
+    {
+        // The await is critical: it causes the component to enter the streaming phase.
+        // Values set AFTER the await are the ones that were silently lost with the old
+        // OnStarting-based persistence, because OnStarting fires before streaming begins.
+        await Task.Delay(100);
+
+        Email = "set-during-streaming";
+        SupplyParameterFromTempDataValue = "tempdata-set-during-streaming";
+        if (TempData is not null)
+        {
+            TempData["Message"] = "streaming-tempdata-message";
+        }
+        _done = true;
+    }
+}