Fix reexecution mechanism.

Author: ilonatommy

src/Components/Endpoints/src/RazorComponentEndpointInvoker.cs

@@ -45,7 +45,11 @@ private async Task RenderComponentCore(HttpContext context)
             Log.InteractivityDisabledForErrorHandling(_logger);
         }
         _renderer.InitializeStreamingRenderingFraming(context, isErrorHandler, hasStatusCodePage);
-        EndpointHtmlRenderer.MarkAsAllowingEnhancedNavigation(context);
+        bool avoidEditingHeaders = hasStatusCodePage && context.Response.StatusCode == StatusCodes.Status404NotFound;
+        if (!avoidEditingHeaders)
+        {
+            EndpointHtmlRenderer.MarkAsAllowingEnhancedNavigation(context);
+        }
 
         var endpoint = context.GetEndpoint() ?? throw new InvalidOperationException($"An endpoint must be set on the '{nameof(HttpContext)}'.");
 
@@ -86,6 +90,8 @@ await _renderer.InitializeStandardComponentServicesAsync(
         await using var writer = new HttpResponseStreamWriter(context.Response.Body, Encoding.UTF8, defaultBufferSize, ArrayPool<byte>.Shared, ArrayPool<char>.Shared);
         using var bufferWriter = new BufferedTextWriter(writer);
 
+        int originalStatusCode = context.Response.StatusCode;
+
         // 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
         // component takes care of switching into your desired render mode when it produces its own output.
@@ -95,6 +101,16 @@ await _renderer.InitializeStandardComponentServicesAsync(
             ParameterView.Empty,
             waitForQuiescence: result.IsPost || isErrorHandler || hasStatusCodePage);
 
+        bool requresReexecution = originalStatusCode != context.Response.StatusCode && hasStatusCodePage;
+        if (requresReexecution)
+        {
+            // 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.
+            await bufferWriter.FlushAsync();
+            return;
+        }
+
         Task quiesceTask;
         if (!result.IsPost)
         {

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

@@ -84,6 +84,7 @@ private void SetNotFoundResponse(object? sender, EventArgs args)
             return;
         }
         _httpContext.Response.StatusCode = StatusCodes.Status404NotFound;
+        _httpContext.Response.ContentType = null;
         SignalRendererToFinishRendering();
     }
 

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

