Set X-MS-COLDSTART header in Flex environments

src/WebJobs.Script.WebHost/Middleware/PlaceholderSpecializationMiddleware.cs

@@ -4,6 +4,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
+using Microsoft.Azure.WebJobs.Script.Extensions;
 
 namespace Microsoft.Azure.WebJobs.Script.WebHost.Middleware
 {
@@ -47,6 +48,13 @@ private async Task InvokeSpecializationCheck(HttpContext httpContext)
 
                 if (Interlocked.CompareExchange(ref _specialized, 1, 0) == 0)
                 {
+                    // For Flex environments, set the cold start header during specialization
+                    // since it's not automatically set by the platform like in Windows Consumption
+                    if (_environment.IsFlexConsumptionSku())
+                    {
+                        httpContext.Request.Headers[ScriptConstants.AntaresColdStartHeaderName] = "1";
+                    }
+
                     Interlocked.Exchange(ref _invoke, _next);
                 }
             }

test/WebJobs.Script.Tests/Middleware/PlaceholderSpecializationMiddlewareTests.cs

@@ -0,0 +1,114 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Azure.WebJobs.Script.WebHost.Middleware;
+using Microsoft.Extensions.Logging.Abstractions;
+using Moq;
+using Xunit;
+
+namespace Microsoft.Azure.WebJobs.Script.Tests.Middleware
+{
+    public class PlaceholderSpecializationMiddlewareTests
+    {
+        [Fact]
+        public async Task Invoke_FlexConsumption_AddsHeader_When_Specialized()
+        {
+            // Arrange
+            var webHostEnvironment = new Mock<IScriptWebHostEnvironment>();
+            var standbyManager = new Mock<IStandbyManager>();
+            var environment = new Mock<IEnvironment>();
+            var httpContext = new DefaultHttpContext();
+
+            // Setup conditions for specialization
+            webHostEnvironment.Setup(e => e.InStandbyMode).Returns(false);
+            environment.Setup(e => e.IsContainerReady()).Returns(true);
+            environment.Setup(e => e.IsFlexConsumptionSku()).Returns(true);
+
+            standbyManager.Setup(s => s.SpecializeHostAsync()).Returns(Task.CompletedTask);
+
+            var nextCalled = false;
+            RequestDelegate next = (ctx) =>
+            {
+                nextCalled = true;
+                return Task.CompletedTask;
+            };
+
+            var middleware = new PlaceholderSpecializationMiddleware(next, webHostEnvironment.Object, standbyManager.Object, environment.Object);
+
+            // Act
+            await middleware.Invoke(httpContext);
+
+            // Assert
+            Assert.True(nextCalled);
+            Assert.Equal("1", httpContext.Request.Headers[ScriptConstants.AntaresColdStartHeaderName]);
+            standbyManager.Verify(s => s.SpecializeHostAsync(), Times.Once);
+        }
+
+        [Fact]
+        public async Task Invoke_NonFlexConsumption_DoesNotAddHeader_When_Specialized()
+        {
+            // Arrange
+            var webHostEnvironment = new Mock<IScriptWebHostEnvironment>();
+            var standbyManager = new Mock<IStandbyManager>();
+            var environment = new Mock<IEnvironment>();
+            var httpContext = new DefaultHttpContext();
+
+            // Setup conditions for specialization but NOT Flex
+            webHostEnvironment.Setup(e => e.InStandbyMode).Returns(false);
+            environment.Setup(e => e.IsContainerReady()).Returns(true);
+            environment.Setup(e => e.IsFlexConsumptionSku()).Returns(false);
+
+            standbyManager.Setup(s => s.SpecializeHostAsync()).Returns(Task.CompletedTask);
+
+            var nextCalled = false;
+            RequestDelegate next = (ctx) =>
+            {
+                nextCalled = true;
+                return Task.CompletedTask;
+            };
+
+            var middleware = new PlaceholderSpecializationMiddleware(next, webHostEnvironment.Object, standbyManager.Object, environment.Object);
+
+            // Act
+            await middleware.Invoke(httpContext);
+
+            // Assert
+            Assert.True(nextCalled);
+            Assert.False(httpContext.Request.Headers.ContainsKey(ScriptConstants.AntaresColdStartHeaderName));
+            standbyManager.Verify(s => s.SpecializeHostAsync(), Times.Once);
+        }
+
+        [Fact]
+        public async Task Invoke_NoSpecialization_DoesNotAddHeader()
+        {
+            // Arrange
+            var webHostEnvironment = new Mock<IScriptWebHostEnvironment>();
+            var standbyManager = new Mock<IStandbyManager>();
+            var environment = new Mock<IEnvironment>();
+            var httpContext = new DefaultHttpContext();
+
+            // Setup conditions for NO specialization
+            webHostEnvironment.Setup(e => e.InStandbyMode).Returns(true);
+            environment.Setup(e => e.IsContainerReady()).Returns(true);
+
+            var nextCalled = false;
+            RequestDelegate next = (ctx) =>
+            {
+                nextCalled = true;
+                return Task.CompletedTask;
+            };
+
+            var middleware = new PlaceholderSpecializationMiddleware(next, webHostEnvironment.Object, standbyManager.Object, environment.Object);
+
+            // Act
+            await middleware.Invoke(httpContext);
+
+            // Assert
+            Assert.True(nextCalled);
+            Assert.False(httpContext.Request.Headers.ContainsKey(ScriptConstants.AntaresColdStartHeaderName));
+            standbyManager.Verify(s => s.SpecializeHostAsync(), Times.Never);
+        }
+    }
+}