Blazor 404 code re-execution is compatible with OnNavigateAsync (#64034)

Author: ilonatommy

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

@@ -168,30 +168,7 @@ private async Task WaitForResultReady(bool waitForQuiescence, PrerenderedCompone
         }
         else if (_nonStreamingPendingTasks.Count > 0)
         {
-            if (_isReExecuted)
-            {
-                HandleNonStreamingTasks();
-            }
-            else
-            {
-                await WaitForNonStreamingPendingTasks();
-            }
-        }
-    }
-
-    public void HandleNonStreamingTasks()
-    {
-        if (NonStreamingPendingTasksCompletion == null)
-        {
-            foreach (var task in _nonStreamingPendingTasks)
-            {
-                _ = GetErrorHandledTask(task);
-            }
-
-            // Clear the pending tasks since we are handling them
-            _nonStreamingPendingTasks.Clear();
-
-            NonStreamingPendingTasksCompletion = Task.CompletedTask;
+            await WaitForNonStreamingPendingTasks();
         }
     }
 

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

@@ -85,8 +85,8 @@ internal async Task InitializeStandardComponentServicesAsync(
     {
         var navigationManager = httpContext.RequestServices.GetRequiredService<NavigationManager>();
         ((IHostEnvironmentNavigationManager)navigationManager)?.Initialize(
-            GetContextBaseUri(httpContext.Request), 
-            GetFullUri(httpContext.Request), 
+            GetContextBaseUri(httpContext.Request),
+            GetFullUri(httpContext.Request),
             uri => GetErrorHandledTask(OnNavigateTo(uri)));
 
         navigationManager?.OnNotFound += (sender, args) => NotFoundEventArgs = args;
@@ -168,11 +168,6 @@ protected override ComponentState CreateComponentState(int componentId, ICompone
 
     protected override void AddPendingTask(ComponentState? componentState, Task task)
     {
-        if (_isReExecuted)
-        {
-            return;
-        }
-
         var streamRendering = componentState is null
             ? false
             : ((EndpointComponentState)componentState).StreamRendering;

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

@@ -1,7 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System.Net.Http;
+using System;
 using Components.TestServer.RazorComponents;
 using Microsoft.AspNetCore.Components.E2ETest;
 using Microsoft.AspNetCore.Components.E2ETest.Infrastructure;
@@ -125,6 +125,23 @@ public void BrowserNavigationToNotExistingPath_ReExecutesTo404(bool streaming)
         AssertReExecutionPageRendered();
     }
 
+    [Fact]
+    public void BrowserNavigationToNotExistingPath_WithOnNavigateAsync_ReExecutesTo404()
+    {
+        AppContext.SetSwitch("Microsoft.AspNetCore.Components.Endpoints.NavigationManager.DisableThrowNavigationException", isEnabled: true);
+        Navigate($"{ServerPathBase}/reexecution/not-existing-page?useOnNavigateAsync=true");
+        AssertReExecutionPageRendered();
+    }
+
+    [Fact]
+    public void BrowserNavigationToNotExistingPath_WithOnNavigateAsync_ReExecutesTo404_CanStream()
+    {
+        AppContext.SetSwitch("Microsoft.AspNetCore.Components.Endpoints.NavigationManager.DisableThrowNavigationException", isEnabled: true);
+        Navigate($"{ServerPathBase}/streaming-reexecution/not-existing-page?useOnNavigateAsync=true");
+        AssertReExecutionPageRendered();
+        Browser.Equal("Streaming completed.", () => Browser.Exists(By.Id("reexecute-streaming-status")).Text);
+    }
+
     private void AssertReExecutionPageRendered() =>
         Browser.Equal("Welcome On Page Re-executed After Not Found Event", () => Browser.Exists(By.Id("test-info")).Text);
 
@@ -192,7 +209,7 @@ public void NotFoundSetOnInitialization_ResponseNotStarted_SSR(bool hasReExecuti
     [InlineData(false, true)]
     [InlineData(false, false)]
     // This tests the application subscribing to OnNotFound event and setting NotFoundEventArgs.Path, opposed to the framework doing it for the app.
-    public void NotFoundSetOnInitialization_ApplicationSubscribesToNotFoundEventToSetNotFoundPath_SSR (bool streaming, bool customRouter)
+    public void NotFoundSetOnInitialization_ApplicationSubscribesToNotFoundEventToSetNotFoundPath_SSR(bool streaming, bool customRouter)
     {
         string streamingPath = streaming ? "-streaming" : "";
         string testUrl = $"{ServerPathBase}/set-not-found-ssr{streamingPath}?useCustomRouter={customRouter}&appSetsEventArgsPath=true";

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

@@ -8,6 +8,7 @@
 using Components.TestServer.RazorComponents;
 using Components.TestServer.RazorComponents.Pages.Forms;
 using Components.TestServer.Services;
+using Microsoft.AspNetCore.Diagnostics;
 using Microsoft.AspNetCore.Mvc;
 
 namespace TestServer;
@@ -64,7 +65,7 @@ public void Configure(IApplicationBuilder app, IWebHostEnvironment env)
                     app.UseExceptionHandler("/Error", createScopeForErrors: true);
                 }
 
-                reexecutionApp.UseStatusCodePagesWithReExecute("/not-found-reexecute", createScopeForStatusCodePages: true);
+                ConfigureReexecutionPipeline(reexecutionApp, "/not-found-reexecute");
                 reexecutionApp.UseStaticFiles();
                 reexecutionApp.UseRouting();
                 RazorComponentEndpointsStartup<TRootComponent>.UseFakeAuthState(reexecutionApp);
@@ -75,11 +76,37 @@ public void Configure(IApplicationBuilder app, IWebHostEnvironment env)
                         .AddAdditionalAssemblies(Assembly.Load("TestContentPackage"));
                 });
             });
