Use IStatusCodeReExecuteFeature.

Author: ilonatommy

src/Components/Endpoints/src/Microsoft.AspNetCore.Components.Endpoints.csproj

@@ -40,6 +40,7 @@
     <Reference Include="Microsoft.AspNetCore.Antiforgery" />
     <Reference Include="Microsoft.AspNetCore.Components.Authorization" />
     <Reference Include="Microsoft.AspNetCore.Components.Web" />
+    <Reference Include="Microsoft.AspNetCore.Diagnostics" />
     <Reference Include="Microsoft.AspNetCore.Diagnostics.Abstractions" />
     <Reference Include="Microsoft.AspNetCore.DataProtection.Extensions" />
     <Reference Include="Microsoft.AspNetCore.Hosting.Abstractions" />

src/Components/Endpoints/src/RazorComponentEndpointInvoker.cs

@@ -39,15 +39,15 @@ private async Task RenderComponentCore(HttpContext context)
     {
         context.Response.ContentType = RazorComponentResultExecutor.DefaultContentType;
         var isErrorHandler = context.Features.Get<IExceptionHandlerFeature>() is not null;
-        var hasStatusCodePage = context.Features.Get<IStatusCodePagesFeature>() is not null;
+        var isReExecuted = context.Features.Get<IStatusCodeReExecuteFeature>() is not null;
         if (isErrorHandler)
         {
             Log.InteractivityDisabledForErrorHandling(_logger);
         }
-        _renderer.InitializeStreamingRenderingFraming(context, isErrorHandler, hasStatusCodePage);
-        bool avoidEditingHeaders = hasStatusCodePage && context.Response.StatusCode == StatusCodes.Status404NotFound;
-        if (!avoidEditingHeaders)
+        _renderer.InitializeStreamingRenderingFraming(context, isErrorHandler, isReExecuted);
+        if (!isReExecuted)
         {
+            // re-executed pages have Headers already set up
             EndpointHtmlRenderer.MarkAsAllowingEnhancedNavigation(context);
         }
 
@@ -91,7 +91,7 @@ await _renderer.InitializeStandardComponentServicesAsync(
         using var bufferWriter = new BufferedTextWriter(writer);
 
         int originalStatusCode = context.Response.StatusCode;
-        bool isErrorHandlerOrHasStatusCodePage = isErrorHandler || hasStatusCodePage;
+        bool isErrorHandlerOrReExecuted = isErrorHandler || isReExecuted;
 
         // Note that we always use Static rendering mode for the top-level output from a RazorComponentResult,
         // because you never want to serialize the invocation of RazorComponentResultHost. Instead, that host
@@ -100,14 +100,15 @@ await _renderer.InitializeStandardComponentServicesAsync(
             context,
             rootComponent,
             ParameterView.Empty,
-            waitForQuiescence: result.IsPost || isErrorHandlerOrHasStatusCodePage);
+            waitForQuiescence: result.IsPost || isErrorHandlerOrReExecuted);
 
-        bool requiresReexecution = originalStatusCode != context.Response.StatusCode && hasStatusCodePage;
-        if (requiresReexecution)
+        bool isReExecutionRequested = context.Features.Get<IStatusCodeReExecuteFeature>() is not null;
+        bool avoidStartingResponse = !isReExecuted && isReExecutionRequested;
+        if (avoidStartingResponse)
         {
-            // If the response is a 404, we don't want to write any content.
-            // This is because the 404 status code is used by the routing middleware
-            // to indicate that no endpoint was found for the request.
+            // re-execution feature was set during rendering,
+            // we should finish early to avoid writing to the response
+            // and let the re-execution middleware take care of it
             return;
         }
 
@@ -162,7 +163,7 @@ await _renderer.InitializeStandardComponentServicesAsync(
         }
 
         // Emit comment containing state.
-        if (!isErrorHandlerOrHasStatusCodePage)
+        if (!isErrorHandlerOrReExecuted)
         {
             var componentStateHtmlContent = await _renderer.PrerenderPersistedStateAsync(context);
             componentStateHtmlContent.WriteTo(bufferWriter, HtmlEncoder.Default);

src/Components/Endpoints/src/Rendering/EndpointHtmlRenderer.EventDispatch.cs

@@ -4,6 +4,7 @@
 using Microsoft.AspNetCore.Components.Endpoints.Rendering;
 using Microsoft.AspNetCore.Components.Rendering;
 using Microsoft.AspNetCore.Components.RenderTree;
+using Microsoft.AspNetCore.Diagnostics;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.WebUtilities;
 using Microsoft.Extensions.DependencyInjection;
@@ -93,6 +94,8 @@ private async Task SetNotFoundResponseAsync(string baseUri)
             _httpContext.Response.StatusCode = StatusCodes.Status404NotFound;
             _httpContext.Response.ContentType = null;
         }
+        var statusCodeFeature = new StatusCodeReExecuteFeature();
+        _httpContext.Features.Set<IStatusCodeReExecuteFeature>(statusCodeFeature);
         SignalRendererToFinishRendering();
     }
 

src/Components/Endpoints/src/Rendering/EndpointHtmlRenderer.Prerendering.cs

