Make EndpointHtmlRenderer track the HttpContext

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

@@ -37,6 +37,7 @@ public async ValueTask<IHtmlAsyncContent> PrerenderComponentAsync(
 
         // Make sure we only initialize the services once, but on every call we wait for that process to complete
         // This does not have to be threadsafe since it's not valid to call this simultaneously from multiple threads.
+        SetHttpContext(httpContext);
         _servicesInitializedTask ??= InitializeStandardComponentServicesAsync(httpContext);
         await _servicesInitializedTask;
 
@@ -46,8 +47,8 @@ public async ValueTask<IHtmlAsyncContent> PrerenderComponentAsync(
         {
             var result = prerenderMode switch
             {
-                RenderMode.Server => NonPrerenderedServerComponent(httpContext, GetOrCreateInvocationId(httpContext), componentType, parameters),
-                RenderMode.ServerPrerendered => await PrerenderedServerComponentAsync(httpContext, GetOrCreateInvocationId(httpContext), componentType, parameters),
+                RenderMode.Server => NonPrerenderedServerComponent(GetOrCreateInvocationId(httpContext), componentType, parameters),
+                RenderMode.ServerPrerendered => await PrerenderedServerComponentAsync(GetOrCreateInvocationId(httpContext), componentType, parameters),
                 RenderMode.Static => await StaticComponentAsync(componentType, parameters),
                 RenderMode.WebAssembly => NonPrerenderedWebAssemblyComponent(componentType, parameters),
                 RenderMode.WebAssemblyPrerendered => await PrerenderedWebAssemblyComponentAsync(componentType, parameters),
@@ -70,6 +71,8 @@ internal async ValueTask<PrerenderedComponentHtmlContent> RenderEndpointComponen
         ParameterView parameters,
         bool waitForQuiescence)
     {
+        SetHttpContext(httpContext);
+
         try
         {
             var component = BeginRenderingComponent(componentType, parameters);
@@ -81,7 +84,7 @@ internal async ValueTask<PrerenderedComponentHtmlContent> RenderEndpointComponen
         }
         catch (NavigationException navigationException)
         {
-            return await HandleNavigationException(httpContext, navigationException);
+            return await HandleNavigationException(_httpContext, navigationException);
         }
     }
 
@@ -137,11 +140,11 @@ private async Task<PrerenderedComponentHtmlContent> StaticComponentAsync(Type ty
         return new PrerenderedComponentHtmlContent(Dispatcher, htmlComponent, null, null);
     }
 
-    private async Task<PrerenderedComponentHtmlContent> PrerenderedServerComponentAsync(HttpContext context, ServerComponentInvocationSequence invocationId, Type type, ParameterView parametersCollection)
+    private async Task<PrerenderedComponentHtmlContent> PrerenderedServerComponentAsync(ServerComponentInvocationSequence invocationId, Type type, ParameterView parametersCollection)
     {
-        if (!context.Response.HasStarted)
+        if (!_httpContext.Response.HasStarted)
         {
-            context.Response.Headers.CacheControl = "no-cache, no-store, max-age=0";
+            _httpContext.Response.Headers.CacheControl = "no-cache, no-store, max-age=0";
         }
 
         // Lazy because we don't actually want to require a whole chain of services including Data Protection
@@ -169,11 +172,11 @@ private async ValueTask<PrerenderedComponentHtmlContent> PrerenderedWebAssemblyC
         return new PrerenderedComponentHtmlContent(Dispatcher, htmlComponent, null, marker);
     }
 
-    private PrerenderedComponentHtmlContent NonPrerenderedServerComponent(HttpContext context, ServerComponentInvocationSequence invocationId, Type type, ParameterView parametersCollection)
+    private PrerenderedComponentHtmlContent NonPrerenderedServerComponent(ServerComponentInvocationSequence invocationId, Type type, ParameterView parametersCollection)
     {
-        if (!context.Response.HasStarted)
+        if (!_httpContext.Response.HasStarted)
         {
-            context.Response.Headers.CacheControl = "no-cache, no-store, max-age=0";
+            _httpContext.Response.Headers.CacheControl = "no-cache, no-store, max-age=0";
         }
 
         // Lazy because we don't actually want to require a whole chain of services including Data Protection

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

@@ -16,10 +16,12 @@ internal partial class EndpointHtmlRenderer
 
     public async ValueTask<IHtmlContent> PrerenderPersistedStateAsync(HttpContext httpContext, PersistedStateSerializationMode serializationMode)
     {
+        SetHttpContext(httpContext);
+
         // First we resolve "infer" mode to a specific mode
         if (serializationMode == PersistedStateSerializationMode.Infer)
         {
-            switch (GetPersistStateRenderMode(httpContext))
+            switch (GetPersistStateRenderMode(_httpContext))
             {
                 case InvokedRenderModes.Mode.None:
                     return ComponentStateHtmlContent.Empty;
@@ -41,15 +43,15 @@ public async ValueTask<IHtmlContent> PrerenderPersistedStateAsync(HttpContext ht
         var store = serializationMode switch
         {
             PersistedStateSerializationMode.Server =>
-                new ProtectedPrerenderComponentApplicationStore(httpContext.RequestServices.GetRequiredService<IDataProtectionProvider>()),
+                new ProtectedPrerenderComponentApplicationStore(_httpContext.RequestServices.GetRequiredService<IDataProtectionProvider>()),
             PersistedStateSerializationMode.WebAssembly =>
                 new PrerenderComponentApplicationStore(),
             _ =>
                 throw new InvalidOperationException("Invalid persistence mode.")
         };
 
         // Finally, persist the state and return the HTML content
-        var manager = httpContext.RequestServices.GetRequiredService<ComponentStatePersistenceManager>();
+        var manager = _httpContext.RequestServices.GetRequiredService<ComponentStatePersistenceManager>();
         await manager.PersistStateAsync(store, Dispatcher);
         return new ComponentStateHtmlContent(store);
     }

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

@@ -18,6 +18,8 @@ internal partial class EndpointHtmlRenderer
 
     public async Task SendStreamingUpdatesAsync(HttpContext httpContext, Task untilTaskCompleted, TextWriter writer)
     {
+        SetHttpContext(httpContext);
+
         if (_streamingUpdatesWriter is not null)
         {
             // The framework is the only caller, so it's OK to have a nonobvious restriction like this
@@ -37,7 +39,7 @@ public async Task SendStreamingUpdatesAsync(HttpContext httpContext, Task untilT
         }
         catch (Exception ex)
         {
-            HandleExceptionAfterResponseStarted(httpContext, writer, ex);
+            HandleExceptionAfterResponseStarted(_httpContext, writer, ex);
 
             // The rest of the pipeline can treat this as a regular unhandled exception
             // TODO: Is this really right? I think we'll terminate the response in an invalid way.

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

@@ -36,6 +36,7 @@ internal partial class EndpointHtmlRenderer : StaticHtmlRenderer, IComponentPrer
     private readonly IServiceProvider _services;
     private Task? _servicesInitializedTask;
 
+    private HttpContext _httpContext = default!; // Always set at the start of an inbound call
     private string? _formHandler;
     private NamedEvent _capturedNamedEvent;
 
@@ -51,6 +52,18 @@ public EndpointHtmlRenderer(IServiceProvider serviceProvider, ILoggerFactory log
         _services = serviceProvider;
     }
 
+    private void SetHttpContext(HttpContext httpContext)
+    {
+        if (_httpContext is null)
+        {
+            _httpContext = httpContext;
+        }
+        else if (_httpContext != httpContext)
+        {
+            throw new InvalidOperationException("The HttpContext cannot change value once assigned.");
+        }
+    }
+
     internal static async Task InitializeStandardComponentServicesAsync(
         HttpContext httpContext,
         string? handler = null,