@@ -8,13 +8,15 @@
 using Microsoft.AspNetCore.Components.Web.HtmlRendering;
 using Microsoft.AspNetCore.Html;
 using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Logging;
 using static Microsoft.AspNetCore.Internal.LinkerFlags;
 
 namespace Microsoft.AspNetCore.Components.Endpoints;
 
 internal partial class EndpointHtmlRenderer
 {
     private static readonly object ComponentSequenceKey = new object();
+    private bool stopAddingTasks;
 
     protected override IComponent ResolveComponentForRenderMode([DynamicallyAccessedMembers(Component)] Type componentType, int? parentComponentId, IComponentActivator componentActivator, IComponentRenderMode renderMode)
     {
@@ -146,6 +148,7 @@ internal async ValueTask<PrerenderedComponentHtmlContent> RenderEndpointComponen
         {
             var component = BeginRenderingComponent(rootComponentType, parameters);
             var result = new PrerenderedComponentHtmlContent(Dispatcher, component);
+            stopAddingTasks = httpContext.Response.StatusCode == StatusCodes.Status404NotFound && waitForQuiescence;
 
             await WaitForResultReady(waitForQuiescence, result);
 
@@ -166,7 +169,52 @@ private async Task WaitForResultReady(bool waitForQuiescence, PrerenderedCompone
         }
         else if (_nonStreamingPendingTasks.Count > 0)
         {
-            await WaitForNonStreamingPendingTasks();
+            if (stopAddingTasks)
+            {
+                HandleNonStreamingTasks();
+            }
+            else
+            {
+                await WaitForNonStreamingPendingTasks();
+            }
+        }
+    }
+
+    public void HandleNonStreamingTasks()
+    {
+        if (NonStreamingPendingTasksCompletion == null)
+        {
+            // Iterate over the tasks and handle their exceptions
+            foreach (var task in _nonStreamingPendingTasks)
+            {
+                _ = GetErrorHandledTask(task); // Fire-and-forget with exception handling
+            }
+
+            // Clear the pending tasks since we are handling them
+            _nonStreamingPendingTasks.Clear();
+
+            // Mark the tasks as completed
+            NonStreamingPendingTasksCompletion = Task.CompletedTask;
+        }
+    }
+
+    private async Task GetErrorHandledTask(Task taskToHandle)
+    {
+        try
+        {
+            await taskToHandle;
+        }
+        catch (Exception ex)
+        {
+            // Ignore errors due to task cancellations.
+            if (!taskToHandle.IsCanceled)
+            {
+                _logger.LogError(
+                    ex,
+                    @"An exception occurred during non-streaming rendering.
+                    This exception will be ignored because the response
+                    is being discarded and the request is being re-executed.");
+            }
         }
     }
 

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

@@ -28,7 +28,8 @@ public void InitializeStreamingRenderingFraming(HttpContext httpContext, bool is
     {
         _isHandlingErrors = isErrorHandler;
         _hasStatusCodePage = hasStatusCodePage;
-        if (IsProgressivelyEnhancedNavigation(httpContext.Request))
+        bool avoidEditingHeaders = hasStatusCodePage && httpContext.Response.StatusCode == StatusCodes.Status404NotFound;
+        if (!avoidEditingHeaders && IsProgressivelyEnhancedNavigation(httpContext.Request))
         {
             var id = Guid.NewGuid().ToString();
             httpContext.Response.Headers.Add(_streamingRenderingFramingHeaderName, id);

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

@@ -44,6 +44,7 @@ internal partial class EndpointHtmlRenderer : StaticHtmlRenderer, IComponentPrer
     private HttpContext _httpContext = default!; // Always set at the start of an inbound call
     private ResourceAssetCollection? _resourceCollection;
     private bool _rendererIsStopped;
+    private readonly ILogger _logger;
 
     // The underlying Renderer always tracks the pending tasks representing *full* quiescence, i.e.,
     // when everything (regardless of streaming SSR) is fully complete. In this subclass we also track
@@ -56,6 +57,7 @@ public EndpointHtmlRenderer(IServiceProvider serviceProvider, ILoggerFactory log
     {
         _services = serviceProvider;
         _options = serviceProvider.GetRequiredService<IOptions<RazorComponentsServiceOptions>>().Value;
+        _logger = loggerFactory.CreateLogger("Microsoft.AspNetCore.Components.RenderTree.Renderer");
     }
 
     internal HttpContext? HttpContext => _httpContext;
@@ -163,6 +165,11 @@ protected override ComponentState CreateComponentState(int componentId, ICompone
 
     protected override void AddPendingTask(ComponentState? componentState, Task task)
     {
+        if (stopAddingTasks)
+        {
+            return;
+        }
+
         var streamRendering = componentState is null
             ? false
             : ((EndpointComponentState)componentState).StreamRendering;

src/Middleware/Diagnostics/src/PublicAPI.Unshipped.txt

@@ -1 +1,4 @@
 #nullable enable
+Microsoft.AspNetCore.Builder.StatusCodePagesOptions.CreateScopeForErrors.get -> bool
+Microsoft.AspNetCore.Builder.StatusCodePagesOptions.CreateScopeForErrors.set -> void
+static Microsoft.AspNetCore.Builder.StatusCodePagesExtensions.UseStatusCodePagesWithReExecute(this Microsoft.AspNetCore.Builder.IApplicationBuilder! app, string! pathFormat, string? queryFormat = null, bool createScopeForErrors = false) -> Microsoft.AspNetCore.Builder.IApplicationBuilder!

src/Middleware/Diagnostics/src/StatusCodePage/StatusCodePagesExtensions.cs

@@ -1,11 +1,13 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Diagnostics.CodeAnalysis;
 using System.Globalization;
 using Microsoft.AspNetCore.Diagnostics;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Routing;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Options;
 
 namespace Microsoft.AspNetCore.Builder;
@@ -160,10 +162,50 @@ public static IApplicationBuilder UseStatusCodePagesWithReExecute(
         return app.UseStatusCodePages(CreateHandler(pathFormat, queryFormat));
     }
 
+    /// <summary>
+    /// Adds a StatusCodePages middleware to the pipeline. Specifies that the response body should be generated by
+    /// re-executing the request pipeline using an alternate path. This path may contain a '{0}' placeholder of the status code.
+    /// </summary>
+    /// <param name="app"></param>
+    /// <param name="pathFormat"></param>
+    /// <param name="queryFormat"></param>
+    /// <param name="createScopeForErrors">Whether or not to create a new <see cref="IServiceProvider"/> scope.</param>
+    /// <returns></returns>
+    [SuppressMessage("ApiDesign", "RS0026:Do not add multiple overloads with optional parameters", Justification = "Required to maintain compatibility")]
+    public static IApplicationBuilder UseStatusCodePagesWithReExecute(
+        this IApplicationBuilder app,
+        string pathFormat,
+        string? queryFormat = null,
+        bool createScopeForErrors = false)
+    {
+        ArgumentNullException.ThrowIfNull(app);
+
+        // Only use this path if there's a global router (in the 'WebApplication' case).
+        if (app.Properties.TryGetValue(RerouteHelper.GlobalRouteBuilderKey, out var routeBuilder) && routeBuilder is not null)
+        {
+            return app.Use(next =>
+            {
+                var newNext = RerouteHelper.Reroute(app, routeBuilder, next);
+                return new StatusCodePagesMiddleware(next,
+                    Options.Create(new StatusCodePagesOptions() {
+                        HandleAsync = CreateHandler(pathFormat, queryFormat, newNext),
+                        CreateScopeForErrors = createScopeForErrors
+                    })).Invoke;
+            });
+        }
+
+        return app.UseStatusCodePages(new StatusCodePagesOptions
+        {
+            HandleAsync = CreateHandler(pathFormat, queryFormat),
+            CreateScopeForErrors = createScopeForErrors
+        });
+    }
+
     private static Func<StatusCodeContext, Task> CreateHandler(string pathFormat, string? queryFormat, RequestDelegate? next = null)
     {
         var handler = async (StatusCodeContext context) =>
         {
+            // context.Options.CreateScopeForErrors
             var originalStatusCode = context.HttpContext.Response.StatusCode;
 
             var newPath = new PathString(
@@ -176,6 +218,10 @@ private static Func<StatusCodeContext, Task> CreateHandler(string pathFormat, st
             var originalQueryString = context.HttpContext.Request.QueryString;
 
             var routeValuesFeature = context.HttpContext.Features.Get<IRouteValuesFeature>();
+            var oldScope = context.Options.CreateScopeForErrors ? context.HttpContext.RequestServices : null;
+            await using AsyncServiceScope? scope = context.Options.CreateScopeForErrors
+                ? context.HttpContext.RequestServices.GetRequiredService<IServiceScopeFactory>().CreateAsyncScope() // or .GetRequiredService<IServiceScopeFactory>().CreateAsyncScope()
+                : null;
 
             // Store the original paths so the app can check it.
             context.HttpContext.Features.Set<IStatusCodeReExecuteFeature>(new StatusCodeReExecuteFeature()
@@ -188,6 +234,11 @@ private static Func<StatusCodeContext, Task> CreateHandler(string pathFormat, st
                 RouteValues = routeValuesFeature?.RouteValues
             });
 
+            if (scope.HasValue)
+            {
+                context.HttpContext.RequestServices = scope.Value.ServiceProvider;
+            }
+
             // An endpoint may have already been set. Since we're going to re-invoke the middleware pipeline we need to reset
             // the endpoint and route values to ensure things are re-calculated.
             HttpExtensions.ClearEndpoint(context.HttpContext);
@@ -210,6 +261,10 @@ private static Func<StatusCodeContext, Task> CreateHandler(string pathFormat, st
                 context.HttpContext.Request.QueryString = originalQueryString;
                 context.HttpContext.Request.Path = originalPath;
                 context.HttpContext.Features.Set<IStatusCodeReExecuteFeature?>(null);
+                if (oldScope != null)
+                {
+                    context.HttpContext.RequestServices = oldScope;
+                }
             }
         };
 

src/Middleware/Diagnostics/src/StatusCodePage/StatusCodePagesOptions.cs

@@ -55,4 +55,11 @@ private static string BuildResponseBody(int httpStatusCode)
     /// The handler that generates the response body for the given <see cref="StatusCodeContext"/>. By default this produces a plain text response that includes the status code.
     /// </summary>
     public Func<StatusCodeContext, Task> HandleAsync { get; set; }
+
+    /// <summary>
+    /// Gets or sets whether the handler needs to create a separate <see cref="IServiceProvider"/> scope and
+    /// replace it on <see cref="HttpContext.RequestServices"/> when re-executing the request.
+    /// </summary>
+    /// <remarks>The default value is <see langword="false"/>.</remarks>
+    public bool CreateScopeForErrors { get; set; }
 }