Fix TempData and SupplyParameterFromSession persistence for streaming SSR case

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

@@ -3,6 +3,7 @@
 
 using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
 
 namespace Microsoft.AspNetCore.Components.Endpoints;
 
@@ -25,17 +26,44 @@ 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 data || !data.WasLoaded)
         {
             return;
         }
-        _tempDataProvider.SaveTempData(httpContext, tempData.Save());
+        _tempDataProvider.SaveTempData(httpContext, data.Save());
+    }
+
+    public void Persist(HttpContext httpContext)
+    {
+        if (!httpContext.Items.TryGetValue(TempDataProviderServiceCollectionExtensions.HttpContextItemKey, out var tempDataObj)
+            || tempDataObj is not ITempData tempData)
+        {
+            return;
+        }
+
+        if (_tempDataProvider is CookieTempDataProvider && httpContext.Response.HasStarted)
+        {
+            var logger = httpContext.RequestServices.GetRequiredService<ILoggerFactory>().CreateLogger(typeof(TempDataService).FullName!);
+            Log.CookieTempDataNotPersistedAfterResponseStarted(logger);
+            return;
+        }
+
+        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/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();
+        }
+        context.RequestServices.GetRequiredService<TempDataService>().Persist(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,10 +36,10 @@ internal CascadingParameterSubscription CreateSubscription(
         SupplyParameterFromSessionAttribute attribute,
         CascadingParameterInfo parameterInfo)
     {
-        if (!_onStartingRegistered && _httpContext is not null)
+        if (_httpContext is not null)
         {
-            _onStartingRegistered = true;
-            _httpContext.Response.OnStarting(PersistAllValues);
+            // Ensure that session cookie is issued to allow for persistence from streaming context
+            SessionEstablishmentHelper.TryRegisterSessionEstablishment(_httpContext);
         }
 
         var sessionKey = attribute.Name ?? parameterInfo.PropertyName;

src/Components/Endpoints/src/SessionEstablishmentHelper.cs

@@ -0,0 +1,73 @@
+// 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.Http;
+using Microsoft.AspNetCore.Http.Features;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
+
+namespace Microsoft.AspNetCore.Components.Endpoints;
+
+internal static partial class SessionEstablishmentHelper
+{
+    private const string SessionEstablishmentKey = "__AspNetCore.Components.Endpoints.SessionEstablishment";
+    private const string LoggedResponseHasStartedKey = "__AspNetCore.Components.Endpoints.SessionEstablishment.LoggedResponseHasStarted";
+    private const string LoggedSessionDoesNotExistKey = "__AspNetCore.Components.Endpoints.SessionEstablishment.LoggedSessionDoesNotExist";
+
+    public static void TryRegisterSessionEstablishment(HttpContext context)
+    {
+        var loggerFactory = context.RequestServices.GetRequiredService<ILoggerFactory>();
+        var session = context.Features.Get<ISessionFeature>()?.Session;
+
+        if (session == null)
+        {
+            if (!context.Items.ContainsKey(LoggedSessionDoesNotExistKey))
+            {
+                Log.SessionDoesNotExist(loggerFactory.CreateLogger(typeof(SessionEstablishmentHelper).FullName!));
+                context.Items[LoggedSessionDoesNotExistKey] = true;
+            }
+            return;
+        }
+
+        if (context.Response.HasStarted)
+        {
+            if (!context.Items.ContainsKey(LoggedResponseHasStartedKey))
+            {
+                Log.ResponseHasStarted(loggerFactory.CreateLogger(typeof(SessionEstablishmentHelper).FullName!));
+                context.Items[LoggedResponseHasStartedKey] = true;
+            }
+            return;
+        }
+
+        if (context.Items.ContainsKey(SessionEstablishmentKey))
+        {
+            return;
+        }
+
+        context.Items[SessionEstablishmentKey] = true;
+        context.Response.OnStarting(static state =>
+        {
+            var session = (ISession)state!;
+            session.Set(SessionEstablishmentKey, []);
+            session.Remove(SessionEstablishmentKey);
+            return Task.CompletedTask;
+        }, session);
+    }
+
+    private static partial class Log
+    {
+        [LoggerMessage(1, LogLevel.Warning,
+            "Session state was not persisted to the next request. " +
+            "The response has already started, so the session cookie can no longer be issued. " +
+            "To avoid this, place at least one [SupplyParameterFromSession] (or use Session.Set directly) before any await that triggers the first response flush.",
+            EventName = "SessionStateNotPersistedAfterResponseStarted")]
+        public static partial void ResponseHasStarted(ILogger logger);
+
+        [LoggerMessage(2, LogLevel.Warning,
+            "Session state was not persisted to the next request. " +
+            "No session is available session middleware was not registered. " +
+            "To fix this, add 'builder.Services.AddSession()' and 'app.UseSession()' to your app.",
+            EventName = "SessionDoesNotExist")]
+        public static partial void SessionDoesNotExist(ILogger logger);
+    }
+}

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

@@ -57,11 +57,12 @@ internal static ITempData GetOrCreateTempData(HttpContext httpContext)
         var tempDataService = httpContext.RequestServices.GetRequiredService<TempDataService>();
         var tempDataInstance = tempDataService.CreateEmpty(httpContext);
         httpContext.Items[HttpContextItemKey] = tempDataInstance;
-        httpContext.Response.OnStarting(() =>
+
+        // Ensure that session cookie is issued to allow for persistence from streaming context
+        if (httpContext.RequestServices.GetService<ITempDataProvider>() is SessionStorageTempDataProvider)
         {
-            tempDataService.Save(httpContext, tempDataInstance);
-            return Task.CompletedTask;
-        });
+            SessionEstablishmentHelper.TryRegisterSessionEstablishment(httpContext);
+        }
 
         return tempDataInstance;
     }

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

