[HealthChecks] Add health check middleware

src/WebJobs.Script.WebHost/Controllers/KeysController.cs

@@ -69,7 +69,7 @@ public async Task<IActionResult> Get()
                 // Extensions that are webhook providers create their default system keys
                 // as part of host initialization (when those keys aren't already present).
                 // So we must delay key retrieval until host initialization is complete.
-                await _hostManager.DelayUntilHostReady();
+                await _hostManager.DelayUntilHostReadyAsync();
             }
 
             Dictionary<string, string> keys = await GetHostSecretsByScope(hostKeyScope);

src/WebJobs.Script.WebHost/Diagnostics/HealthChecks/HealthCheckResponseWriter.cs

@@ -0,0 +1,60 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using System.Threading.Tasks;
+using HealthChecks.UI.Client;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+using Microsoft.Extensions.Primitives;
+
+namespace Microsoft.Azure.WebJobs.Script.WebHost.Diagnostics.HealthChecks
+{
+    public class HealthCheckResponseWriter
+    {
+        private static readonly JsonSerializerOptions _options = CreateJsonOptions();
+
+        public static Task WriteResponseAsync(HttpContext httpContext, HealthReport report)
+        {
+            ArgumentNullException.ThrowIfNull(httpContext);
+            ArgumentNullException.ThrowIfNull(report);
+
+            // We will write a detailed report if ?expand=true is present.
+            if (httpContext.Request.Query.TryGetValue("expand", out StringValues value)
+                && bool.TryParse(value, out bool expand) && expand)
+            {
+                return UIResponseWriter.WriteHealthCheckUIResponse(httpContext, report);
+            }
+
+            return WriteMinimalResponseAsync(httpContext, report);
+        }
+
+        private static Task WriteMinimalResponseAsync(HttpContext httpContext, HealthReport report)
+        {
+            MinimalResponse body = new(report.Status);
+            return JsonSerializer.SerializeAsync(
+                httpContext.Response.Body, body, _options, httpContext.RequestAborted);
+        }
+
+        private static JsonSerializerOptions CreateJsonOptions()
+        {
+            var options = new JsonSerializerOptions
+            {
+                AllowTrailingCommas = true,
+                PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
+                DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull
+            };
+
+            options.Converters.Add(new JsonStringEnumConverter());
+
+            return options;
+        }
+
+        internal readonly struct MinimalResponse(HealthStatus status)
+        {
+            public HealthStatus Status { get; } = status;
+        }
+    }
+}

src/WebJobs.Script.WebHost/Diagnostics/HealthChecks/HealthCheckWaitMiddleware.cs

@@ -0,0 +1,40 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Azure.WebJobs.Script.WebHost.Models;
+using Microsoft.Extensions.Primitives;
+
+namespace Microsoft.Azure.WebJobs.Script.WebHost.Diagnostics.HealthChecks
+{
+    public class HealthCheckWaitMiddleware(RequestDelegate next, IScriptHostManager manager)
+    {
+        private readonly RequestDelegate _next = next ?? throw new ArgumentNullException(nameof(next));
+        private readonly IScriptHostManager _manager = manager ?? throw new ArgumentNullException(nameof(manager));
+
+        public async Task InvokeAsync(HttpContext context)
+        {
+            ArgumentNullException.ThrowIfNull(next);
+
+            // If specified, the ?wait={seconds} query param will wait for an
+            // active script host for that duration. This is to avoid excessive polling
+            // when waiting for the initial readiness probe.
+            if (context.Request.Query.TryGetValue("wait", out StringValues wait))
+            {
+                if (!int.TryParse(wait.ToString(), out int waitSeconds) || waitSeconds < 0)
+                {
+                    context.Response.StatusCode = StatusCodes.Status400BadRequest;
+                    await context.Response.WriteAsJsonAsync(
+                        ErrorResponse.BadArgument("'wait' query param must be a positive integer", $"wait={wait}"));
+                    return;
+                }
+
+                await _manager.DelayUntilHostReadyAsync(waitSeconds);
+            }
+
+            await _next(context);
+        }
+    }
+}

src/WebJobs.Script.WebHost/Extensions/HttpContextExtensions.cs

