Change priority of re-execution handling and allow router to stream NotFound contents

src/Components/Components/src/Routing/Router.cs

@@ -19,6 +19,7 @@ namespace Microsoft.AspNetCore.Components.Routing;
 /// <summary>
 /// A component that supplies route data corresponding to the current navigation state.
 /// </summary>
+[StreamRendering]
 public partial class Router : IComponent, IHandleAfterRender, IDisposable
 {
     // Dictionary is intentionally used instead of ReadOnlyDictionary to reduce Blazor size

src/Components/Endpoints/src/RazorComponentEndpointInvoker.cs

@@ -102,13 +102,6 @@ await _renderer.InitializeStandardComponentServicesAsync(
             ParameterView.Empty,
             waitForQuiescence: result.IsPost || isErrorHandlerOrReExecuted);
 
-        bool avoidStartingResponse = hasStatusCodePage && !isReExecuted && context.Response.StatusCode == StatusCodes.Status404NotFound;
-        if (avoidStartingResponse)
-        {
-            // the request is going to be re-executed, we should avoid writing to the response
-            return;
-        }
-
         Task quiesceTask;
         if (!result.IsPost)
         {

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

@@ -77,27 +77,21 @@ private Task ReturnErrorResponse(string detailedMessage)
             : Task.CompletedTask;
     }
 
-    private async Task SetNotFoundResponseAsync(string baseUri)
+    private void SetNotFoundResponse(object? sender, EventArgs args)
     {
         if (_httpContext.Response.HasStarted)
         {
-            var defaultBufferSize = 16 * 1024;
-            await using var writer = new HttpResponseStreamWriter(_httpContext.Response.Body, Encoding.UTF8, defaultBufferSize, ArrayPool<byte>.Shared, ArrayPool<char>.Shared);
-            using var bufferWriter = new BufferedTextWriter(writer);
-            var notFoundUri = $"{baseUri}not-found";
-            HandleNavigationAfterResponseStarted(bufferWriter, _httpContext, notFoundUri);
-            await bufferWriter.FlushAsync();
+            // We're expecting the Router to continue streaming the NotFound contents
         }
         else
         {
             _httpContext.Response.StatusCode = StatusCodes.Status404NotFound;
-            _httpContext.Response.ContentType = null;
-        }
 
-        // When the application triggers a NotFound event, we continue rendering the current batch.
-        // However, after completing this batch, we do not want to process any further UI updates,
-        // as we are going to return a 404 status and discard the UI updates generated so far.
-        SignalRendererToFinishRendering();
+            // When the application triggers a NotFound event, we continue rendering the current batch.
+            // However, after completing this batch, we do not want to process any further UI updates,
+            // as we are going to return a 404 status and discard the UI updates generated so far.
+            SignalRendererToFinishRendering();
+        }
     }
 
     private async Task OnNavigateTo(string uri)

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

@@ -85,7 +85,7 @@ internal async Task InitializeStandardComponentServicesAsync(
 
         if (navigationManager != null)
         {
-            navigationManager.OnNotFound += async (sender, args) => await SetNotFoundResponseAsync(navigationManager.BaseUri);
+            navigationManager.OnNotFound += SetNotFoundResponse;
         }
 
         var authenticationStateProvider = httpContext.RequestServices.GetService<AuthenticationStateProvider>();

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

@@ -80,13 +80,6 @@ public void NavigatesWithoutInteractivityByRequestRedirection(bool controlFlowBy
         Browser.Equal("Routing test cases", () => Browser.Exists(By.Id("test-info")).Text);
     }
 
-    [Fact]
-    public void CanRenderNotFoundPageAfterStreamingStarted()
-    {
-        Navigate($"{ServerPathBase}/streaming-set-not-found");
-        Browser.Equal("Default Not Found Page", () => Browser.Title);
-    }
-
     [Theory]
     [InlineData(true)]
     [InlineData(false)]
@@ -127,36 +120,35 @@ private void Assert404ReExecuted() =>
     [Theory]
     [InlineData(true)]
     [InlineData(false)]
-    public void CanRenderNotFoundPageNoStreaming(bool useCustomNotFoundPage)
+    public void CanRenderNotFoundPage(bool streamingStarted)
     {
-        string query = useCustomNotFoundPage ? "&useCustomNotFoundPage=true" : "";
-        Navigate($"{ServerPathBase}/set-not-found?shouldSet=true{query}");
-
-        if (useCustomNotFoundPage)
-        {
-            var infoText = Browser.FindElement(By.Id("test-info")).Text;
-            Assert.Contains("Welcome On Custom Not Found Page", infoText);
-            // custom page should have a custom layout
-            var aboutLink = Browser.FindElement(By.Id("about-link")).Text;
-            Assert.Contains("About", aboutLink);
-        }
-        else
-        {
-            var bodyText = Browser.FindElement(By.TagName("body")).Text;
-            Assert.Contains("There's nothing here", bodyText);
-        }
+        string streamingPath = streamingStarted ? "-streaming" : "";
+        Navigate($"{ServerPathBase}/set-not-found-ssr{streamingPath}?useCustomNotFoundPage=true");
+
+        var infoText = Browser.FindElement(By.Id("test-info")).Text;
+        Assert.Contains("Welcome On Custom Not Found Page", infoText);
+        // custom page should have a custom layout
+        var aboutLink = Browser.FindElement(By.Id("about-link")).Text;
+        Assert.Contains("About", aboutLink);
     }
 
     [Theory]
-    [InlineData(true)]
     [InlineData(false)]
-    public void CanRenderNotFoundPageWithStreaming(bool useCustomNotFoundPage)
+    [InlineData(true)]
+    public void DoesNotReExecuteIf404WasHandled(bool streamingStarted)
     {
-        // when streaming started, we always render page under "not-found" path
-        string query = useCustomNotFoundPage ? "?useCustomNotFoundPage=true" : "";
-        Navigate($"{ServerPathBase}/streaming-set-not-found{query}");
+        string streamingPath = streamingStarted ? "-streaming" : "";
+        Navigate($"{ServerPathBase}/reexecution/set-not-found-ssr{streamingPath}");
+        AssertNotFoundFragmentRendered();
+    }
+
+    private void AssertNotFoundFragmentRendered() =>
+        Browser.Equal("There's nothing here", () => Browser.FindElement(By.Id("not-found-fragment")).Text);
 
-        string expectedTitle = "Default Not Found Page";
-        Browser.Equal(expectedTitle, () => Browser.Title);
+    [Fact]
+    public void StatusCodePagesWithReExecution()
+    {
+        Navigate($"{ServerPathBase}/reexecution/trigger-404");
+        Browser.Equal("Re-executed page", () => Browser.Title);
     }
 }

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

