Add MCP probe to comprehensive health endpoint

src/Config/HealthCheck/HealthCheckConstants.cs

@@ -12,6 +12,7 @@ public static class HealthCheckConstants
         public const string DATASOURCE = "data-source";
         public const string REST = "rest";
         public const string GRAPHQL = "graphql";
+        public const string MCP = "mcp";
         public const string EMBEDDING = "embedding";
         public const int ERROR_RESPONSE_TIME_MS = -1;
         public const int DEFAULT_THRESHOLD_RESPONSE_TIME_MS = 1000;

src/Service.Tests/Configuration/HealthEndpointTests.cs

@@ -110,6 +110,7 @@ public async Task ComprehensiveHealthEndpoint_ValidateContents(
                     ValidateConfigurationDetailsHealthCheckResponse(responseProperties, enableGlobalRest, enableGlobalGraphql, enableGlobalMcp);
                     ValidateIfAttributePresentInResponse(responseProperties, enableDatasourceHealth, HealthCheckConstants.DATASOURCE);
                     ValidateIfAttributePresentInResponse(responseProperties, enableEntityHealth, HealthCheckConstants.ENDPOINT);
+                    ValidateIfAttributePresentInResponse(responseProperties, enableGlobalMcp, HealthCheckConstants.MCP);
                     if (enableEntityHealth)
                     {
                         ValidateEntityRestAndGraphQLResponse(responseProperties, enableEntityRest, enableEntityGraphQL, enableGlobalRest, enableGlobalGraphql);
@@ -223,6 +224,52 @@ public async Task TestFailureHealthCheckGraphQLResponseAsync()
             Assert.IsNotNull(errorMessageFromGraphQL);
         }
 
+        /// <summary>
+        /// Simulates the function call to HttpUtilities.ExecuteMcpQueryAsync.
+        /// while setting up mock HTTP client to simulate the response from the server to send OK code.
+        /// Validates the response to ensure no error message is received.
+        /// </summary>
+        [TestMethod]
+        public async Task TestHealthCheckMcpResponseAsync()
+        {
+            // Arrange
+            RuntimeConfig runtimeConfig = SetupCustomConfigFile(true, true, true, true, true, true, true, true);
+            HttpUtilities httpUtilities = SetupMcpTest(runtimeConfig);
+
+            // Act
+            // Simulate an MCP initialize POST request to the endpoint.
+            // Response should be null as error message is not expected to be returned.
+            string errorMessageFromMcp = await httpUtilities.ExecuteMcpQueryAsync(
+                mcpUriSuffix: runtimeConfig.McpPath,
+                incomingRoleHeader: string.Empty,
+                incomingRoleToken: string.Empty);
+
+            // Assert
+            Assert.IsNull(errorMessageFromMcp);
+        }
+
+        /// <summary>
+        /// Simulates the function call to HttpUtilities.ExecuteMcpQueryAsync.
+        /// while setting up mock HTTP client to simulate the response from the server to send InternalServerError code.
+        /// Validates the response to ensure error message is received.
+        /// </summary>
+        [TestMethod]
+        public async Task TestFailureHealthCheckMcpResponseAsync()
+        {
+            // Arrange
+            RuntimeConfig runtimeConfig = SetupCustomConfigFile(true, true, true, true, true, true, true, true);
+            HttpUtilities httpUtilities = SetupMcpTest(runtimeConfig, HttpStatusCode.InternalServerError);
+
+            // Act
+            string errorMessageFromMcp = await httpUtilities.ExecuteMcpQueryAsync(
+                mcpUriSuffix: runtimeConfig.McpPath,
+                incomingRoleHeader: string.Empty,
+                incomingRoleToken: string.Empty);
+
+            // Assert
+            Assert.IsNotNull(errorMessageFromMcp);
+        }
+
         /// <summary>
         /// Tests the serialization behavior of <see cref="RuntimeHealthCheckConfig"/> for the <see cref="RuntimeHealthCheckConfig.MaxQueryParallelism"/> property."
         /// </summary>
@@ -366,6 +413,51 @@ private static HttpUtilities SetupGraphQLTest(RuntimeConfig runtimeConfig, HttpS
                 mockHttpClientFactory.Object);
         }
 
