[MCP] Respond to stdio initialize before metadata inference

src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs

@@ -6,6 +6,7 @@
 using Azure.DataApiBuilder.Config.ObjectModel;
 using Azure.DataApiBuilder.Core.AuthenticationHelpers.AuthenticationSimulator;
 using Azure.DataApiBuilder.Core.Configurations;
+using Azure.DataApiBuilder.Core.Services.MetadataProviders;
 using Azure.DataApiBuilder.Core.Telemetry;
 using Azure.DataApiBuilder.Mcp.Model;
 using Azure.DataApiBuilder.Mcp.Telemetry;
@@ -31,6 +32,8 @@ public class McpStdioServer : IMcpStdioServer
         private readonly IMcpStdioToolListChangedNotifier? _toolListChangedNotifier;
         private readonly TextReader? _inputReader;
         private readonly string _protocolVersion;
+        private readonly object _initializationLock = new();
+        private Task? _initializationTask;
 
         private const int MAX_LINE_LENGTH = 1024 * 1024; // 1 MB limit for incoming JSON-RPC requests
 
@@ -148,7 +151,7 @@ public async Task RunAsync(CancellationToken cancellationToken)
                                 break;
 
                             case "tools/list":
-                                HandleListTools(id);
+                                await HandleListToolsAsync(id, cancellationToken);
                                 break;
 
                             case "tools/call":
@@ -299,8 +302,10 @@ private void HandleInitialize(JsonElement? id, JsonElement root)
         /// <param name="id">
         /// The request identifier extracted from the incoming JSON-RPC request. Used to correlate the response with the request.
         /// </param>
-        private void HandleListTools(JsonElement? id)
+        private async Task HandleListToolsAsync(JsonElement? id, CancellationToken cancellationToken)
         {
+            await EnsureToolsInitializedAsync(cancellationToken);
+
             List<object> toolsWire = new();
 
             foreach (Tool tool in _toolRegistry.GetAdvertisedTools())
@@ -316,6 +321,27 @@ private void HandleListTools(JsonElement? id)
             WriteResult(id, new { tools = toolsWire });
         }
 
+        private Task EnsureToolsInitializedAsync(CancellationToken cancellationToken)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+
+            lock (_initializationLock)
+            {
+                return _initializationTask ??= InitializeToolsAsync(cancellationToken);
+            }
+        }
+
+        private async Task InitializeToolsAsync(CancellationToken cancellationToken)
+        {
+            IMetadataProviderFactory metadataProviderFactory =
+                _serviceProvider.GetRequiredService<IMetadataProviderFactory>();
+            await metadataProviderFactory.InitializeAsync(cancellationToken);
+
+            IMcpToolRegistryRefreshService? registryRefreshService =
+                _serviceProvider.GetService<IMcpToolRegistryRefreshService>();
+            registryRefreshService?.EnsureInitialized(cancellationToken);
+        }
+
         /// <summary>
         /// Handles the "logging/setLevel" JSON-RPC method by updating the runtime log level.
         /// </summary>
@@ -463,6 +489,8 @@ private async Task HandleCallToolAsync(JsonElement? id, JsonElement root, Cancel
                 return;
             }
 