@@ -19,7 +19,7 @@ internal partial class EndpointHtmlRenderer
 
     protected override IComponent ResolveComponentForRenderMode([DynamicallyAccessedMembers(Component)] Type componentType, int? parentComponentId, IComponentActivator componentActivator, IComponentRenderMode renderMode)
     {
-        if (_isHandlingErrors || _hasStatusCodePage)
+        if (_isHandlingErrors || _isReExecuted)
         {
             // Ignore the render mode boundary in error scenarios.
             return componentActivator.CreateInstance(componentType);
@@ -167,7 +167,7 @@ private async Task WaitForResultReady(bool waitForQuiescence, PrerenderedCompone
         }
         else if (_nonStreamingPendingTasks.Count > 0)
         {
-            if (_hasStatusCodePage)
+            if (_isReExecuted)
             {
                 HandleNonStreamingTasks();
             }

src/Components/Endpoints/src/Rendering/EndpointHtmlRenderer.Streaming.cs

@@ -23,14 +23,13 @@ internal partial class EndpointHtmlRenderer
     private HashSet<int>? _visitedComponentIdsInCurrentStreamingBatch;
     private string? _ssrFramingCommentMarkup;
     private bool _isHandlingErrors;
-    private bool _hasStatusCodePage;
+    private bool _isReExecuted;
 
-    public void InitializeStreamingRenderingFraming(HttpContext httpContext, bool isErrorHandler, bool hasStatusCodePage)
+    public void InitializeStreamingRenderingFraming(HttpContext httpContext, bool isErrorHandler, bool isReExecuted)
     {
         _isHandlingErrors = isErrorHandler;
-        _hasStatusCodePage = hasStatusCodePage;
-        bool avoidEditingHeaders = hasStatusCodePage && httpContext.Response.StatusCode == StatusCodes.Status404NotFound;
-        if (!avoidEditingHeaders && IsProgressivelyEnhancedNavigation(httpContext.Request))
+        _isReExecuted = isReExecuted;
+        if (!isReExecuted && IsProgressivelyEnhancedNavigation(httpContext.Request))
         {
             var id = Guid.NewGuid().ToString();
             httpContext.Response.Headers.Add(_streamingRenderingFramingHeaderName, id);

src/Components/Endpoints/src/Rendering/EndpointHtmlRenderer.cs

@@ -165,7 +165,7 @@ protected override ComponentState CreateComponentState(int componentId, ICompone
 
     protected override void AddPendingTask(ComponentState? componentState, Task task)
     {
-        if (_hasStatusCodePage)
+        if (_isReExecuted)
         {
             return;
         }

src/Components/Endpoints/src/Results/RazorComponentResultExecutor.cs

@@ -48,8 +48,8 @@ private static Task RenderComponentToResponse(
         return endpointHtmlRenderer.Dispatcher.InvokeAsync(async () =>
         {
             var isErrorHandler = httpContext.Features.Get<IExceptionHandlerFeature>() is not null;
-            var hasStatusCodePage = httpContext.Features.Get<IStatusCodePagesFeature>() is not null;
-            endpointHtmlRenderer.InitializeStreamingRenderingFraming(httpContext, isErrorHandler, hasStatusCodePage);
+            var isReExecuted = httpContext.Features.Get<IStatusCodeReExecuteFeature>() is not null;
+            endpointHtmlRenderer.InitializeStreamingRenderingFraming(httpContext, isErrorHandler, isReExecuted);
             EndpointHtmlRenderer.MarkAsAllowingEnhancedNavigation(httpContext);
 
             // We could pool these dictionary instances if we wanted, and possibly even the ParameterView

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

@@ -6,13 +6,14 @@
 using System.Linq;
 using System.Text;
 using System.Threading.Tasks;
-using Microsoft.AspNetCore.Components.E2ETest.Infrastructure.ServerFixtures;
-using Microsoft.AspNetCore.Components.E2ETest.Infrastructure;
-using TestServer;
 using Components.TestServer.RazorComponents;
+using Components.TestServer.RazorComponents.Pages.StreamingRendering;
+using Microsoft.AspNetCore.Components.E2ETest.Infrastructure;
+using Microsoft.AspNetCore.Components.E2ETest.Infrastructure.ServerFixtures;
 using Microsoft.AspNetCore.E2ETesting;
-using Xunit.Abstractions;
 using OpenQA.Selenium;
+using TestServer;
+using Xunit.Abstractions;
 
 namespace Microsoft.AspNetCore.Components.E2ETests.ServerRenderingTests;
 
@@ -21,16 +22,19 @@ public class StatusCodePagesTest(BrowserFixture browserFixture, BasicTestAppServ
 {
 
     [Theory]
-    [InlineData(true)]
-    [InlineData(false)]
-    public void StatusCodePagesWithReexecution(bool setNotFound)
+    [InlineData(false, true)]
+    [InlineData(true, false)]
+    [InlineData(false, false)]
+    public void StatusCodePagesWithReExecution(bool setNotFound, bool streaming)
     {
-        Navigate($"{ServerPathBase}/reexecution/set-not-found?shouldSet={setNotFound}");
+        string streamingPath = streaming ? "streaming-" : "";
+        Navigate($"{ServerPathBase}/reexecution/{streamingPath}set-not-found?shouldSet={setNotFound}");
 
         string expectedTitle = setNotFound ? "Re-executed page" : "Original page";
         Browser.Equal(expectedTitle, () => Browser.Title);
         var infoText = Browser.FindElement(By.Id("test-info")).Text;
-        string expectedInfoText = setNotFound ? "Welcome On Page Re-executed After Not Found Event" : "Any content";
+        // streaming when response started does not re-execute
+        string expectedInfoText = streaming ? "Default Not Found Page" : setNotFound ? "Welcome On Page Re-executed After Not Found Event" : "Any content";
         Assert.Contains(expectedInfoText, infoText);
     }