Fix: move setting response dispatch to invoker.

Author: ilonatommy

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

@@ -384,12 +384,12 @@ private void OnLocationChanged(object sender, LocationChangedEventArgs args)
 
     private void OnNotFound(object sender, NotFoundEventArgs args)
     {
-        if (_renderHandle.IsInitialized)
+        if (_renderHandle.IsInitialized && NotFoundPage != null)
         {
+            // setting the path signals to the endpoint renderer that router handled rendering
+            args.Path = _notFoundPageRoute;
             Log.DisplayingNotFound(_logger);
             RenderNotFound();
-            // setting the path signals to the endpoint renderer that router handled rendering
-            args.Path = _notFoundPageRoute ?? "not-found-handled-by-router";
         }
     }
 

src/Components/Endpoints/src/RazorComponentEndpointInvoker.cs

@@ -136,6 +136,11 @@ await _renderer.InitializeStandardComponentServicesAsync(
             }
         }
 
+        if (_renderer.NotFoundEventArgs != null)
+        {
+            _renderer.SetNotFoundWhenResponseNotStarted();
+        }
+
         if (!quiesceTask.IsCompleted)
         {
             // An incomplete QuiescenceTask indicates there may be streaming rendering updates.
@@ -155,6 +160,10 @@ await _renderer.InitializeStandardComponentServicesAsync(
         if (!quiesceTask.IsCompletedSuccessfully)
         {
             await _renderer.SendStreamingUpdatesAsync(context, quiesceTask, bufferWriter);
+            if (_renderer.NotFoundEventArgs != null)
+            {
+                await _renderer.SetNotFoundWhenResponseHasStarted();
+            }
         }
         else
         {

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

@@ -79,37 +79,38 @@ private Task ReturnErrorResponse(string detailedMessage)
             : Task.CompletedTask;
     }
 
-    internal async Task SetNotFoundResponseAsync(string baseUri, NotFoundEventArgs args)
+    internal void SetNotFoundWhenResponseNotStarted()
     {
-        if (_httpContext.Response.HasStarted ||
-            // POST waits for quiescence -> rendering the NotFoundPage would be queued for the next batch
-            // but we want to send the signal to the renderer to stop rendering future batches -> use client rendering
-            HttpMethods.IsPost(_httpContext.Request.Method))
-        {
-            if (string.IsNullOrEmpty(_notFoundUrl))
-            {
-                _notFoundUrl = GetNotFoundUrl(baseUri, args);
-            }
-            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);
-            HandleNotFoundAfterResponseStarted(bufferWriter, _httpContext, _notFoundUrl);
-            await bufferWriter.FlushAsync();
-        }
-        else
+        _httpContext.Response.StatusCode = StatusCodes.Status404NotFound;
+
+        // 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();
+    }
+
+    internal async Task SetNotFoundWhenResponseHasStarted()
+    {
+        if (string.IsNullOrEmpty(_notFoundUrl))
         {
-            _httpContext.Response.StatusCode = StatusCodes.Status404NotFound;
+            var baseUri = $"{_httpContext.Request.Scheme}://{_httpContext.Request.Host}{_httpContext.Request.PathBase}/";
+            _notFoundUrl = GetNotFoundUrl(baseUri, NotFoundEventArgs);
         }
+        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);
+        HandleNotFoundAfterResponseStarted(bufferWriter, _httpContext, _notFoundUrl);
+        await bufferWriter.FlushAsync();
 
         // 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 string GetNotFoundUrl(string baseUri, NotFoundEventArgs args)
+    private string GetNotFoundUrl(string baseUri, NotFoundEventArgs? args)
     {
-        string? path = args.Path;
+        string? path = args?.Path;
         if (string.IsNullOrEmpty(path))
         {
             var pathFormat = _httpContext.Items[nameof(StatusCodePagesOptions)] as string;

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

@@ -63,7 +63,7 @@ public EndpointHtmlRenderer(IServiceProvider serviceProvider, ILoggerFactory log
     }
 
     internal HttpContext? HttpContext => _httpContext;
-    internal NotFoundEventArgs? NotFoundEventArgs { get; set; }
+    internal NotFoundEventArgs? NotFoundEventArgs { get; private set; }
 
     internal void SetHttpContext(HttpContext httpContext)
     {
@@ -86,11 +86,7 @@ internal async Task InitializeStandardComponentServicesAsync(
         var navigationManager = httpContext.RequestServices.GetRequiredService<NavigationManager>();
         ((IHostEnvironmentNavigationManager)navigationManager)?.Initialize(GetContextBaseUri(httpContext.Request), GetFullUri(httpContext.Request), OnNavigateTo);
 
-        navigationManager?.OnNotFound += (sender, args) =>
-        {
-            NotFoundEventArgs = args;
-            _ = GetErrorHandledTask(SetNotFoundResponseAsync(navigationManager.BaseUri, args));
-        };
+        navigationManager?.OnNotFound += (sender, args) => NotFoundEventArgs = args;
 
         var authenticationStateProvider = httpContext.RequestServices.GetService<AuthenticationStateProvider>();
         if (authenticationStateProvider is IHostEnvironmentAuthenticationStateProvider hostEnvironmentAuthenticationStateProvider)

src/Components/Endpoints/test/EndpointHtmlRendererTest.cs

@@ -1823,7 +1823,7 @@ protected override void ProcessPendingRender()
         public async Task SetNotFoundResponseAsync(HttpContext httpContext, NotFoundEventArgs args)
         {
             SetHttpContext(httpContext);
-            await SetNotFoundResponseAsync(httpContext.Request.PathBase, args);
+            await SetNotFoundWhenResponseHasStarted();
         }
     }
 

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

@@ -239,7 +239,7 @@ private void AssertNotFoundRendered_ResponseStarted_Or_POST(bool hasReExecutionM
     public void NotFoundSetOnFormSubmit_ResponseNotStarted_SSR(bool hasReExecutionMiddleware, bool hasCustomNotFoundPageSet)
     {
         string reexecution = hasReExecutionMiddleware ? "/reexecution" : "";
-        string testUrl = $"{ServerPathBase}{reexecution}/post-not-found-ssr-streaming?useCustomNotFoundPage={hasCustomNotFoundPageSet}";
+        string testUrl = $"{ServerPathBase}{reexecution}/post-not-found-ssr?useCustomNotFoundPage={hasCustomNotFoundPageSet}";
         Navigate(testUrl);
         Browser.FindElement(By.Id("not-found-form")).FindElement(By.TagName("button")).Click();