+            await EnsureToolsInitializedAsync(ct);
+
             if (!_toolRegistry.TryGetTool(toolName!, out IMcpTool? tool) || tool is null)
             {
                 WriteError(id, McpStdioJsonRpcErrorCodes.INVALID_PARAMS, $"Tool not found: {toolName}");

src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs

@@ -64,6 +64,9 @@ public async Task InitializedClient_FileReload_EmitsOneNotificationAndReturnsUpd
                     .SetupGet(provider => provider.EntityToDatabaseObject)
                     .Returns(new Dictionary<string, DatabaseObject>());
                 Mock<IMetadataProviderFactory> metadataProviderFactory = new();
+                metadataProviderFactory
+                    .Setup(factory => factory.InitializeAsync())
+                    .Returns(Task.CompletedTask);
                 metadataProviderFactory
                     .Setup(factory => factory.GetMetadataProvider(It.IsAny<string>()))
                     .Returns(sqlMetadataProvider.Object);
@@ -73,12 +76,6 @@ public async Task InitializedClient_FileReload_EmitsOneNotificationAndReturnsUpd
                 ChannelTextWriter stdout = new();
                 using McpStdoutWriter stdoutWriter = new(stdout);
                 McpStdioToolListChangedNotifier notifier = new(stdoutWriter);
-                using ServiceProvider serviceProvider = new ServiceCollection()
-                    .AddSingleton(stdoutWriter)
-                    .AddSingleton<IMcpStdioToolListChangedNotifier>(notifier)
-                    .AddSingleton(configProvider.Object)
-                    .BuildServiceProvider();
-
                 McpToolRegistryRefreshService refreshService = new(
                     configProvider.Object,
                     Array.Empty<IMcpTool>(),
@@ -88,6 +85,13 @@ public async Task InitializedClient_FileReload_EmitsOneNotificationAndReturnsUpd
                     NullLogger<McpToolRegistryRefreshService>.Instance,
                     hotReloadEventHandler);
                 refreshService.EnsureInitialized();
+                using ServiceProvider serviceProvider = new ServiceCollection()
+                    .AddSingleton(stdoutWriter)
+                    .AddSingleton<IMcpStdioToolListChangedNotifier>(notifier)
+                    .AddSingleton(configProvider.Object)
+                    .AddSingleton<IMetadataProviderFactory>(metadataProviderFactory.Object)
+                    .AddSingleton<IMcpToolRegistryRefreshService>(refreshService)
+                    .BuildServiceProvider();
 
                 McpStdioServer server = new(registry, serviceProvider, stdin);
                 using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10));

src/Service.Tests/UnitTests/McpStdioHelperTests.cs

@@ -48,22 +48,18 @@ public void RunMcpStdioHost_DoesNotStartWebHost()
                 "MCP stdio mode should not stop a host that was never started.");
             Assert.AreEqual(1, stdioServer.RunAsyncCallCount,
                 "MCP stdio mode should still run the stdio JSON-RPC loop.");
-            Assert.AreEqual(1, refreshService.EnsureInitializedCallCount,
-                "MCP stdio mode should initialize the shared tool registry before running the loop.");
+            Assert.AreEqual(0, refreshService.EnsureInitializedCallCount,
+                "MCP stdio mode should defer shared tool registry initialization until the protocol loop needs tools.");
             CollectionAssert.AreEqual(
-                new[] { "metadata", "registry" },
+                Array.Empty<string>(),
                 metadataProviderFactory.InitializationOrder,
-                "MCP stdio mode should initialize metadata before publishing the registry.");
-            Assert.IsTrue(metadataProviderFactory.CancellationToken.CanBeCanceled,
-                "Metadata initialization must receive the loader's shutdown cancellation token.");
-            Assert.AreEqual(metadataProviderFactory.CancellationToken, refreshService.CancellationToken,
-                "Metadata initialization and registry publication must share the serialized operation's token.");
+                "MCP stdio mode should not infer metadata before the protocol loop starts.");
             Assert.AreEqual(lifetime.ApplicationStopping, stdioServer.CancellationToken,
                 "The stdio loop should keep using the host lifetime cancellation token.");
             Assert.AreEqual(1, host.DisposeCallCount,
                 "MCP stdio mode should dispose the host after the stdio loop exits.");
-            Assert.AreEqual(1, metadataProviderFactory.InitializeAsyncCallCount,
-                "MCP stdio mode must initialize metadata exactly once through the shared runtime initialization path.");
+            Assert.AreEqual(0, metadataProviderFactory.InitializeAsyncCallCount,
+                "MCP stdio mode must not initialize metadata before the protocol loop needs tools.");
             Assert.IsTrue(serviceProvider.GetRequiredService<FileSystemRuntimeConfigLoader>().ShutdownResourcesDisposed,
                 "MCP stdio shutdown must drain the loader before disposing the host.");
         }