+        private static HttpUtilities SetupMcpTest(RuntimeConfig runtimeConfig, HttpStatusCode httpStatusCode = HttpStatusCode.OK)
+        {
+            // Arrange
+            // Create a mock entity map with a single entity for testing and load in RuntimeConfigProvider
+            MockFileSystem fileSystem = new();
+            fileSystem.AddFile(FileSystemRuntimeConfigLoader.DEFAULT_CONFIG_FILE_NAME, new MockFileData(runtimeConfig.ToJson()));
+            FileSystemRuntimeConfigLoader loader = new(fileSystem);
+            RuntimeConfigProvider provider = new(loader);
+            Mock<IMetadataProviderFactory> metadataProviderFactory = new();
+            string expectedMcpPayload = Azure.DataApiBuilder.Service.HealthCheck.Utilities.CreateHttpMcpQuery();
+
+            // Mock the handler to return the supplied status code for the MCP initialize POST request.
+            Mock<HttpMessageHandler> mockHandler = new();
+            mockHandler.Protected()
+                .Setup<Task<HttpResponseMessage>>(
+                    "SendAsync",
+                    ItExpr.Is<HttpRequestMessage>(req =>
+                        req.Method == HttpMethod.Post &&
+                        req.RequestUri == new Uri($"{BASE_DAB_URL}{runtimeConfig.McpPath}") &&
+                        req.Content != null &&
+                        req.Content.ReadAsStringAsync().Result.Equals(expectedMcpPayload) &&
+                        req.Headers.Accept.Any(v => v.MediaType == "application/json") &&
+                        req.Headers.Accept.Any(v => v.MediaType == "text/event-stream")),
+                    ItExpr.IsAny<CancellationToken>())
+                .ReturnsAsync(new HttpResponseMessage(httpStatusCode)
+                {
+                    Content = new StringContent("{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{}}")
+                });
+
+            Mock<IHttpClientFactory> mockHttpClientFactory = new();
+            mockHttpClientFactory.Setup(x => x.CreateClient("ContextConfiguredHealthCheckClient"))
+                .Returns(new HttpClient(mockHandler.Object)
+                {
+                    BaseAddress = new Uri($"{BASE_DAB_URL}")
+                });
+
+            Mock<ILogger<HttpUtilities>> logger = new();
+
+            return new(
+                logger.Object,
+                metadataProviderFactory.Object,
+                provider,
+                mockHttpClientFactory.Object);
+        }
+
         private static void ValidateEntityRestAndGraphQLResponse(
             Dictionary<string, JsonElement> responseProperties,
             bool enableEntityRest,

src/Service/HealthCheck/HealthCheckHelper.cs

@@ -169,6 +169,7 @@ private async Task UpdateHealthCheckDetailsAsync(ComprehensiveHealthCheckReport
             comprehensiveHealthCheckReport.Checks = new List<HealthCheckResultEntry>();
             await UpdateDataSourceHealthCheckResultsAsync(comprehensiveHealthCheckReport, runtimeConfig);
             await UpdateEntityHealthCheckResultsAsync(comprehensiveHealthCheckReport, runtimeConfig, roleHeader, roleToken);
+            await UpdateMcpHealthCheckResultsAsync(comprehensiveHealthCheckReport, runtimeConfig, roleHeader, roleToken);
             await UpdateEmbeddingsHealthCheckResultsAsync(comprehensiveHealthCheckReport, runtimeConfig);
         }
 
@@ -213,6 +214,51 @@ private async Task UpdateDataSourceHealthCheckResultsAsync(ComprehensiveHealthCh
             return (HealthCheckConstants.ERROR_RESPONSE_TIME_MS, errorMessage);
         }
 
+        // Updates the MCP Health Check Result in the response.
+        // The check verifies that the MCP endpoint is reachable and responds within the threshold.
+        // It runs only when MCP is enabled in the runtime configuration.
+        private async Task UpdateMcpHealthCheckResultsAsync(ComprehensiveHealthCheckReport comprehensiveHealthCheckReport, RuntimeConfig runtimeConfig, string roleHeader, string roleToken)
+        {
+            if (comprehensiveHealthCheckReport.Checks is null || !runtimeConfig.IsMcpEnabled)
+            {
+                return;
+            }
+
+            (int, string?) response = await ExecuteMcpQueryAsync(runtimeConfig.McpPath, roleHeader, roleToken);
+            bool isResponseTimeWithinThreshold = response.Item1 >= 0 && response.Item1 < HealthCheckConstants.DEFAULT_THRESHOLD_RESPONSE_TIME_MS;
+
+            comprehensiveHealthCheckReport.Checks.Add(new HealthCheckResultEntry
+            {
+                Name = HealthCheckConstants.MCP,
+                ResponseTimeData = new ResponseTimeData
+                {
+                    ResponseTimeMs = response.Item1,
+                    ThresholdMs = HealthCheckConstants.DEFAULT_THRESHOLD_RESPONSE_TIME_MS
+                },
+                Tags = [HealthCheckConstants.MCP],
+                Exception = response.Item2 ?? (!isResponseTimeWithinThreshold ? TIME_EXCEEDED_ERROR_MESSAGE : null),
+                Status = isResponseTimeWithinThreshold ? HealthStatus.Healthy : HealthStatus.Unhealthy
+            });
+        }
+
+        // Executes the MCP Query and keeps track of the response time and error message.
+        private async Task<(int, string?)> ExecuteMcpQueryAsync(string mcpUriSuffix, string roleHeader, string roleToken)
+        {
+            if (string.IsNullOrEmpty(mcpUriSuffix))
+            {
+                return (HealthCheckConstants.ERROR_RESPONSE_TIME_MS, "MCP is enabled but no MCP path is configured.");
+            }
+
+            Stopwatch stopwatch = new();
+            stopwatch.Start();
+            string? errorMessage = await _httpUtility.ExecuteMcpQueryAsync(mcpUriSuffix, roleHeader, roleToken);
+            stopwatch.Stop();
+
+            return string.IsNullOrEmpty(errorMessage)
+                ? ((int)stopwatch.ElapsedMilliseconds, errorMessage)
+                : (HealthCheckConstants.ERROR_RESPONSE_TIME_MS, errorMessage);
+        }
+
         // Updates the Entity Health Check Results in the response.
         // Goes through the entities one by one and executes the rest and graphql checks (if enabled).
         // Stored procedures are excluded from health checks because they require parameters and are not guaranteed to be deterministic.

src/Service/HealthCheck/HttpUtilities.cs

@@ -136,6 +136,82 @@ public HttpUtilities(
             }
         }
 