+            app.Map("/streaming-reexecution", reexecutionApp =>
+            {
+                ConfigureReexecutionPipeline(reexecutionApp, "/not-found-reexecute-streaming");
+                reexecutionApp.UseRouting();
+                reexecutionApp.UseAntiforgery();
+                reexecutionApp.UseEndpoints(endpoints =>
+                {
+                    endpoints.MapRazorComponents<TRootComponent>()
+                        .AddAdditionalAssemblies(Assembly.Load("TestContentPackage"));
+                });
+            });
 
             ConfigureSubdirPipeline(app, env);
         });
     }
 
+    private void ConfigureReexecutionPipeline(IApplicationBuilder pipeline, string pathFormat)
+    {
+        pipeline.UseStatusCodePagesWithReExecute(pathFormat, createScopeForStatusCodePages: true);
+        pipeline.Use(async (context, next) =>
+        {
+            var reexecute = context.Features.Get<IStatusCodeReExecuteFeature>();
+            if (reexecute is not null && !string.IsNullOrEmpty(reexecute.OriginalQueryString))
+            {
+                context.Request.QueryString = new QueryString(reexecute.OriginalQueryString);
+            }
+
+            await next();
+        });
+    }
+
     private void ConfigureSubdirPipeline(IApplicationBuilder app, IWebHostEnvironment env)
     {
         if (!env.IsDevelopment())

src/Components/test/testassets/Components.TestServer/RazorComponents/App.razor

@@ -3,6 +3,10 @@
 @using Components.WasmMinimal.Pages.NotFound
 @using TestContentPackage.NotFound
 @using Components.TestServer.RazorComponents
+@using Microsoft.AspNetCore.Components
+@using Microsoft.AspNetCore.Components.Routing
+@using Microsoft.AspNetCore.Components.Web
+@using System.Threading.Tasks
 
 @code {
     [Parameter]
@@ -17,6 +21,10 @@
     [SupplyParameterFromQuery(Name = "appSetsEventArgsPath")]
     public bool AppSetsEventArgsPath { get; set; }
 
+    [Parameter]
+    [SupplyParameterFromQuery(Name = "useOnNavigateAsync")]
+    public bool ShouldDelayOnNavigateAsync { get; set; }
+
     private Type? NotFoundPageType { get; set; }
     private NavigationManager _navigationManager = default!;
 
@@ -70,6 +78,21 @@
             _navigationManager.OnNotFound -= OnNotFoundEvent;
         }
     }
+
+    private Task HandleOnNavigateAsync(NavigationContext args)
+    {
+        if (!ShouldDelayOnNavigateAsync)
+        {
+            return Task.CompletedTask;
+        }
+
+        return PerformOnNavigateAsyncWork();
+    }
+
+    private async Task PerformOnNavigateAsyncWork()
+    {
+        await Task.Yield();
+    }
 }
 
 <!DOCTYPE html>
@@ -93,7 +116,7 @@
     {
         @if (NotFoundPageType is not null)
         {
-            <Router AppAssembly="@typeof(App).Assembly" AdditionalAssemblies="new[] { typeof(TestContentPackage.NotFound.NotFoundPage).Assembly }" NotFoundPage="NotFoundPageType">
+            <Router AppAssembly="@typeof(App).Assembly" AdditionalAssemblies="new[] { typeof(TestContentPackage.NotFound.NotFoundPage).Assembly }" NotFoundPage="NotFoundPageType" OnNavigateAsync="HandleOnNavigateAsync">
                 <Found Context="routeData">
                     <RouteView RouteData="@routeData" />
                     <FocusOnNavigate RouteData="@routeData" Selector="[data-focus-on-navigate]" />
@@ -102,7 +125,7 @@
         }
         else
         {
-            <Router AppAssembly="@typeof(App).Assembly" AdditionalAssemblies="new[] { typeof(TestContentPackage.NotFound.NotFoundPage).Assembly }">
+            <Router AppAssembly="@typeof(App).Assembly" AdditionalAssemblies="new[] { typeof(TestContentPackage.NotFound.NotFoundPage).Assembly }" OnNavigateAsync="HandleOnNavigateAsync">
                 <Found Context="routeData">
                     <RouteView RouteData="@routeData" />
                     <FocusOnNavigate RouteData="@routeData" Selector="[data-focus-on-navigate]" />

src/Components/test/testassets/TestContentPackage/NotFound/ReexecutedPageStreaming.razor

@@ -0,0 +1,17 @@
+@page "/not-found-reexecute-streaming"
+@attribute [StreamRendering(true)]
+
+<ReExecutedComponent/>
+
+<p id="reexecute-streaming-status">@_status</p>
+
+@code {
+    private string _status = "Streaming in progress...";
+
+    protected override async Task OnInitializedAsync()
+    {
+        await Task.Yield();
+        _status = "Streaming completed.";
+        await InvokeAsync(StateHasChanged);
+    }
+}