@@ -112,21 +108,31 @@ public void RunMcpStdioHost_Fails_ReportsOnStandardErrorAndDisposesHost(bool fai
 
             string reported = capturedError.ToString();
 
-            Assert.IsFalse(result, "A host failure should be reported through the bool contract.");
-            Assert.AreEqual(failDuringStdio ? 1 : 0, stdioServer.RunAsyncCallCount,
-                "The stdio loop must run only when metadata initialization succeeds.");
+            Assert.AreEqual(!failDuringStdio, result,
+                "Only failures from the stdio loop should be reported by the host helper before lazy tool initialization.");
+            Assert.AreEqual(1, stdioServer.RunAsyncCallCount,
+                "The stdio loop must run even when metadata initialization would fail later.");
             TestMcpToolRegistryRefreshService refreshService =
                 (TestMcpToolRegistryRefreshService)serviceProvider.GetRequiredService<IMcpToolRegistryRefreshService>();
-            Assert.AreEqual(failDuringStdio ? 1 : 0, refreshService.EnsureInitializedCallCount,
-                "The registry must not publish tools after metadata initialization fails.");
+            Assert.AreEqual(0, refreshService.EnsureInitializedCallCount,
+                "The host helper must not publish tools before the stdio loop needs them.");
             Assert.AreEqual(1, host.DisposeCallCount,
                 "The host must still be disposed when initialization or the loop fails.");
             Assert.IsTrue(serviceProvider.GetRequiredService<FileSystemRuntimeConfigLoader>().ShutdownResourcesDisposed,
                 "Failure reporting must not bypass the loader's shutdown drain.");
-            StringAssert.Contains(reported, "MCP stdio host",
-                "The operator needs to know which host failed, not only that one did.");
-            StringAssert.Contains(reported, failure.Message,
-                "GetAwaiter().GetResult() rethrows the original exception, so the cause must survive.");
+            if (failDuringStdio)
+            {
+                StringAssert.Contains(reported, "MCP stdio host",
+                    "The operator needs to know which host failed, not only that one did.");
+                StringAssert.Contains(reported, failure.Message,
+                    "GetAwaiter().GetResult() rethrows the original exception, so the cause must survive.");
+            }
+            else
+            {
+                Assert.AreEqual(string.Empty, reported,
+                    "Lazy metadata failures must not be reported before a tool request starts initialization.");
+            }
+
             Assert.AreEqual(string.Empty, capturedOut.ToString(),
                 "stdout is the JSON-RPC channel; a stray byte on it corrupts the protocol.");
         }