@@ -23,7 +23,7 @@ public static async Task WaitForRunningHostAsync(this HttpContext httpContext, I
                 // If the host is not ready, we'll wait a bit for it to initialize.
                 // This might happen if http requests come in while the host is starting
                 // up for the first time, or if it is restarting.
-                bool hostReady = await hostManager.DelayUntilHostReady(timeoutSeconds, pollingIntervalMilliseconds);
+                bool hostReady = await hostManager.DelayUntilHostReadyAsync(timeoutSeconds, pollingIntervalMilliseconds);
 
                 if (!hostReady)
                 {

src/WebJobs.Script.WebHost/Extensions/ScriptHostManagerExtensions.cs

@@ -7,7 +7,7 @@ namespace Microsoft.Azure.WebJobs.Script
 {
     public static class ScriptHostManagerExtensions
     {
-        public static async Task<bool> DelayUntilHostReady(this IScriptHostManager hostManager, int timeoutSeconds = ScriptConstants.HostTimeoutSeconds, int pollingIntervalMilliseconds = ScriptConstants.HostPollingIntervalMilliseconds)
+        public static async Task<bool> DelayUntilHostReadyAsync(this IScriptHostManager hostManager, int timeoutSeconds = ScriptConstants.HostTimeoutSeconds, int pollingIntervalMilliseconds = ScriptConstants.HostPollingIntervalMilliseconds)
         {
             if (HostIsInitialized(hostManager))
             {

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

@@ -54,7 +54,7 @@ private static async Task InvokeAwaitingHost(HttpContext context, RequestDelegat
             {
                 Logger.InitiatingHostAvailabilityCheck(logger);
 
-                bool hostReady = await scriptHostManager.DelayUntilHostReady();
+                bool hostReady = await scriptHostManager.DelayUntilHostReadyAsync();
                 if (!hostReady)
                 {
                     Logger.HostUnavailableAfterCheck(logger);

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

@@ -161,7 +161,7 @@ public async Task HostWarmupAsync(HttpRequest request)
                 await _hostManager.RestartHostAsync(CancellationToken.None);
 
                 // This call is here for sanity, but we should be fully initialized.
-                await _hostManager.DelayUntilHostReady();
+                await _hostManager.DelayUntilHostReadyAsync();
             }
         }
 

src/WebJobs.Script.WebHost/Models/ErrorResponse.cs

@@ -0,0 +1,63 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System.Collections.Generic;
+using System.Text.Json.Serialization;
+using Newtonsoft.Json;
+
+namespace Microsoft.Azure.WebJobs.Script.WebHost.Models
+{
+    /// <summary>
+    /// Represents an error response.
+    /// See https://github.com/Azure/azure-resource-manager-rpc/blob/master/v1.0/common-api-details.md#error-response-content.
+    /// </summary>
+    /// <param name="Code">
+    /// The error code. This is NOT the HTTP status code.
+    /// Unlocalized string which can be used to programmatically identify the error.
+    /// The code should be Pascal-cased, and should serve to uniquely identify a particular class of error,
+    /// for example "BadArgument".
+    /// </param>
+    /// <param name="Message">
+    /// The error message. Describes the error in detail and provides debugging information.
+    /// If Accept-Language is set in the request, it must be localized to that language.
+    /// </param>]
+    public record ErrorResponse(
+        [property: JsonProperty("code")][property: JsonPropertyName("code")] string Code,
+        [property: JsonProperty("message")][property: JsonPropertyName("message")] string Message)
+    {
+        /// <summary>
+        /// Gets the target of the particular error. For example, the name of the property in error.
+        /// </summary>
+        [JsonProperty("target")]
+        [JsonPropertyName("target")]
+        public string Target { get; init; }
+
+        /// <summary>
+        /// Gets the details of this error.
+        /// </summary>
+        [JsonProperty("details")]
+        [JsonPropertyName("details")]
+        public IEnumerable<ErrorResponse> Details { get; init; } = [];
+
+        /// <summary>
+        /// Gets the additional information for this error.
+        /// </summary>
+        [JsonProperty("additionalInfo")]
+        [JsonPropertyName("additionalInfo")]
+        public IEnumerable<ErrorAdditionalInfo> AdditionalInfo { get; init; } = [];
+
+        public static ErrorResponse BadArgument(string message, string target = null)
+        {
+            return new("BadArgument", message) { Target = target };
+        }
+    }
+
+    /// <summary>
+    /// Represents additional information for an error.
+    /// </summary>
+    /// <param name="Type">The type of additional information.</param>
+    /// <param name="Info">The additional error information.</param>
+    public record ErrorAdditionalInfo(
+        [property: JsonProperty("type")][property: JsonPropertyName("type")] string Type,
+        [property: JsonProperty("info")][property: JsonPropertyName("info")] object Info);
+}

src/WebJobs.Script.WebHost/Standby/StandbyManager.cs

@@ -124,7 +124,7 @@ public async Task SpecializeHostCoreAsync()
 
             using (_metricsLogger.LatencyEvent(MetricEventNames.SpecializationDelayUntilHostReady))
             {
-                await _scriptHostManager.DelayUntilHostReady();
+                await _scriptHostManager.DelayUntilHostReadyAsync();
             }
         }
 

src/WebJobs.Script.WebHost/WebJobs.Script.WebHost.csproj

@@ -50,6 +50,7 @@
   </ItemGroup>
 
   <ItemGroup>
+    <PackageReference Include="AspNetCore.HealthChecks.UI.Client" Version="9.0.0" />
     <PackageReference Include="Azure.Data.Tables" Version="12.8.3" />
     <PackageReference Include="Azure.Identity" Version="1.11.4" />
     <PackageReference Include="Azure.Security.KeyVault.Secrets" Version="4.6.0" />

src/WebJobs.Script/runtimeassemblies.json

@@ -71,6 +71,14 @@
       "name": "Grpc.Net.Common",
       "resolutionPolicy": "private"
     },
+    {
+      "name": "HealthChecks.UI.Client",
+      "resolutionPolicy": "private"
+    },
+    {
+      "name": "HealthChecks.UI.Core",
+      "resolutionPolicy": "private"
+    },
     {
       "name": "Microsoft.AI.DependencyCollector",
       "resolutionPolicy": "private"
@@ -607,6 +615,10 @@
       "name": "Microsoft.Extensions.Caching.Memory",
       "resolutionPolicy": "minorMatchOrLower"
     },
+    {
+      "name": "Microsoft.Extensions.Compliance.Abstractions",
+      "resolutionPolicy": "minorMatchOrLower"
+    },
     {
       "name": "Microsoft.Extensions.Configuration",
       "resolutionPolicy": "minorMatchOrLower"
@@ -779,6 +791,10 @@
       "name": "Microsoft.Extensions.Primitives",
       "resolutionPolicy": "minorMatchOrLower"
     },
+    {
+      "name": "Microsoft.Extensions.Telemetry.Abstractions",
+      "resolutionPolicy": "minorMatchOrLower"
+    },
     {
       "name": "Microsoft.Extensions.WebEncoders",
       "resolutionPolicy": "minorMatchOrLower"

test/WebJobs.Script.Tests.Integration/Controllers/ControllerScenarioTestFixture.cs

@@ -76,7 +76,7 @@ public virtual async Task InitializeAsync()
             HttpClient.BaseAddress = new Uri("https://localhost/");
 
             var manager = HttpServer.Host.Services.GetService<IScriptHostManager>();
-            await manager.DelayUntilHostReady();
+            await manager.DelayUntilHostReadyAsync();
         }
 
         public Task DisposeAsync()

test/WebJobs.Script.Tests.Shared/TestHelpers.cs

@@ -35,6 +35,28 @@ public static partial class TestHelpers
         private const string Chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
         private static readonly Random Random = new Random();
 
+        /// <summary>
+        /// Helper method to inline an action delegate.
+        /// </summary>
+        /// <param name="act">The action.</param>
+        /// <returns>The provided action.</returns>
+        /// <remarks>
+        /// This is intended to be used with a fluent assertion.
+        /// <c>Act(() => { }).Should().Something();</c>.
+        /// </remarks>
+        public static Action Act(Action act) => act;
+
+        /// <summary>
+        /// Helper method to inline an action delegate.
+        /// </summary>
+        /// <param name="act">The action.</param>
+        /// <returns>The provided action.</returns>
+        /// <remarks>
+        /// This is intended to be used with a fluent assertion.
+        /// <c>Act(() => { }).Should().Something();</c>.
+        /// </remarks>
+        public static Func<T> Act<T>(Func<T> act) => act;
+
         public static Task WaitOneAsync(this WaitHandle waitHandle)
         {
             ArgumentNullException.ThrowIfNull(waitHandle);