@@ -37,11 +37,7 @@ public void CanRenderNotFoundInteractive(string renderingMode, bool useCustomNot
 
         if (useCustomNotFoundPage)
         {
-            var infoText = Browser.FindElement(By.Id("test-info")).Text;
-            Assert.Contains("Welcome On Custom Not Found Page", infoText);
-            // custom page should have a custom layout
-            var aboutLink = Browser.FindElement(By.Id("about-link")).Text;
-            Assert.Contains("About", aboutLink);
+            AssertCustomNotFoundPageRendered();
         }
         else
         {
@@ -140,4 +136,67 @@ public void CanNavigateBetweenStaticPagesViaEnhancedNav()
         Browser.Equal("Global interactivity page: Static via attribute", () => h1.Text);
         Browser.Equal("static", () => Browser.Exists(By.Id("execution-mode")).Text);
     }
+
+    [Theory]
+    [InlineData(true)]
+    [InlineData(false)]
+    public void CanRenderNotFoundPage_SSR(bool streamingStarted)
+    {
+        string streamingPath = streamingStarted ? "-streaming" : "";
+        Navigate($"{ServerPathBase}/set-not-found-ssr{streamingPath}?useCustomNotFoundPage=true");
+        AssertCustomNotFoundPageRendered();
+    }
+
+    [Theory]
+    [InlineData("ServerNonPrerendered")]
+    [InlineData("WebAssemblyNonPrerendered")]
+    public void CanRenderNotFoundPage_Interactive(string renderMode)
+    {
+        Navigate($"{ServerPathBase}/set-not-found?useCustomNotFoundPage=true&renderMode={renderMode}");
+        AssertCustomNotFoundPageRendered();
+    }
+
+    private void AssertCustomNotFoundPageRendered()
+    {
+        var infoText = Browser.FindElement(By.Id("test-info")).Text;
+        Assert.Contains("Welcome On Custom Not Found Page", infoText);
+        // custom page should have a custom layout
+        var aboutLink = Browser.FindElement(By.Id("about-link")).Text;
+        Assert.Contains("About", aboutLink);
+    }
+
+    [Theory]
+    [InlineData(false)]
+    [InlineData(true)]
+    public void DoesNotReExecuteIf404WasHandled_SSR(bool streamingStarted)
+    {
+        string streamingPath = streamingStarted ? "-streaming" : "";
+        Navigate($"{ServerPathBase}/reexecution/set-not-found-ssr{streamingPath}");
+        AssertNotFoundFragmentRendered();
+    }
+
+    [Theory]
+    [InlineData("ServerNonPrerendered")]
+    [InlineData("WebAssemblyNonPrerendered")]
+    public void DoesNotReExecuteIf404WasHandled_Interactive(string renderMode)
+    {
+        Navigate($"{ServerPathBase}/reexecution/set-not-found?renderMode={renderMode}");
+        AssertNotFoundFragmentRendered();
+    }
+
+    private void AssertNotFoundFragmentRendered()
+    {
+        var body = Browser.FindElement(By.TagName("body"));
+        var notFound = Browser.FindElement(By.Id("not-found-fragment")).Text;
+        Browser.Equal("There's nothing here", () => Browser.FindElement(By.Id("not-found-fragment")).Text);
+    }
+
+    [Fact]
+    public void StatusCodePagesWithReExecution()
+    {
+        Navigate($"{ServerPathBase}/reexecution/trigger-404");
+        Assert404ReExecuted();
+    }
+    private void Assert404ReExecuted() =>
+        Browser.Equal("Welcome On Page Re-executed After Not Found Event", () => Browser.Exists(By.Id("test-info")).Text);
 }