+        // Executes the MCP query by sending an initialize JSON-RPC POST request to the MCP endpoint.
+        public async Task<string?> ExecuteMcpQueryAsync(string mcpUriSuffix, string incomingRoleHeader, string incomingRoleToken)
+        {
+            string? errorMessage = null;
+            try
+            {
+                if (string.IsNullOrWhiteSpace(mcpUriSuffix))
+                {
+                    const string msg = "MCP path is not configured.";
+                    _logger.LogError(msg);
+                    return msg;
+                }
+
+                if (_httpClient.BaseAddress is null)
+                {
+                    const string msg = "Health check HTTP client BaseAddress is not configured.";
+                    _logger.LogError(msg);
+                    return msg;
+                }
+
+                // Ensure the configured MCP path cannot override the host (e.g. absolute URIs).
+                if (!Uri.TryCreate(mcpUriSuffix, UriKind.Relative, out _))
+                {
+                    _logger.LogError("Blocked outbound request due to invalid or unsafe URI.");
+                    return "Blocked outbound request due to invalid or unsafe URI.";
+                }
+
+                Uri requestUri = new(_httpClient.BaseAddress, mcpUriSuffix);
+
+                if (!Program.CheckSanityOfUrl(requestUri.AbsoluteUri))
+                {
+                    _logger.LogError("Blocked outbound request due to invalid or unsafe URI.");
+                    return "Blocked outbound request due to invalid or unsafe URI.";
+                }
+
+                string jsonPayload = Utilities.CreateHttpMcpQuery();
+                HttpContent content = new StringContent(jsonPayload, Encoding.UTF8, Utilities.JSON_CONTENT_TYPE);
+
+                HttpRequestMessage message = new(method: HttpMethod.Post, requestUri: requestUri)
+                {
+                    Content = content
+                };
+
+                // The MCP Streamable HTTP transport requires the client to accept both
+                // JSON and SSE responses.
+                message.Headers.Add("Accept", "application/json, text/event-stream");
+
+                if (!string.IsNullOrEmpty(incomingRoleToken))
+                {
+                    message.Headers.Add(AuthenticationOptions.CLIENT_PRINCIPAL_HEADER, incomingRoleToken);
+                }
+
+                if (!string.IsNullOrEmpty(incomingRoleHeader))
+                {
+                    message.Headers.Add(AuthorizationResolver.CLIENT_ROLE_HEADER, incomingRoleHeader);
+                }
+
+                HttpResponseMessage response = await _httpClient.SendAsync(message);
+                if (response.IsSuccessStatusCode)
+                {
+                    _logger.LogTrace($"The MCP HealthEndpoint query executed successfully with code {response.StatusCode}.");
+                }
+                else
+                {
+                    errorMessage = $"The MCP HealthEndpoint query failed with code: {response.StatusCode}.";
+                }
+
+                return errorMessage;
+            }
+            catch (Exception ex)
+            {
+                _logger.LogError($"An exception occurred while executing the health check MCP query: {ex.Message}");
+                return ex.Message;
+            }
+        }
+
         // Executes the GraphQL query by sending a POST request to the API.
         // Internally calls the metadata provider to fetch the column names to create the graphql payload.
         public async Task<string?> ExecuteGraphQLQueryAsync(string graphqlUriSuffix, string entityName, Entity entity, string incomingRoleHeader, string incomingRoleToken)