@@ -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/SessionEstablishmentHelperTest.cs

@@ -0,0 +1,171 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+#nullable enable
+
+using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Http.Features;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Logging.Testing;
+
+namespace Microsoft.AspNetCore.Components.Endpoints;
+
+public class SessionEstablishmentHelperTest
+{
+    [Fact]
+    public void TryRegisterSessionEstablishment_LogsSessionDoesNotExist_WhenNoSessionFeature()
+    {
+        var sink = new TestSink();
+        var httpContext = CreateHttpContext(sink, session: null, responseHasStarted: false);
+
+        SessionEstablishmentHelper.TryRegisterSessionEstablishment(httpContext);
+
+        var write = Assert.Single(sink.Writes);
+        Assert.Equal(LogLevel.Warning, write.LogLevel);
+        Assert.Equal("SessionDoesNotExist", write.EventId.Name);
+    }
+
+    [Fact]
+    public void TryRegisterSessionEstablishment_LogsResponseHasStarted_WhenSessionAvailableButResponseStarted()
+    {
+        var sink = new TestSink();
+        var httpContext = CreateHttpContext(sink, session: new TestSession(), responseHasStarted: true);
+
+        SessionEstablishmentHelper.TryRegisterSessionEstablishment(httpContext);
+
+        var write = Assert.Single(sink.Writes);
+        Assert.Equal(LogLevel.Warning, write.LogLevel);
+        Assert.Equal("SessionStateNotPersistedAfterResponseStarted", write.EventId.Name);
+    }
+
+    [Fact]
+    public void TryRegisterSessionEstablishment_DoesNotLog_WhenSessionAvailableAndResponseNotStarted()
+    {
+        var sink = new TestSink();
+        var httpContext = CreateHttpContext(sink, session: new TestSession(), responseHasStarted: false);
+
+        SessionEstablishmentHelper.TryRegisterSessionEstablishment(httpContext);
+
+        Assert.Empty(sink.Writes);
+    }
+
+    [Fact]
+    public void TryRegisterSessionEstablishment_LogsSessionDoesNotExistOncePerRequest()
+    {
+        var sink = new TestSink();
+        var httpContext = CreateHttpContext(sink, session: null, responseHasStarted: false);
+
+        SessionEstablishmentHelper.TryRegisterSessionEstablishment(httpContext);
+        SessionEstablishmentHelper.TryRegisterSessionEstablishment(httpContext);
+
+        var write = Assert.Single(sink.Writes);
+        Assert.Equal("SessionDoesNotExist", write.EventId.Name);
+    }
+
+    [Fact]
+    public void TryRegisterSessionEstablishment_LogsSessionDoesNotExistOncePerEachRequest()
+    {
+        var sink = new TestSink();
+        var first = CreateHttpContext(sink, session: null, responseHasStarted: false);
+        var second = CreateHttpContext(sink, session: null, responseHasStarted: false);
+
+        SessionEstablishmentHelper.TryRegisterSessionEstablishment(first);
+        SessionEstablishmentHelper.TryRegisterSessionEstablishment(second);
+
+        Assert.Equal(2, sink.Writes.Count);
+        Assert.All(sink.Writes, write => Assert.Equal("SessionDoesNotExist", write.EventId.Name));
+    }
+
+    [Fact]
+    public void TryRegisterSessionEstablishment_LogsResponseHasStartedOncePerRequest()
+    {
+        var sink = new TestSink();
+        var httpContext = CreateHttpContext(sink, session: new TestSession(), responseHasStarted: true);
+
+        SessionEstablishmentHelper.TryRegisterSessionEstablishment(httpContext);
+        SessionEstablishmentHelper.TryRegisterSessionEstablishment(httpContext);
+
+        var write = Assert.Single(sink.Writes);
+        Assert.Equal("SessionStateNotPersistedAfterResponseStarted", write.EventId.Name);
+    }
+
+    [Fact]
+    public void TryRegisterSessionEstablishment_LogsResponseHasStartedOncePerEachRequest()
+    {
+        var sink = new TestSink();
+        var first = CreateHttpContext(sink, session: new TestSession(), responseHasStarted: true);
+        var second = CreateHttpContext(sink, session: new TestSession(), responseHasStarted: true);
+
+        SessionEstablishmentHelper.TryRegisterSessionEstablishment(first);
+        SessionEstablishmentHelper.TryRegisterSessionEstablishment(second);
+
+        Assert.Equal(2, sink.Writes.Count);
+        Assert.All(sink.Writes, write => Assert.Equal("SessionStateNotPersistedAfterResponseStarted", write.EventId.Name));
+    }
+
+    private static DefaultHttpContext CreateHttpContext(TestSink sink, ISession? session, bool responseHasStarted)
+    {
+        var responseFeature = new TestHttpResponseFeature { HasStarted = responseHasStarted };
+        return CreateHttpContext(sink, session, responseFeature);
+    }
+
+    private static DefaultHttpContext CreateHttpContext(TestSink sink, ISession? session, TestHttpResponseFeature responseFeature)
+    {
+        var services = new ServiceCollection();
+        services.AddSingleton<ILoggerFactory>(new TestLoggerFactory(sink, enabled: true));
+
+        var httpContext = new DefaultHttpContext
+        {
+            RequestServices = services.BuildServiceProvider(),
+        };
+
+        if (session is not null)
+        {
+            httpContext.Features.Set<ISessionFeature>(new TestSessionFeature(session));
+        }
+
+        httpContext.Features.Set<IHttpResponseFeature>(responseFeature);
+        return httpContext;
+    }
+
+    private sealed class TestSession : ISession
+    {
+        private readonly Dictionary<string, byte[]> _store = new();
+
+        public bool IsAvailable => true;
+        public string Id => "test-session";
+        public IEnumerable<string> Keys => _store.Keys;
+
+        public void Clear() => _store.Clear();
+        public Task CommitAsync(CancellationToken cancellationToken = default) => Task.CompletedTask;
+        public Task LoadAsync(CancellationToken cancellationToken = default) => Task.CompletedTask;
+        public void Remove(string key) => _store.Remove(key);
+        public void Set(string key, byte[] value) => _store[key] = value;
+        public bool TryGetValue(string key, [System.Diagnostics.CodeAnalysis.NotNullWhen(true)] out byte[]? value) => _store.TryGetValue(key, out value);
+    }
+
+    private sealed class TestSessionFeature : ISessionFeature
+    {
+        public TestSessionFeature(ISession session) => Session = session;
+
+        public ISession Session { get; set; }
+    }
+
+    private sealed class TestHttpResponseFeature : IHttpResponseFeature
+    {
+        public int OnStartingCount { get; private set; }
+
+        public int StatusCode { get; set; } = 200;
+        public string? ReasonPhrase { get; set; }
+        public IHeaderDictionary Headers { get; set; } = new HeaderDictionary();
+        public Stream Body { get; set; } = new MemoryStream();
+        public bool HasStarted { get; set; }
+
+        public void OnCompleted(Func<object, Task> callback, object state)
+        {
+        }
+
+        public void OnStarting(Func<object, Task> callback, object state) => OnStartingCount++;
+    }
+}