src/Components/test/testassets/Components.Shared/ComponentThatSetsNotFound.razor

@@ -0,0 +1,26 @@
+@inject NavigationManager NavigationManager
+
+@if (!WaitForInteractivity || RendererInfo.IsInteractive)
+{
+    <PageTitle>Original page</PageTitle>
+
+    <p id="test-info">Any content</p>
+
+}
+
+@code{
+    [Parameter]
+    public bool StartStreaming { get; set; } = false;
+
+    [Parameter]
+    public bool WaitForInteractivity { get; set; } = false;
+
+    protected async override Task OnInitializedAsync()
+    {
+        if (StartStreaming)
+        {
+            await Task.Yield();
+        }
+        NavigationManager.NotFound();
+    }
+}

src/Components/test/testassets/Components.Shared/Components.Shared.csproj

@@ -0,0 +1,17 @@
+<Project Sdk="Microsoft.NET.Sdk.Razor">
+
+  <PropertyGroup>
+    <TargetFramework>$(DefaultNetCoreTargetFramework)</TargetFramework>
+    <Nullable>enable</Nullable>
+    <ImplicitUsings>enable</ImplicitUsings>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <SupportedPlatform Include="browser" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <Reference Include="Microsoft.AspNetCore.Components.Web" />
+  </ItemGroup>
+
+</Project>

src/Components/test/testassets/Components.Shared/PageThatSetsNotFound-no-streaming.razor

@@ -0,0 +1,11 @@
+@page "/reexecution/set-not-found-ssr"
+@page "/set-not-found-ssr"
+@attribute [StreamRendering(false)]
+
+@*
+    this page is used in global interactivity and no interactivity scenarios
+    the content is rendered on the server without streaming and might become
+    interactive later if interactivity was enabled in the app
+*@
+
+<Components.Shared.ComponentThatSetsNotFound />

src/Components/test/testassets/Components.Shared/PageThatSetsNotFound-streaming.razor

@@ -0,0 +1,11 @@
+@page "/reexecution/set-not-found-ssr-streaming"
+@page "/set-not-found-ssr-streaming"
+@attribute [StreamRendering(true)]
+
+@*
+    this page is used in global interactivity and no interactivity scenarios
+    the content is rendered on the server with streaming and might become
+    interactive later if interactivity was enabled in the app
+*@
+
+<Components.Shared.ComponentThatSetsNotFound StartStreaming="true"/>

src/Components/test/testassets/Components.Shared/_Imports.razor

@@ -0,0 +1 @@
+@using Microsoft.AspNetCore.Components.Web

src/Components/test/testassets/Components.TestServer/Components.TestServer.csproj

@@ -35,6 +35,7 @@
     <ProjectReference Include="..\BasicTestApp\BasicTestApp.csproj" />
     <ProjectReference Include="..\Components.WasmMinimal\Components.WasmMinimal.csproj" />
     <ProjectReference Include="..\Components.WasmRemoteAuthentication\Components.WasmRemoteAuthentication.csproj" />
+    <ProjectReference Include="..\Components.Shared\Components.Shared.csproj" />
   </ItemGroup>
 
   <ItemGroup>

src/Components/test/testassets/Components.TestServer/RazorComponentEndpointsNoInteractivityStartup.cs

@@ -50,6 +50,15 @@ public void Configure(IApplicationBuilder app, IWebHostEnvironment env)
         {
             app.Map("/reexecution", reexecutionApp =>
             {
+                app.Map("/trigger-404", trigger404App =>
+                {
+                    trigger404App.Run(async context =>
+                    {
+                        context.Response.StatusCode = 404;
+                        await context.Response.WriteAsync("Triggered a 404 status code.");
+                    });
+                });
+
                 if (!env.IsDevelopment())
                 {
                     app.UseExceptionHandler("/Error", createScopeForErrors: true);
@@ -62,7 +71,8 @@ public void Configure(IApplicationBuilder app, IWebHostEnvironment env)
                 reexecutionApp.UseAntiforgery();
                 reexecutionApp.UseEndpoints(endpoints =>
                 {
-                    endpoints.MapRazorComponents<TRootComponent>();
+                    endpoints.MapRazorComponents<TRootComponent>()
+                        .AddAdditionalAssemblies(Assembly.Load("Components.Shared"));
                 });
             });
 
@@ -83,7 +93,8 @@ private void ConfigureSubdirPipeline(IApplicationBuilder app, IWebHostEnvironmen
         app.UseAntiforgery();
         app.UseEndpoints(endpoints =>
         {
-            endpoints.MapRazorComponents<TRootComponent>();
+            endpoints.MapRazorComponents<TRootComponent>()
+                .AddAdditionalAssemblies(Assembly.Load("Components.Shared"));
         });
     }
 }