src/Service/HealthCheck/Utilities.cs

@@ -6,6 +6,7 @@
 using System.Data.Common;
 using System.Text.Json;
 using Azure.DataApiBuilder.Config.ObjectModel;
+using Azure.DataApiBuilder.Mcp.Core;
 using Microsoft.Data.SqlClient;
 using Microsoft.Extensions.Logging;
 using MySqlConnector;
@@ -72,6 +73,28 @@ public static string CreateHttpRestQuery(string entityName, int first)
             return $"/{entityName}?$first={first}";
         }
 
+        public static string CreateHttpMcpQuery()
+        {
+            // Create a minimal MCP request (initialize) as a valid JSON-RPC request.
+            // 'initialize' is used because other methods (e.g. 'tools/list') require an active
+            // session in the MCP Streamable HTTP transport.
+            string protocolVersion = McpProtocolDefaults.ResolveProtocolVersion(configuration: null);
+            var payload = new
+            {
+                jsonrpc = "2.0",
+                id = 1,
+                method = "initialize",
+                @params = new
+                {
+                    protocolVersion,
+                    capabilities = new { },
+                    clientInfo = new { name = "dab-health-check", version = "1.0.0" }
+                }
+            };
+
+            return JsonSerializer.Serialize(payload);
+        }
+
         public static string NormalizeConnectionString(string connectionString, DatabaseType dbType, ILogger? logger = null)
         {
             try