Restore default request service feature after middleware completes

extensions/Worker.Extensions.Http.AspNetCore/src/AspNetMiddleware/WorkerRequestServicesMiddleware.cs

@@ -27,15 +27,25 @@ public async Task Invoke(HttpContext context)
 
         FunctionContext functionContext = await _coordinator.SetHttpContextAsync(invocationId, context);
 
+        // Retrieve the existing service provider feature
+        var existingFeature = context.Features.Get<IServiceProvidersFeature>();
+
         // Explicitly set the RequestServices to prevent a new scope from being created internally.
         // This also prevents the scope from being disposed when the request is complete. We want this to 
         // be disposed in the Functions middleware, not here.
-        var servicesFeature = new RequestServicesFeature(context, null)
+        await using var servicesFeature = new RequestServicesFeature(context, null)
         {
             RequestServices = functionContext.InstanceServices
         };
         context.Features.Set<IServiceProvidersFeature>(servicesFeature);
 
-        await _next(context);
+        try
+        {
+            await _next(context);
+        }
+        finally
+        {
+            context.Features.Set(existingFeature);
+        }
     }
 }

test/DotNetWorkerTests/AspNetCore/WorkerRequestServicesMiddlewareTests.cs

@@ -48,10 +48,45 @@ public async Task ServiceProviders_Equal()
             var middleware = new WorkerRequestServicesMiddleware(next, httpCoordinator.Object);
             await middleware.Invoke(httpContext);
         }
-
-        private class MyService
+
+        [Fact]
+        public async Task DefaultServiceProvider_Restored()
         {
+            var services = new ServiceCollection();
+            services.AddScoped<MyService>();
+
+            var httpContextProvider = services.BuildServiceProvider().CreateScope().ServiceProvider;
+
+            // FunctionContext services will be scoped.
+            var functionContextProvider = services.BuildServiceProvider().CreateScope().ServiceProvider;
+
+            RequestDelegate next = (ctxt) =>
+            {
+                Assert.Same(functionContextProvider, ctxt.RequestServices);
+
+                return Task.CompletedTask;
+            };
+
+            var functionContext = new TestFunctionContext(new TestFunctionDefinition(), new TestFunctionInvocation(), CancellationToken.None,
+                serviceProvider: functionContextProvider);
+
+            var httpContext = new DefaultHttpContext();
+            httpContext.RequestServices = httpContextProvider;
+            httpContext.Request.Headers.Add(Constants.CorrelationHeader, functionContext.InvocationId);
+
+            var httpCoordinator = new Mock<IHttpCoordinator>();
+            httpCoordinator
+                .Setup(p => p.SetHttpContextAsync(functionContext.InvocationId, httpContext))
+                .ReturnsAsync(functionContext);
+
+            var middleware = new WorkerRequestServicesMiddleware(next, httpCoordinator.Object);
+            await middleware.Invoke(httpContext);
+
+            Assert.Same(httpContextProvider, httpContext.RequestServices);
+            Assert.NotSame(functionContextProvider, httpContext.RequestServices);
         }
+
+        private class MyService;
     }
 }