@@ -138,9 +144,13 @@ public void RunMcpStdioHost_Fails_ReportsOnStandardErrorAndDisposesHost(bool fai
         /// real stream without installing a writer that would outlive the call.
         /// </summary>
         [TestMethod]
-        public void RunMcpStdioHost_StartupFails_WhenStandardErrorSuppressed_LeavesConsoleUnchanged()
+        public void RunMcpStdioHost_LoopFails_WhenStandardErrorSuppressed_LeavesConsoleUnchanged()
         {
-            TestMcpStdioServer stdioServer = new();
+            Exception failure = InferenceFailure();
+            TestMcpStdioServer stdioServer = new()
+            {
+                RunAsyncException = failure
+            };
             TestMetadataProviderFactory metadataProviderFactory = new()
             {
                 InitializeAsyncException = InferenceFailure()
@@ -168,8 +178,8 @@ public void RunMcpStdioHost_StartupFails_WhenStandardErrorSuppressed_LeavesConso
             Assert.IsFalse(result, "The bool contract holds whether or not stderr was suppressed.");
             Assert.IsTrue(consoleErrorUntouched,
                 "The report must not leave a replacement writer installed on Console.Error.");
-            Assert.AreEqual(0, stdioServer.RunAsyncCallCount,
-                "The stdio loop must not run after startup failed.");
+            Assert.AreEqual(1, stdioServer.RunAsyncCallCount,
+                "The stdio loop failure should be reported without changing Console.Error.");
         }
 
         [DataTestMethod]
@@ -192,14 +202,23 @@ public void RunMcpStdioHost_Canceled_PropagatesCancellationAndDrainsLoader(bool
                 BuildServices(stdioServer, metadataProviderFactory, out _);
             TestHost host = new(serviceProvider);
 
-            OperationCanceledException actual = Assert.ThrowsException<OperationCanceledException>(
-                () => McpStdioHelper.RunMcpStdioHost(host));
+            if (cancelDuringStdio)
+            {
+                OperationCanceledException actual = Assert.ThrowsException<OperationCanceledException>(
+                    () => McpStdioHelper.RunMcpStdioHost(host));
+
+                Assert.AreSame(failure, actual, "Cancellation must propagate to Program.StartEngine, not become a startup failure.");
+            }
+            else
+            {
+                Assert.IsTrue(McpStdioHelper.RunMcpStdioHost(host),
+                    "Metadata cancellation should be deferred until a tool request starts lazy initialization.");
+            }
 
-            Assert.AreSame(failure, actual, "Cancellation must propagate to Program.StartEngine, not become a startup failure.");
-            Assert.AreEqual(cancelDuringStdio ? 1 : 0, stdioServer.RunAsyncCallCount);
+            Assert.AreEqual(1, stdioServer.RunAsyncCallCount);
             TestMcpToolRegistryRefreshService refreshService =
                 (TestMcpToolRegistryRefreshService)serviceProvider.GetRequiredService<IMcpToolRegistryRefreshService>();
-            Assert.AreEqual(cancelDuringStdio ? 1 : 0, refreshService.EnsureInitializedCallCount);
+            Assert.AreEqual(0, refreshService.EnsureInitializedCallCount);
             Assert.AreEqual(1, host.DisposeCallCount);
             Assert.IsTrue(serviceProvider.GetRequiredService<FileSystemRuntimeConfigLoader>().ShutdownResourcesDisposed,
                 "Cancellation must still drain the loader before disposing the host.");

src/Service.Tests/UnitTests/McpStdioServerProtocolTests.cs

@@ -13,8 +13,11 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.DataApiBuilder.Config;
+using Azure.DataApiBuilder.Config.DatabasePrimitives;
 using Azure.DataApiBuilder.Config.ObjectModel;
 using Azure.DataApiBuilder.Core.Configurations;
+using Azure.DataApiBuilder.Core.Services;
+using Azure.DataApiBuilder.Core.Services.MetadataProviders;
 using Azure.DataApiBuilder.Core.Telemetry;
 using Azure.DataApiBuilder.Mcp.Core;
 using Azure.DataApiBuilder.Mcp.Model;
@@ -410,6 +413,7 @@ private static (McpStdioServer Server, StringWriter Output, IServiceProvider Ser
             services.AddSingleton(registry);
             services.AddSingleton(runtimeConfigProvider ?? new StubRuntimeConfigProvider(CreateRuntimeConfig()));
             services.AddSingleton<IConfiguration>(configuration ?? new ConfigurationBuilder().Build());
+            services.AddSingleton<IMetadataProviderFactory, NoOpMetadataProviderFactory>();
 
             if (logLevelController is not null)
             {
@@ -549,6 +553,26 @@ public Task<CallToolResult> ExecuteAsync(
             }
         }
 
+        private sealed class NoOpMetadataProviderFactory : IMetadataProviderFactory
+        {
+            public Task InitializeAsync() => Task.CompletedTask;
+
+            public void InitializeAsync(
+                Dictionary<string, Dictionary<string, DatabaseObject>> entityToDatabaseObjectMap,
+                Dictionary<string, Dictionary<string, string>> graphQLStoredProcedureExposedNameToEntityNameMap)
+            {
+            }
+
+            public ISqlMetadataProvider GetMetadataProvider(string dataSourceName)
+                => throw new NotImplementedException();
+
+            public IEnumerable<ISqlMetadataProvider> ListMetadataProviders()
+                => Array.Empty<ISqlMetadataProvider>();
+
+            public List<Exception> GetAllMetadataExceptions()
+                => new();
+        }
+
         private sealed class RecordingLogLevelController : ILogLevelController
         {
             private readonly bool _updateResult;