diff --git a/.gitignore b/.gitignore index 4de9ba1a81..267f3af4bf 100644 --- a/.gitignore +++ b/.gitignore @@ -36,6 +36,7 @@ dab-config*.json # Local-Only files .env +/docs/design/McpToolRegistryHotReload.md # Verify test files *.received.* \ No newline at end of file diff --git a/src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs b/src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs index f688eeb80a..2d23c6585c 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs @@ -2,7 +2,6 @@ // Licensed under the MIT License. using Azure.DataApiBuilder.Config.ObjectModel; -using Azure.DataApiBuilder.Mcp.Model; using Microsoft.Extensions.Logging; namespace Azure.DataApiBuilder.Mcp.Core @@ -18,16 +17,16 @@ public class CustomMcpToolFactory /// /// The runtime configuration containing entity definitions. /// Optional logger for diagnostic information. - /// Enumerable of custom tools generated from configuration. - public static IEnumerable CreateCustomTools(RuntimeConfig config, ILogger? logger = null) + /// Enumerable of dynamic custom tools generated from configuration. + public static IEnumerable CreateCustomTools(RuntimeConfig config, ILogger? logger = null) { if (config.Entities == null) { logger?.LogWarning("No entities found in runtime configuration for custom tool generation."); - return Enumerable.Empty(); + return Enumerable.Empty(); } - List customTools = new(); + List customTools = new(); foreach ((string entityName, Entity entity) in config.Entities) { @@ -48,10 +47,11 @@ public static IEnumerable CreateCustomTools(RuntimeConfig config, ILog } catch (Exception ex) { - logger?.LogError( - ex, - "Failed to create custom tool for entity '{EntityName}'. Skipping.", - entityName); + // Preserve entity context without logging here. The caller owns failure + // logging and can include whether startup failed or a snapshot was retained. + throw new InvalidOperationException( + $"Failed to create custom MCP tool for entity '{entityName}'.", + ex); } } } diff --git a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs index 573e6ea6a9..1d78e81680 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs @@ -12,6 +12,7 @@ using Azure.DataApiBuilder.Core.Resolvers; using Azure.DataApiBuilder.Core.Resolvers.Factories; using Azure.DataApiBuilder.Core.Services; +using Azure.DataApiBuilder.Core.Services.MetadataProviders; using Azure.DataApiBuilder.Mcp.Model; using Azure.DataApiBuilder.Mcp.Utils; using Azure.DataApiBuilder.Service.Exceptions; @@ -28,13 +29,8 @@ namespace Azure.DataApiBuilder.Mcp.Core /// /// Dynamic custom MCP tool generated from stored procedure entity configuration. /// Each custom tool represents a single stored procedure exposed as a dedicated MCP tool. - /// - /// Note: The entity configuration is captured at tool construction time. If the RuntimeConfig - /// is hot-reloaded, GetToolMetadata() will return cached metadata (name, description, parameters) - /// from the original configuration. This is acceptable because: - /// 1. MCP clients typically call tools/list once at startup - /// 2. ExecuteAsync always validates against the current runtime configuration - /// 3. Cached metadata improves performance for repeated metadata requests + /// A new instance is created for each MCP registry generation so its cached metadata remains + /// aligned with the runtime configuration used to advertise it. /// public class DynamicCustomTool : IMcpTool { @@ -50,6 +46,7 @@ public DynamicCustomTool(string entityName, Entity entity) { EntityName = entityName ?? throw new ArgumentNullException(nameof(entityName)); _entity = entity ?? throw new ArgumentNullException(nameof(entity)); + ToolName = ConvertToToolName(entityName); // Validate that this is a stored procedure if (_entity.Source.Type != EntitySourceType.StoredProcedure) @@ -65,6 +62,12 @@ public DynamicCustomTool(string entityName, Entity entity) /// public ToolType ToolType { get; } = ToolType.Custom; + /// + /// Returns true because creates an instance only when + /// the source entity has mcp.custom-tool enabled for the candidate configuration. + /// Each registry generation recreates that membership, so an extant dynamic tool is + /// enabled by construction. Execution still revalidates enablement against current state. + /// public bool IsEnabled(RuntimeConfig config) => true; /// @@ -73,16 +76,38 @@ public DynamicCustomTool(string entityName, Entity entity) public string EntityName { get; } /// - /// Initializes the tool's input schema using DB metadata from the service provider. - /// Called after DI initialization to enrich the tool schema with DB-discovered parameters - /// and type information that aren't available at construction time. - /// Falls back silently to config-based schema if DB metadata is unavailable. + /// Gets the normalized MCP tool name without materializing the complete metadata schema. + /// + internal string ToolName { get; } + + /// + /// Initializes the input schema using an explicit configuration and metadata-provider + /// generation. Falls back to config-based metadata when database metadata is unavailable. + /// + public bool InitializeMetadata( + RuntimeConfig config, + IMetadataProviderFactory metadataProviderFactory) + { + return InitializeMetadata(config, metadataProviderFactory, out _); + } + + /// + /// Initializes the input schema using an explicit configuration and metadata-provider + /// generation and reports why configuration metadata was used when database enrichment + /// is unavailable. /// - /// The application service provider with initialized metadata providers. - public void InitializeMetadata(IServiceProvider serviceProvider) + public bool InitializeMetadata( + RuntimeConfig config, + IMetadataProviderFactory metadataProviderFactory, + out string fallbackReason) { - ArgumentNullException.ThrowIfNull(serviceProvider); - _cachedInputSchema = BuildInputSchemaFromDbMetadata(serviceProvider); + ArgumentNullException.ThrowIfNull(config); + ArgumentNullException.ThrowIfNull(metadataProviderFactory); + _cachedInputSchema = BuildInputSchemaFromDbMetadata( + config, + metadataProviderFactory, + out fallbackReason); + return _cachedInputSchema.HasValue; } /// @@ -90,15 +115,14 @@ public void InitializeMetadata(IServiceProvider serviceProvider) /// public Tool GetToolMetadata() { - string toolName = ConvertToToolName(EntityName); - string description = _entity.Description ?? $"Executes the {toolName} stored procedure"; + string description = _entity.Description ?? $"Executes the {ToolName} stored procedure"; // Build input schema based on parameters JsonElement inputSchema = BuildInputSchema(); return new Tool { - Name = toolName, + Name = ToolName, Description = description, InputSchema = inputSchema }; @@ -113,7 +137,7 @@ public async Task ExecuteAsync( CancellationToken cancellationToken = default) { ILogger? logger = serviceProvider.GetService>(); - string toolName = GetToolMetadata().Name; + string toolName = ToolName; try { @@ -259,6 +283,10 @@ public async Task ExecuteAsync( cancellationToken.ThrowIfCancellationRequested(); queryResult = await queryEngine.ExecuteAsync(context, dataSourceName).ConfigureAwait(false); } + catch (OperationCanceledException) + { + throw; + } catch (DataApiBuilderException dabEx) { logger?.LogError(dabEx, "Error executing custom tool {ToolName} for entity {Entity}", toolName, EntityName); @@ -322,33 +350,32 @@ private JsonElement BuildInputSchema() /// Builds the input schema from DB metadata (StoredProcedureDefinition.Parameters). /// Returns null if metadata cannot be resolved (caller should fall back to config-based schema). /// - private JsonElement? BuildInputSchemaFromDbMetadata(IServiceProvider serviceProvider) + private JsonElement? BuildInputSchemaFromDbMetadata( + RuntimeConfig config, + IMetadataProviderFactory metadataProviderFactory, + out string fallbackReason) { - RuntimeConfigProvider? configProvider = serviceProvider.GetService(); - if (configProvider is null) - { - return null; - } - - RuntimeConfig config = configProvider.GetConfig(); - if (!McpMetadataHelper.TryResolveMetadata( EntityName, config, - serviceProvider, + metadataProviderFactory, out _, out DatabaseObject dbObject, out _, - out _)) + out fallbackReason)) { return null; } if (dbObject is not DatabaseStoredProcedure storedProcedure) { + fallbackReason = + $"Database object '{dbObject.FullName}' for entity '{EntityName}' is not a stored procedure."; return null; } + fallbackReason = string.Empty; + StoredProcedureDefinition spDefinition = storedProcedure.StoredProcedureDefinition; if (spDefinition.Parameters is null || spDefinition.Parameters.Count == 0) { diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs index 20040588fe..f40497a0bb 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs @@ -2,8 +2,6 @@ // Licensed under the MIT License. using System.Text.Json; -using Azure.DataApiBuilder.Config.ObjectModel; -using Azure.DataApiBuilder.Core.Configurations; using Azure.DataApiBuilder.Mcp.Model; using Azure.DataApiBuilder.Mcp.Utils; using Microsoft.Extensions.DependencyInjection; @@ -32,13 +30,9 @@ internal static IServiceCollection ConfigureMcpServer(this IServiceCollection se throw new InvalidOperationException("Tool registry is not available."); } - RuntimeConfigProvider runtimeConfigProvider = request.Services!.GetRequiredService(); - RuntimeConfig runtimeConfig = runtimeConfigProvider.GetConfig(); - List tools = toolRegistry.GetEnabledTools(runtimeConfig).ToList(); - return ValueTask.FromResult(new ListToolsResult { - Tools = tools + Tools = toolRegistry.GetAdvertisedTools().ToList() }); }) .WithCallToolHandler(async (RequestContext request, CancellationToken ct) => @@ -97,6 +91,10 @@ internal static IServiceCollection ConfigureMcpServer(this IServiceCollection se options.ServerInfo = new() { Name = McpProtocolDefaults.MCP_SERVER_NAME, Version = McpProtocolDefaults.MCP_SERVER_VERSION }; options.Capabilities ??= new(); options.Capabilities.Tools ??= new(); + // WithListToolsHandler enables tool discovery, but HTTP session broadcast is not + // implemented. Do not promise list-change notifications to HTTP clients. Stdio + // advertises and implements this capability in its separate initialize handler. + options.Capabilities.Tools.ListChanged = false; options.ServerInstructions = !string.IsNullOrWhiteSpace(instructions) ? instructions : null; }); diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpServiceCollectionExtensions.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpServiceCollectionExtensions.cs index c88cae148d..e22490d110 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpServiceCollectionExtensions.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpServiceCollectionExtensions.cs @@ -6,6 +6,7 @@ using Azure.DataApiBuilder.Core.Configurations; using Azure.DataApiBuilder.Mcp.Model; using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Hosting; namespace Azure.DataApiBuilder.Mcp.Core { @@ -33,14 +34,17 @@ public static IServiceCollection AddDabMcpServer(this IServiceCollection service // Register core MCP services services.AddSingleton(); - services.AddHostedService(); + services.AddSingleton(); + services.AddSingleton(serviceProvider => + serviceProvider.GetRequiredService()); + services.AddSingleton(serviceProvider => + serviceProvider.GetRequiredService()); - // Auto-discover and register all MCP tools + // Auto-discover MCP tool implementations from this assembly. Configuration-generated + // DynamicCustomTool objects are created separately by McpToolRegistryRefreshService; + // independently registered IMcpTool extensions remain in DI across generations. RegisterAllMcpTools(services); - // Register custom tools from configuration - RegisterCustomTools(services, runtimeConfig); - // Configure MCP server and propagate runtime description to MCP initialize instructions. services.ConfigureMcpServer(runtimeConfig.Runtime?.Mcp?.Description); @@ -66,16 +70,5 @@ private static void RegisterAllMcpTools(IServiceCollection services) } } - /// - /// Registers custom MCP tools generated from stored procedure entity configurations. - /// - private static void RegisterCustomTools(IServiceCollection services, RuntimeConfig config) - { - // Create custom tools and register each as a singleton - foreach (IMcpTool customTool in CustomMcpToolFactory.CreateCustomTools(config)) - { - services.AddSingleton(customTool); - } - } } } diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs index aef3b63dcc..0f25f3b253 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs @@ -28,6 +28,7 @@ public class McpStdioServer : IMcpStdioServer private readonly McpToolRegistry _toolRegistry; private readonly IServiceProvider _serviceProvider; private readonly McpStdoutWriter _stdoutWriter; + private readonly IMcpStdioToolListChangedNotifier? _toolListChangedNotifier; private readonly TextReader? _inputReader; private readonly string _protocolVersion; @@ -50,6 +51,7 @@ public McpStdioServer(McpToolRegistry toolRegistry, IServiceProvider serviceProv // notifications/message frames are serialized through one lock. // Falls back to a fresh instance if DI didn't register one (defensive). _stdoutWriter = _serviceProvider.GetService() ?? new McpStdoutWriter(); + _toolListChangedNotifier = _serviceProvider.GetService(); // Allow protocol version to be configured via IConfiguration, using centralized defaults. IConfiguration? configuration = _serviceProvider.GetService(); @@ -66,6 +68,7 @@ public async Task RunAsync(CancellationToken cancellationToken) // By default read via Console.In so the loop honors the configured // Console.InputEncoding in stdio mode. TextReader reader = _inputReader ?? Console.In; + bool initializeResponseWritten = false; while (!cancellationToken.IsCancellationRequested) { @@ -128,9 +131,20 @@ public async Task RunAsync(CancellationToken cancellationToken) { case "initialize": HandleInitialize(id, root); + // This assignment is reached only after WriteResult succeeds. + initializeResponseWritten = true; break; case "notifications/initialized": + // This notification completes the MCP handshake only after the + // server successfully wrote its initialize response. Ignore an + // out-of-order notification rather than enabling capabilities the + // client has not negotiated. + if (initializeResponseWritten) + { + _toolListChangedNotifier?.MarkInitialized(); + } + break; case "tools/list": @@ -183,6 +197,7 @@ private void HandleInitialize(JsonElement? id, JsonElement root) string? clientRequestedProtocolVersion = GetClientProtocolVersion(root); string negotiatedProtocolVersion = McpProtocolDefaults.ResolveInitializeResponseProtocolVersion(_protocolVersion, clientRequestedProtocolVersion); + bool supportsToolListChanged = _toolListChangedNotifier is not null; // Get the description from runtime config if available string? description = null; @@ -212,7 +227,7 @@ private void HandleInitialize(JsonElement? id, JsonElement root) protocolVersion = negotiatedProtocolVersion, capabilities = new { - tools = new { listChanged = true }, + tools = new { listChanged = supportsToolListChanged }, logging = new { } }, serverInfo = new @@ -230,7 +245,7 @@ private void HandleInitialize(JsonElement? id, JsonElement root) protocolVersion = negotiatedProtocolVersion, capabilities = new { - tools = new { listChanged = true }, + tools = new { listChanged = supportsToolListChanged }, logging = new { } }, serverInfo = new @@ -248,7 +263,7 @@ private void HandleInitialize(JsonElement? id, JsonElement root) protocolVersion = negotiatedProtocolVersion, capabilities = new { - tools = new { listChanged = true }, + tools = new { listChanged = supportsToolListChanged }, logging = new { } }, serverInfo = new @@ -287,16 +302,9 @@ private void HandleInitialize(JsonElement? id, JsonElement root) private void HandleListTools(JsonElement? id) { List toolsWire = new(); - int count = 0; - - // Resolve runtime config to filter out disabled tools. - RuntimeConfigProvider runtimeConfigProvider = _serviceProvider.GetRequiredService(); - RuntimeConfig runtimeConfig = runtimeConfigProvider.GetConfig(); - IEnumerable tools = _toolRegistry.GetEnabledTools(runtimeConfig); - foreach (Tool tool in tools) + foreach (Tool tool in _toolRegistry.GetAdvertisedTools()) { - count++; toolsWire.Add(new { name = tool.Name, diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs new file mode 100644 index 0000000000..dc71ed3f02 --- /dev/null +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs @@ -0,0 +1,164 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Text.Json; +using Azure.DataApiBuilder.Mcp.Model; +using Microsoft.Extensions.Logging; +using Microsoft.Extensions.Logging.Abstractions; +using ModelContextProtocol.Protocol; + +namespace Azure.DataApiBuilder.Mcp.Core +{ + /// + /// Stdio-specific lifecycle contract used by the JSON-RPC server to mark the client ready for + /// unsolicited tool-list change notifications. + /// + public interface IMcpStdioToolListChangedNotifier : IMcpToolListChangedNotifier + { + /// + /// Marks the MCP initialization handshake complete. + /// + void MarkInitialized(); + } + + /// + /// Writes MCP notifications/tools/list_changed frames for an initialized stdio client. + /// + public sealed class McpStdioToolListChangedNotifier : IMcpStdioToolListChangedNotifier + { + private static readonly string _notificationJson = JsonSerializer.Serialize(new + { + jsonrpc = McpStdioJsonRpcErrorCodes.JSON_RPC_VERSION, + method = NotificationMethods.ToolListChangedNotification, + @params = new { } + }); + + private readonly McpStdoutWriter _stdoutWriter; + private readonly ILogger _logger; + private readonly Func _tryScheduleWorker; + private int _isInitialized; + private int _notificationPending; + private int _notificationWorkerScheduled; + + public McpStdioToolListChangedNotifier( + McpStdoutWriter stdoutWriter, + ILogger? logger = null) + : this(stdoutWriter, logger, TryScheduleOnThreadPool) + { + } + + internal McpStdioToolListChangedNotifier( + McpStdoutWriter stdoutWriter, + ILogger? logger, + Func tryScheduleWorker) + { + _stdoutWriter = stdoutWriter ?? throw new ArgumentNullException(nameof(stdoutWriter)); + _logger = logger ?? NullLogger.Instance; + _tryScheduleWorker = tryScheduleWorker ?? + throw new ArgumentNullException(nameof(tryScheduleWorker)); + } + + /// + public void MarkInitialized() + { + Interlocked.Exchange(ref _isInitialized, 1); + } + + /// + public void NotifyToolsListChanged() + { + if (Volatile.Read(ref _isInitialized) == 0) + { + return; + } + + // One pending invalidation is sufficient: after receiving it, the client requests the + // latest complete snapshot. This keeps queued state bounded while stdout is blocked. + Interlocked.Exchange(ref _notificationPending, 1); + ScheduleNotificationWorker(); + } + + private void ScheduleNotificationWorker() + { + if (Interlocked.CompareExchange(ref _notificationWorkerScheduled, 1, 0) != 0) + { + return; + } + + Action worker = ProcessPendingNotifications; + if (!_tryScheduleWorker(worker)) + { + // Do not clear _notificationPending: this invalidation is still required even if + // no later configuration change occurs. A dedicated background thread is a rare + // fallback for ThreadPool queue rejection and preserves the nonblocking contract. + _logger.LogWarning( + "Failed to queue an MCP tool-list change notification on the thread pool. " + + "Starting a dedicated fallback worker."); + StartDedicatedFallbackWorker(worker); + } + } + + private static bool TryScheduleOnThreadPool(Action worker) + { + return ThreadPool.QueueUserWorkItem( + static callback => callback(), + worker, + preferLocal: false); + } + + private void StartDedicatedFallbackWorker(Action worker) + { + try + { + Thread fallbackWorker = new( + static callback => ((Action)callback!).Invoke()) + { + IsBackground = true, + Name = "DAB MCP tool-list notification fallback" + }; + fallbackWorker.Start(worker); + } + catch (Exception ex) + { + // Retain the pending flag and reopen the scheduling gate. A later notification can + // retry delivery if the process could not create the fallback thread. + Volatile.Write(ref _notificationWorkerScheduled, 0); + _logger.LogError( + ex, + "Failed to start the MCP tool-list notification fallback worker. " + + "The notification remains pending."); + } + } + + private void ProcessPendingNotifications() + { + try + { + while (Interlocked.Exchange(ref _notificationPending, 0) != 0) + { + try + { + _stdoutWriter.WriteLine(_notificationJson); + } + catch (Exception ex) + { + _logger.LogError( + ex, + "Failed to write an MCP tool-list change notification."); + } + } + } + finally + { + Volatile.Write(ref _notificationWorkerScheduled, 0); + + // A publication can race with worker shutdown after the final pending-flag + // exchange. Reschedule so that invalidation is never lost in that window. + if (Volatile.Read(ref _notificationPending) != 0) + { + ScheduleNotificationWorker(); + } + } + } + } +} diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpStdoutWriter.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpStdoutWriter.cs index 7a30fccec3..0bd4aee6aa 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpStdoutWriter.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpStdoutWriter.cs @@ -24,7 +24,7 @@ public sealed class McpStdoutWriter : IDisposable { private readonly object _lock = new(); private TextWriter? _writer; - private bool _disposed; + private int _disposed; /// /// Production constructor. The underlying stdout stream is opened @@ -51,9 +51,14 @@ internal McpStdoutWriter(TextWriter writer) /// public void WriteLine(string line) { + if (Volatile.Read(ref _disposed) != 0) + { + return; + } + lock (_lock) { - if (_disposed) + if (Volatile.Read(ref _disposed) != 0) { return; } @@ -65,17 +70,30 @@ public void WriteLine(string line) public void Dispose() { - lock (_lock) + if (Interlocked.Exchange(ref _disposed, 1) != 0) { - if (_disposed) - { - return; - } + return; + } + + // Stdout can block indefinitely when the MCP client stops reading. Never make host + // disposal wait behind an in-flight notification write. Marking this instance as + // disposed prevents new writes; if the writer lock is available, release its + // resources immediately. Otherwise the process-owned stdout stream is left for the + // operating system to reclaim when the blocked process exits. + if (!Monitor.TryEnter(_lock)) + { + return; + } - _disposed = true; + try + { _writer?.Dispose(); _writer = null; } + finally + { + Monitor.Exit(_lock); + } } private void EnsureInitialized() diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs index f8275c61d4..9ad2c40dca 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs @@ -1,7 +1,11 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +using System.Collections.Immutable; using System.Net; +using System.Text; +using System.Text.Json; +using System.Text.Json.Serialization; using Azure.DataApiBuilder.Config.ObjectModel; using Azure.DataApiBuilder.Mcp.Model; using Azure.DataApiBuilder.Service.Exceptions; @@ -15,18 +19,134 @@ namespace Azure.DataApiBuilder.Mcp.Core /// public class McpToolRegistry { - private readonly Dictionary _tools = new(StringComparer.OrdinalIgnoreCase); + private static readonly JsonSerializerOptions _discoveryJsonOptions = new() + { + DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull + }; + + private readonly object _writerLock = new(); + private McpToolRegistrySnapshot _snapshot = McpToolRegistrySnapshot.Empty; + + /// + /// Replaces the complete registry with a snapshot built for . + /// The candidate is validated and materialized before it is atomically published. + /// + internal McpToolRegistryUpdateResult ReplaceAll(IEnumerable tools, RuntimeConfig config) + { + return PublishCandidate(CreateCandidate(tools, config, CancellationToken.None)); + } + + /// + /// Builds and validates a complete replacement without publishing it. + /// + internal static McpToolRegistryCandidate CreateCandidate( + IEnumerable tools, + RuntimeConfig config, + CancellationToken cancellationToken) + { + ArgumentNullException.ThrowIfNull(tools); + ArgumentNullException.ThrowIfNull(config); + cancellationToken.ThrowIfCancellationRequested(); + + ImmutableDictionary.Builder toolBuilder = + ImmutableDictionary.CreateBuilder(StringComparer.OrdinalIgnoreCase); + List advertisedMetadata = new(); + + foreach (IMcpTool tool in tools) + { + cancellationToken.ThrowIfCancellationRequested(); + ArgumentNullException.ThrowIfNull(tool); + + Tool metadata = CloneMetadata(tool.GetToolMetadata()); + string toolName = ValidateToolName(metadata); + + if (toolBuilder.TryGetValue(toolName, out IMcpTool? existingTool)) + { + if (ReferenceEquals(existingTool, tool)) + { + continue; + } + + throw CreateDuplicateToolException(toolName, existingTool, tool); + } + + toolBuilder.Add(toolName, tool); + if (tool.IsEnabled(config)) + { + advertisedMetadata.Add(metadata); + } + } + + ImmutableArray advertisedTools = SortMetadata(advertisedMetadata); + string discoveryJson = CreateDiscoveryJson(advertisedTools); + return new McpToolRegistryCandidate( + Tools: toolBuilder.ToImmutable(), + AdvertisedToolCount: advertisedTools.Length, + DiscoveryJson: discoveryJson, + DiscoveryCanonicalJson: CreateDiscoveryCanonicalJson(discoveryJson)); + } /// - /// Registers a tool in the registry + /// Atomically publishes a previously built and validated candidate. /// - /// Thrown when tool name is invalid or duplicate - public void RegisterTool(IMcpTool tool) + internal McpToolRegistryUpdateResult PublishCandidate(McpToolRegistryCandidate candidate) { - Tool metadata = tool.GetToolMetadata(); - string toolName = metadata.Name?.Trim() ?? string.Empty; + ArgumentNullException.ThrowIfNull(candidate); - // Reject empty or whitespace-only tool names + lock (_writerLock) + { + McpToolRegistrySnapshot current = _snapshot; + McpToolRegistrySnapshot replacement = new( + Version: current.Version + 1, + Tools: candidate.Tools, + AdvertisedToolCount: candidate.AdvertisedToolCount, + DiscoveryJson: candidate.DiscoveryJson, + DiscoveryCanonicalJson: candidate.DiscoveryCanonicalJson); + + Interlocked.Exchange(ref _snapshot, replacement); + + return new McpToolRegistryUpdateResult( + Version: replacement.Version, + DiscoveryChanged: !string.Equals( + current.DiscoveryCanonicalJson, + replacement.DiscoveryCanonicalJson, + StringComparison.Ordinal), + RegisteredToolCount: replacement.Tools.Count, + AdvertisedToolCount: replacement.AdvertisedToolCount); + } + } + + /// + /// Gets the metadata snapshot advertised by tools/list. + /// + /// + /// Returns defensive deep clones so callers cannot mutate the private snapshot shared by + /// concurrent readers. Candidate construction serializes an order-preserving discovery + /// representation for serving and a separate canonical representation for semantic change + /// comparison. Discovery deserializes the serving representation instead of serializing + /// every tool again on each request, while still allocating caller-owned objects. + /// + public IReadOnlyList GetAdvertisedTools() + { + McpToolRegistrySnapshot snapshot = Volatile.Read(ref _snapshot); + return JsonSerializer.Deserialize( + snapshot.DiscoveryJson, + _discoveryJsonOptions) + ?? throw new InvalidOperationException( + "Failed to clone advertised MCP tool metadata."); + } + + /// + /// Tries to get a tool by name + /// + public bool TryGetTool(string toolName, out IMcpTool? tool) + { + return Volatile.Read(ref _snapshot).Tools.TryGetValue(toolName, out tool); + } + + private static string ValidateToolName(Tool metadata) + { + string toolName = metadata.Name ?? string.Empty; if (string.IsNullOrWhiteSpace(toolName)) { throw new DataApiBuilderException( @@ -35,68 +155,235 @@ public void RegisterTool(IMcpTool tool) subStatusCode: DataApiBuilderException.SubStatusCodes.ErrorInInitialization); } - // Check for duplicate tool names (case-insensitive) - if (_tools.TryGetValue(toolName, out IMcpTool? existingTool)) + if (!string.Equals(toolName, toolName.Trim(), StringComparison.Ordinal)) { - // If the same tool instance is already registered, skip silently. - // This can happen when both McpToolRegistryInitializer (hosted service) - // and McpStdioHelper register tools during stdio mode startup. - if (ReferenceEquals(existingTool, tool)) - { - return; - } - - string existingToolType = existingTool.ToolType == ToolType.BuiltIn ? "built-in" : "custom"; - string newToolType = tool.ToolType == ToolType.BuiltIn ? "built-in" : "custom"; - throw new DataApiBuilderException( - message: $"Duplicate MCP tool name '{toolName}' detected. " + - $"A {existingToolType} tool with this name is already registered. " + - $"Cannot register {newToolType} tool with the same name. " + - $"Tool names must be unique across all tool types.", + message: "MCP tool name cannot contain leading or trailing whitespace.", statusCode: HttpStatusCode.ServiceUnavailable, subStatusCode: DataApiBuilderException.SubStatusCodes.ErrorInInitialization); } - _tools[toolName] = tool; + return toolName; } - /// - /// Gets metadata for all registered tools that are enabled in the given runtime configuration. - /// - public IEnumerable GetEnabledTools(RuntimeConfig config) + private static DataApiBuilderException CreateDuplicateToolException( + string toolName, + IMcpTool existingTool, + IMcpTool newTool) { - return _tools.Values - .Where(t => t.IsEnabled(config)) - .Select(t => t.GetToolMetadata()); + string existingToolType = existingTool.ToolType == ToolType.BuiltIn ? "built-in" : "custom"; + string newToolType = newTool.ToolType == ToolType.BuiltIn ? "built-in" : "custom"; + + return new DataApiBuilderException( + message: $"Duplicate MCP tool name '{toolName}' detected. " + + $"A {existingToolType} tool with this name is already registered. " + + $"Cannot register {newToolType} tool with the same name. " + + $"Tool names must be unique across all tool types.", + statusCode: HttpStatusCode.ServiceUnavailable, + subStatusCode: DataApiBuilderException.SubStatusCodes.ErrorInInitialization); } - /// - /// Tries to get a tool by name - /// - public bool TryGetTool(string toolName, out IMcpTool? tool) + private static ImmutableArray SortMetadata(IEnumerable metadata) { - return _tools.TryGetValue(toolName, out tool); + return metadata + .OrderBy(tool => tool.Name, StringComparer.OrdinalIgnoreCase) + .ThenBy(tool => tool.Name, StringComparer.Ordinal) + .ToImmutableArray(); + } + + private static string CreateDiscoveryJson(ImmutableArray metadata) + { + return JsonSerializer.Serialize( + metadata.ToArray(), + _discoveryJsonOptions); + } + + private static string CreateDiscoveryCanonicalJson(string discoveryJson) + { + using JsonDocument serializedMetadata = JsonDocument.Parse(discoveryJson); + using MemoryStream canonicalJson = new(); + using (Utf8JsonWriter writer = new(canonicalJson)) + { + writer.WriteStartArray(); + foreach (JsonElement tool in serializedMetadata.RootElement.EnumerateArray()) + { + WriteCanonicalJson(writer, tool, CanonicalJsonContext.Tool); + } + + writer.WriteEndArray(); + } + + return Encoding.UTF8.GetString(canonicalJson.ToArray()); + } + + private static Tool CloneMetadata(Tool metadata) + { + ArgumentNullException.ThrowIfNull(metadata); + + byte[] serializedMetadata = JsonSerializer.SerializeToUtf8Bytes( + metadata, + _discoveryJsonOptions); + return JsonSerializer.Deserialize(serializedMetadata, _discoveryJsonOptions) + ?? throw new InvalidOperationException("Failed to clone MCP tool metadata."); + } + + private static void WriteCanonicalJson( + Utf8JsonWriter writer, + JsonElement element, + CanonicalJsonContext context) + { + switch (element.ValueKind) + { + case JsonValueKind.Object: + writer.WriteStartObject(); + foreach (JsonProperty property in element + .EnumerateObject() + .OrderBy(property => property.Name, StringComparer.Ordinal)) + { + writer.WritePropertyName(property.Name); + WriteCanonicalJson( + writer, + property.Value, + GetCanonicalPropertyContext(context, property)); + } + + writer.WriteEndObject(); + break; + + case JsonValueKind.Array: + writer.WriteStartArray(); + if (context != CanonicalJsonContext.SchemaStringSet || + !TryWriteOrderInsensitiveJsonSchemaStringArray(writer, element)) + { + CanonicalJsonContext itemContext = context == CanonicalJsonContext.SchemaArray + ? CanonicalJsonContext.Schema + : CanonicalJsonContext.Data; + foreach (JsonElement item in element.EnumerateArray()) + { + WriteCanonicalJson(writer, item, itemContext); + } + } + + writer.WriteEndArray(); + break; + + case JsonValueKind.String: + writer.WriteStringValue(element.GetString()); + break; + + case JsonValueKind.Number: + case JsonValueKind.True: + case JsonValueKind.False: + case JsonValueKind.Null: + element.WriteTo(writer); + break; + + default: + throw new InvalidOperationException( + $"Unsupported JSON value kind '{element.ValueKind}' in MCP tool metadata."); + } } /// - /// Initializes and registers all MCP tools, enriching custom tools with DB metadata schemas. - /// Shared by both HTTP hosted-service and stdio startup paths. + /// Only the tool's direct input/output schema fields introduce schema context. Within a + /// schema, follow known subschema keywords; defaults, constants, examples, enum instances, + /// and unknown extensions remain ordinary data at every nesting depth. /// - public static void InitializeAndRegisterTools( - IEnumerable tools, - McpToolRegistry registry, - IServiceProvider serviceProvider) + private static CanonicalJsonContext GetCanonicalPropertyContext( + CanonicalJsonContext context, + JsonProperty property) { - foreach (IMcpTool tool in tools) + return context switch { - if (tool is DynamicCustomTool customTool) + CanonicalJsonContext.Tool when property.Name is "inputSchema" or "outputSchema" => + CanonicalJsonContext.Schema, + // Map keys are instance property/definition names, not schema keywords. Legacy + // dependencies can also contain string arrays, which the schema context preserves. + CanonicalJsonContext.SchemaMap => CanonicalJsonContext.Schema, + CanonicalJsonContext.Schema => property.Name switch { - customTool.InitializeMetadata(serviceProvider); + "required" or "type" or "enum" => CanonicalJsonContext.SchemaStringSet, + "properties" or "patternProperties" or "$defs" or "definitions" or + "dependentSchemas" or "dependencies" => CanonicalJsonContext.SchemaMap, + "allOf" or "anyOf" or "oneOf" or "prefixItems" => CanonicalJsonContext.SchemaArray, + // Before draft 2020-12, items also allowed an ordered array of tuple schemas. + "items" when property.Value.ValueKind == JsonValueKind.Array => CanonicalJsonContext.SchemaArray, + "items" or "additionalItems" or "additionalProperties" or "unevaluatedItems" or + "unevaluatedProperties" or "contains" or "propertyNames" or "not" or + "if" or "then" or "else" or "contentSchema" => CanonicalJsonContext.Schema, + _ => CanonicalJsonContext.Data + }, + _ => CanonicalJsonContext.Data + }; + } + + private static bool TryWriteOrderInsensitiveJsonSchemaStringArray( + Utf8JsonWriter writer, + JsonElement element) + { + // Called only for a schema's own required/type/enum keyword. Non-string enum entries + // are instance data and must not be traversed as schemas when this returns false. + List values = new(); + foreach (JsonElement item in element.EnumerateArray()) + { + if (item.ValueKind != JsonValueKind.String) + { + return false; } - registry.RegisterTool(tool); + values.Add(item.GetString()!); + } + + values.Sort(StringComparer.Ordinal); + foreach (string value in values) + { + writer.WriteStringValue(value); } + + return true; + } + + private enum CanonicalJsonContext + { + Data, + Tool, + Schema, + SchemaMap, + SchemaArray, + SchemaStringSet + } + + private sealed record McpToolRegistrySnapshot( + long Version, + ImmutableDictionary Tools, + int AdvertisedToolCount, + string DiscoveryJson, + string DiscoveryCanonicalJson) + { + public static McpToolRegistrySnapshot Empty { get; } = new( + Version: 0, + Tools: ImmutableDictionary.Create(StringComparer.OrdinalIgnoreCase), + AdvertisedToolCount: 0, + DiscoveryJson: "[]", + DiscoveryCanonicalJson: "[]"); } } + + /// + /// Describes the result of atomically replacing an MCP registry snapshot. + /// + internal readonly record struct McpToolRegistryUpdateResult( + long Version, + bool DiscoveryChanged, + int RegisteredToolCount, + int AdvertisedToolCount); + + /// + /// A fully materialized and validated registry generation awaiting publication. + /// + internal sealed record McpToolRegistryCandidate( + ImmutableDictionary Tools, + int AdvertisedToolCount, + string DiscoveryJson, + string DiscoveryCanonicalJson); } diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs deleted file mode 100644 index a7c323a967..0000000000 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs +++ /dev/null @@ -1,36 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. - -using Azure.DataApiBuilder.Mcp.Model; -using Microsoft.Extensions.DependencyInjection; -using Microsoft.Extensions.Hosting; - -namespace Azure.DataApiBuilder.Mcp.Core -{ - /// - /// Hosted service to initialize the MCP tool registry - /// - public class McpToolRegistryInitializer : IHostedService - { - private readonly IServiceProvider _serviceProvider; - private readonly McpToolRegistry _toolRegistry; - - public McpToolRegistryInitializer(IServiceProvider serviceProvider, McpToolRegistry toolRegistry) - { - _serviceProvider = serviceProvider; - _toolRegistry = toolRegistry; - } - - public Task StartAsync(CancellationToken cancellationToken) - { - IEnumerable tools = _serviceProvider.GetServices(); - McpToolRegistry.InitializeAndRegisterTools(tools, _toolRegistry, _serviceProvider); - return Task.CompletedTask; - } - - public Task StopAsync(CancellationToken cancellationToken) - { - return Task.CompletedTask; - } - } -} diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs new file mode 100644 index 0000000000..1e3675ec24 --- /dev/null +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs @@ -0,0 +1,252 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using Azure.DataApiBuilder.Config; +using Azure.DataApiBuilder.Config.ObjectModel; +using Azure.DataApiBuilder.Core.Configurations; +using Azure.DataApiBuilder.Core.Services.MetadataProviders; +using Azure.DataApiBuilder.Mcp.Model; +using Microsoft.Extensions.Hosting; +using Microsoft.Extensions.Logging; +using static Azure.DataApiBuilder.Config.DabConfigEvents; +using static Azure.DataApiBuilder.Mcp.Model.McpEnums; + +namespace Azure.DataApiBuilder.Mcp.Core +{ + /// + /// Shared initialization contract used by hosted HTTP startup and the manually started stdio host. + /// + public interface IMcpToolRegistryRefreshService + { + /// + /// Initializes the registry for the current runtime configuration. Repeated calls for the + /// same successfully applied configuration are no-ops. + /// + void EnsureInitialized(); + + /// + /// Initializes the registry with cooperative cancellation. Implementations that do not + /// override this member retain their existing initialization behavior. + /// + void EnsureInitialized(CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + EnsureInitialized(); + } + } + + /// + /// Builds complete MCP tool-registry generations at startup and after ordered config reloads. + /// + public sealed class McpToolRegistryRefreshService : + IMcpToolRegistryRefreshService, + IHostedService + { + private readonly RuntimeConfigProvider _runtimeConfigProvider; + private readonly IReadOnlyList _registeredTools; + private readonly McpToolRegistry _toolRegistry; + private readonly IMetadataProviderFactory _metadataProviderFactory; + private readonly IReadOnlyList _notifiers; + private readonly ILogger _logger; + private readonly object _refreshLock = new(); + // RuntimeConfigLoader publishes a new object for every parsed generation. Reference + // identity is therefore the generation token used by both idempotency and stale guards. + private RuntimeConfig? _lastAppliedConfig; + + public McpToolRegistryRefreshService( + RuntimeConfigProvider runtimeConfigProvider, + IEnumerable tools, + McpToolRegistry toolRegistry, + IMetadataProviderFactory metadataProviderFactory, + IEnumerable notifiers, + ILogger logger, + HotReloadEventHandler? hotReloadEventHandler = null) + { + _runtimeConfigProvider = runtimeConfigProvider; + // Configuration-generated DynamicCustomTool instances are created separately for each + // generation. Every tool explicitly registered in DI remains an independent extension + // and must be retained regardless of its declared ToolType. + _registeredTools = tools.ToArray(); + _toolRegistry = toolRegistry; + _metadataProviderFactory = metadataProviderFactory; + _notifiers = notifiers.ToArray(); + _logger = logger; + + hotReloadEventHandler?.Subscribe( + MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED, + OnConfigChanged); + } + + /// + public void EnsureInitialized() + { + EnsureInitialized(CancellationToken.None); + } + + /// + public void EnsureInitialized(CancellationToken cancellationToken) + { + if (RefreshRegistry( + forceRebuildForCurrentConfig: false, + cancellationToken)) + { + NotifyToolsListChanged(); + } + } + + /// + public Task StartAsync(CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + + // Startup.Configure initializes database metadata after hosted services start. + // The HTTP startup orchestrator calls EnsureInitialized once that dependency is + // ready. Keeping this hosted-service registration ensures this singleton is created + // early enough to subscribe to ordered hot-reload events without publishing a + // config-only schema first. + return Task.CompletedTask; + } + + /// + public Task StopAsync(CancellationToken cancellationToken) + { + return Task.CompletedTask; + } + + private void OnConfigChanged(object? sender, HotReloadEventArgs args) + { + try + { + args.CancellationToken.ThrowIfCancellationRequested(); + // The runtime config becomes current before its ordered dependency events run. + // An out-of-band EnsureInitialized call can therefore observe this config while + // the metadata provider still represents the previous generation. Always rebuild + // at the ordered MCP event, after metadata and authorization have been refreshed. + if (RefreshRegistry( + forceRebuildForCurrentConfig: true, + args.CancellationToken)) + { + // Transport notification is deliberately outside _refreshLock. Implementations + // must enqueue any potentially blocking I/O so the ordered reload pipeline can + // continue to GraphQL and logging handlers. + NotifyToolsListChanged(); + } + } + catch (OperationCanceledException) when (args.CancellationToken.IsCancellationRequested) + { + // Host shutdown canceled this generation before publication. + } + catch (Exception ex) + { + _logger.LogError( + ex, + "Failed to refresh the MCP tool registry after a runtime configuration change. " + + "The previous registry snapshot remains active."); + } + } + + /// + /// when an initialized client should be notified after the writer + /// lock is released; otherwise . + /// + private bool RefreshRegistry( + bool forceRebuildForCurrentConfig, + CancellationToken cancellationToken) + { + lock (_refreshLock) + { + cancellationToken.ThrowIfCancellationRequested(); + RuntimeConfig config = _runtimeConfigProvider.GetConfig(); + if (!forceRebuildForCurrentConfig && ReferenceEquals(config, _lastAppliedConfig)) + { + return false; + } + + List customTools = CustomMcpToolFactory + .CreateCustomTools(config, _logger) + .ToList(); + + foreach (DynamicCustomTool customTool in customTools) + { + cancellationToken.ThrowIfCancellationRequested(); + bool initializedFromDatabase = customTool.InitializeMetadata( + config, + _metadataProviderFactory, + out string fallbackReason); + if (!initializedFromDatabase) + { + _logger.LogWarning( + "Using configuration-derived input schema for custom MCP tool " + + "'{ToolName}' on entity '{EntityName}'. Reason: {FallbackReason}", + customTool.ToolName, + customTool.EntityName, + fallbackReason); + } + } + + McpToolRegistryCandidate candidate = McpToolRegistry.CreateCandidate( + _registeredTools.Concat(customTools), + config, + cancellationToken); + + cancellationToken.ThrowIfCancellationRequested(); + if (!ReferenceEquals(config, _runtimeConfigProvider.GetConfig())) + { + _logger.LogWarning( + "Discarded a stale MCP tool registry candidate because a newer runtime " + + "configuration became active during the rebuild."); + return false; + } + + bool isInitialGeneration = _lastAppliedConfig is null; + McpToolRegistryUpdateResult result = _toolRegistry.PublishCandidate(candidate); + _lastAppliedConfig = config; + + _logger.LogInformation( + "Published MCP tool registry version {Version} with {BuiltInToolCount} " + + "built-in tools, {RegisteredCustomToolCount} DI-registered custom tools, " + + "{GeneratedCustomToolCount} configuration-generated custom tools, {RegisteredToolCount} " + + "registered tools, and {AdvertisedToolCount} advertised tools. " + + "Discovery changed: {DiscoveryChanged}.", + result.Version, + _registeredTools.Count(tool => tool.ToolType == ToolType.BuiltIn), + _registeredTools.Count(tool => tool.ToolType != ToolType.BuiltIn), + customTools.Count, + result.RegisteredToolCount, + result.AdvertisedToolCount, + result.DiscoveryChanged); + + return !isInitialGeneration && result.DiscoveryChanged; + } + } + + private void NotifyToolsListChanged() + { + foreach (IMcpToolListChangedNotifier notifier in _notifiers) + { + try + { + notifier.NotifyToolsListChanged(); + } + catch (Exception ex) + { + _logger.LogError( + ex, + "Failed to notify an MCP client that the advertised tool list changed."); + } + } + } + } + + /// + /// Transport-specific notification sink for MCP tool discovery changes. + /// + public interface IMcpToolListChangedNotifier + { + /// + /// Enqueues notification of a connected, initialized client that it should refresh + /// tools/list. Implementations must not block on transport I/O. + /// + void NotifyToolsListChanged(); + } +} diff --git a/src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs b/src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs index 2d79649bbb..6ad9323ff6 100644 --- a/src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs +++ b/src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs @@ -3,6 +3,8 @@ using Azure.DataApiBuilder.Config.DatabasePrimitives; using Azure.DataApiBuilder.Config.ObjectModel; +using Azure.DataApiBuilder.Core.Services; +using Azure.DataApiBuilder.Core.Services.MetadataProviders; using Azure.DataApiBuilder.Service.Exceptions; // Added for DataApiBuilderException using Microsoft.Extensions.DependencyInjection; @@ -46,7 +48,7 @@ public static bool TryResolveMetadata( string entityName, RuntimeConfig config, IServiceProvider serviceProvider, - out Azure.DataApiBuilder.Core.Services.ISqlMetadataProvider sqlMetadataProvider, + out ISqlMetadataProvider sqlMetadataProvider, out DatabaseObject dbObject, out string dataSourceName, out string error, @@ -58,21 +60,59 @@ public static bool TryResolveMetadata( dataSourceName = string.Empty; error = string.Empty; - if (string.IsNullOrWhiteSpace(entityName)) + if (!TryValidateEntityName(entityName, out error)) { - error = "Entity name cannot be null or empty."; return false; } // Use GetService (not GetRequiredService) so the helper honours its Try* contract. - Azure.DataApiBuilder.Core.Services.MetadataProviders.IMetadataProviderFactory? metadataProviderFactory = - serviceProvider.GetService(); + IMetadataProviderFactory? metadataProviderFactory = + serviceProvider.GetService(); if (metadataProviderFactory is null) { error = "Metadata provider factory is not registered."; return false; } + return TryResolveMetadata( + entityName, + config, + metadataProviderFactory, + out sqlMetadataProvider, + out dbObject, + out dataSourceName, + out error, + cancellationToken); + } + + /// + /// Resolves database metadata using the exact runtime configuration and metadata-provider + /// generation supplied by the caller. + /// + public static bool TryResolveMetadata( + string entityName, + RuntimeConfig config, + IMetadataProviderFactory metadataProviderFactory, + out ISqlMetadataProvider sqlMetadataProvider, + out DatabaseObject dbObject, + out string dataSourceName, + out string error, + CancellationToken cancellationToken = default) + { + ArgumentNullException.ThrowIfNull(config); + ArgumentNullException.ThrowIfNull(metadataProviderFactory); + + cancellationToken.ThrowIfCancellationRequested(); + sqlMetadataProvider = default!; + dbObject = default!; + dataSourceName = string.Empty; + error = string.Empty; + + if (!TryValidateEntityName(entityName, out error)) + { + return false; + } + // Resolve datasource name for the entity. try { @@ -114,12 +154,25 @@ public static bool TryResolveMetadata( // Validate entity exists in metadata mapping. if (!sqlMetadataProvider.EntityToDatabaseObject.TryGetValue(entityName, out DatabaseObject? temp) || temp is null) { - error = $"Entity '{entityName}' is not defined in the configuration."; + error = $"Database metadata for entity '{entityName}' was not available from " + + $"data source '{dataSourceName}'."; return false; } dbObject = temp; return true; } + + private static bool TryValidateEntityName(string? entityName, out string error) + { + if (string.IsNullOrWhiteSpace(entityName)) + { + error = "Entity name cannot be null or empty."; + return false; + } + + error = string.Empty; + return true; + } } } diff --git a/src/Config/ConfigFileWatcher.cs b/src/Config/ConfigFileWatcher.cs index e1afb39838..3e2abb90df 100644 --- a/src/Config/ConfigFileWatcher.cs +++ b/src/Config/ConfigFileWatcher.cs @@ -6,6 +6,20 @@ namespace Azure.DataApiBuilder.Config; +/// +/// Internal lifecycle contract that allows event delivery to stop independently from potentially +/// blocking disposal of the underlying operating-system watcher. +/// +internal interface IConfigFileWatcher : IDisposable +{ + event EventHandler? NewFileContentsDetected; + + /// + /// Disables new file-system events and detaches the underlying change callback. + /// + void StopWatching(); +} + /// /// This class is responsible for monitoring the config file from the /// local file system. This watcher maintains a file hash to only emit @@ -20,9 +34,11 @@ namespace Azure.DataApiBuilder.Config; /// /// /// -public class ConfigFileWatcher : IDisposable +public class ConfigFileWatcher : IConfigFileWatcher { + private readonly object _lifecycleLock = new(); private bool _disposed; + private bool _stopped; /// /// Watches a specific file for modifications and alerts @@ -93,12 +109,18 @@ private void OnConfigFileChange(object sender, FileSystemEventArgs e) { try { - if (_fileWatcher is not null) + IFileSystemWatcher? fileWatcher; + lock (_lifecycleLock) + { + fileWatcher = _stopped ? null : _fileWatcher; + } + + if (fileWatcher is not null) { // Multiple file change notifications may be raised for a single file change. // Use file hashes to ensure that HotReload operation is only executed when a net-new // runtime config is detected. - byte[] updatedRuntimeConfigFileHash = FileUtilities.ComputeHash(_fileWatcher.FileSystem, filePath: Path.Combine(WatchedDirectory, WatchedFile)); + byte[] updatedRuntimeConfigFileHash = FileUtilities.ComputeHash(fileWatcher.FileSystem, filePath: Path.Combine(WatchedDirectory, WatchedFile)); if (!_runtimeConfigHash.SequenceEqual(updatedRuntimeConfigFileHash)) { _runtimeConfigHash = updatedRuntimeConfigFileHash; @@ -129,18 +151,43 @@ private void OnConfigFileChange(object sender, FileSystemEventArgs e) /// public void Dispose() { - if (_disposed) + IFileSystemWatcher? fileWatcher; + lock (_lifecycleLock) { - return; + if (_disposed) + { + return; + } + + _disposed = true; + StopWatchingCore(); + fileWatcher = _fileWatcher; + _fileWatcher = null; } - _disposed = true; + fileWatcher?.Dispose(); + } + + void IConfigFileWatcher.StopWatching() + { + lock (_lifecycleLock) + { + StopWatchingCore(); + } + } + + private void StopWatchingCore() + { + if (_stopped) + { + return; + } + _stopped = true; if (_fileWatcher is not null) { _fileWatcher.EnableRaisingEvents = false; _fileWatcher.Changed -= OnConfigFileChange; - _fileWatcher.Dispose(); } } } diff --git a/src/Config/DabConfigEvents.cs b/src/Config/DabConfigEvents.cs index f69193b583..9162a70526 100644 --- a/src/Config/DabConfigEvents.cs +++ b/src/Config/DabConfigEvents.cs @@ -15,6 +15,7 @@ public static class DabConfigEvents public const string POSTGRESQL_QUERY_EXECUTOR_ON_CONFIG_CHANGED = "POSTGRESQL_QUERY_EXECUTOR_ON_CONFIG_CHANGED"; public const string DOCUMENTOR_ON_CONFIG_CHANGED = "DOCUMENTOR_ON_CONFIG_CHANGED"; public const string AUTHZ_RESOLVER_ON_CONFIG_CHANGED = "AUTHZ_RESOLVER_ON_CONFIG_CHANGED"; + public const string MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED = "MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED"; public const string GRAPHQL_SCHEMA_REFRESH_ON_CONFIG_CHANGED = "GRAPHQL_SCHEMA_REFRESH_ON_CONFIG_CHANGED"; public const string GRAPHQL_SCHEMA_EVICTION_ON_CONFIG_CHANGED = "GRAPHQL_SCHEMA_EVICTION_ON_CONFIG_CHANGED"; public const string GRAPHQL_SCHEMA_CREATOR_ON_CONFIG_CHANGED = "GRAPHQL_SCHEMA_CREATOR_ON_CONFIG_CHANGED"; diff --git a/src/Config/FileSystemRuntimeConfigLoader.cs b/src/Config/FileSystemRuntimeConfigLoader.cs index e3529c696f..5c487ab4dc 100644 --- a/src/Config/FileSystemRuntimeConfigLoader.cs +++ b/src/Config/FileSystemRuntimeConfigLoader.cs @@ -32,7 +32,16 @@ namespace Azure.DataApiBuilder.Config; /// public class FileSystemRuntimeConfigLoader : RuntimeConfigLoader, IDisposable { - private bool _disposed; + private readonly SemaphoreSlim _hotReloadGate = new(initialCount: 1, maxCount: 1); + private readonly CancellationTokenSource _disposeCancellation = new(); + private readonly object _operationLock = new(); + private readonly object _watcherLock = new(); + private readonly Func _configFileWatcherFactory; + private TaskCompletionSource _activeOperationsDrained = CreateCompletedDrainSignal(); + private Task _shutdownCompleted = Task.CompletedTask; + private int _activeOperationCount; + private int _disposed; + private int _shutdownResourcesDisposed; /// /// This stores either the default config name e.g. dab-config.json /// or user provided config file which could be a relative file path, @@ -53,7 +62,7 @@ public class FileSystemRuntimeConfigLoader : RuntimeConfigLoader, IDisposable /// /// Watches the config file for changes and triggers hot-reload when a change is detected. /// - private ConfigFileWatcher? _configFileWatcher; + private IConfigFileWatcher? _configFileWatcher; /// /// File system abstraction used to interact with the runtime config file. @@ -96,9 +105,34 @@ public FileSystemRuntimeConfigLoader( string? connectionString = null, bool isCliLoader = false, ILogger? logger = null) + : this( + fileSystem, + handler, + baseConfigFilePath, + connectionString, + isCliLoader, + logger, + static (watcherFileSystem, directoryName, configFileName) => + new ConfigFileWatcher( + new FileSystemWatcherWrapper(watcherFileSystem), + directoryName, + configFileName)) + { + } + + internal FileSystemRuntimeConfigLoader( + IFileSystem fileSystem, + HotReloadEventHandler? handler, + string baseConfigFilePath, + string? connectionString, + bool isCliLoader, + ILogger? logger, + Func configFileWatcherFactory) : base(handler, connectionString) { - _fileSystem = fileSystem; + _fileSystem = fileSystem ?? throw new ArgumentNullException(nameof(fileSystem)); + _configFileWatcherFactory = configFileWatcherFactory ?? + throw new ArgumentNullException(nameof(configFileWatcherFactory)); _baseConfigFilePath = baseConfigFilePath; ConfigFilePath = GetFinalConfigFilePath(); _isCliLoader = isCliLoader; @@ -106,23 +140,123 @@ public FileSystemRuntimeConfigLoader( } /// - /// Disposes the config file watcher to release file handles and stop - /// monitoring the config file for changes. + /// Stops admitting new work and requests cancellation of active work. Coordinated hosts call + /// before disposing dependencies when they need to drain active work. + /// Synchronous disposal does not wait indefinitely for an uncooperative event subscriber. /// public void Dispose() { - if (_disposed) + _ = BeginShutdown(); + } + + /// + /// Stops accepting hot-reload work, requests cancellation of the active operation, and waits + /// until all serialized work has exited. Host shutdown calls this before singleton disposal. + /// + public async Task StopAsync(CancellationToken cancellationToken) + { + Task shutdownCompleted = BeginShutdown(); + await shutdownCompleted.WaitAsync(cancellationToken).ConfigureAwait(false); + } + + internal bool ShutdownResourcesDisposed => + Volatile.Read(ref _shutdownResourcesDisposed) != 0; + + private Task BeginShutdown() + { + bool firstShutdownRequest; + TaskCompletionSource? cancellationCallbacksCompleted = null; + Task shutdownCompleted; + lock (_operationLock) { - return; + firstShutdownRequest = Interlocked.Exchange(ref _disposed, 1) == 0; + if (firstShutdownRequest) + { + cancellationCallbacksCompleted = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + Task drainCompleted = Task.WhenAll( + _activeOperationsDrained.Task, + cancellationCallbacksCompleted.Task); + _shutdownCompleted = DisposeSynchronizationResourcesAfterDrainAsync( + drainCompleted); + } + + shutdownCompleted = _shutdownCompleted; } - _disposed = true; + if (firstShutdownRequest) + { + // Cancellation callbacks are user-extensible and may block. CancelAsync marks the + // token canceled without running those callbacks on the host's stopping thread. + _ = RequestOperationCancellationAsync(cancellationCallbacksCompleted!); + StopAndDisposeConfigFileWatcher(); + } + + return shutdownCompleted; + } - if (_configFileWatcher is not null) + private async Task DisposeSynchronizationResourcesAfterDrainAsync(Task drainCompleted) + { + await drainCompleted.ConfigureAwait(false); + + // Every operation admitted before shutdown, including gate waiters, has exited and every + // cancellation callback has completed. SemaphoreSlim and CancellationTokenSource can now + // be disposed without racing Wait, Release, token registration, or CancelAsync. + _hotReloadGate.Dispose(); + _disposeCancellation.Dispose(); + Volatile.Write(ref _shutdownResourcesDisposed, 1); + } + + private async Task RequestOperationCancellationAsync( + TaskCompletionSource cancellationCallbacksCompleted) + { + try + { + await _disposeCancellation.CancelAsync().ConfigureAwait(false); + } + catch (Exception ex) + { + SendLogToBufferOrLogger( + LogLevel.Warning, + $"A hot-reload cancellation callback failed during shutdown due to {ex.Message}"); + } + finally + { + cancellationCallbacksCompleted.TrySetResult(); + } + } + + private void StopAndDisposeConfigFileWatcher() + { + + IConfigFileWatcher? configFileWatcher; + lock (_watcherLock) { - _configFileWatcher.NewFileContentsDetected -= OnNewFileContentsDetected; - _configFileWatcher.Dispose(); + configFileWatcher = _configFileWatcher; _configFileWatcher = null; + + if (configFileWatcher is not null) + { + configFileWatcher.NewFileContentsDetected -= OnNewFileContentsDetected; + } + } + + if (configFileWatcher is not null) + { + try + { + configFileWatcher.StopWatching(); + } + catch (Exception ex) + { + SendLogToBufferOrLogger( + LogLevel.Warning, + $"Unable to disable the configuration file watcher during shutdown due to {ex.Message}"); + } + + // Underlying FileSystemWatcher disposal can block while an OS callback completes. + // Dispose it on a background worker so host shutdown never waits for an active reload. + ScheduleConfigFileWatcherDisposal(configFileWatcher); } } @@ -159,57 +293,264 @@ public string GetConfigFileName() /// private bool TrySetupConfigFileWatcher() { - // File watching / hot-reload isn't used for the CLI. - if (_isCliLoader) + lock (_watcherLock) { + // File watching / hot-reload isn't used for the CLI and must not start once disposal + // begins, including when disposal races with initial configuration loading. + if (_isCliLoader || IsDisposed) + { + return false; + } + + // If the file watcher is already set up, we don't need to do it again. + if (_configFileWatcher is not null) + { + return false; + } + + if (RuntimeConfig is not null) + { + try + { + _configFileWatcher = _configFileWatcherFactory( + _fileSystem, + GetConfigDirectoryName(), + GetConfigFileName()); + _configFileWatcher.NewFileContentsDetected += OnNewFileContentsDetected; + } + catch (Exception ex) + { + // Need to remove the dependencies in startup on the RuntimeConfigProvider + // before we can have an ILogger here. + Console.WriteLine($"Attempt to configure config file watcher for hot reload failed due to: {ex.Message}."); + } + + return _configFileWatcher is not null; + } + return false; } + } - // If the file watcher is already set up, we don't need to do it again. - if (_configFileWatcher is not null) + /// + /// When a change is detected in the Config file being watched this trigger + /// function is called and handles the hot reload logic when appropriate, + /// ie: in a local development scenario. + /// + private void OnNewFileContentsDetected(object? sender, EventArgs e) + { + ProcessHotReloadNotification(); + } + + /// + /// Processes one file-change notification while serializing the complete hot-reload pipeline + /// for this loader instance. The gate begins before the current configuration is inspected and + /// remains held through all synchronous + /// handlers so dependencies cannot be mixed across generations. + /// + /// + /// Optional observer invoked immediately before waiting for the serialization gate. This is + /// used by deterministic concurrency tests to prove a second notification reached the gate. + /// + internal void ProcessHotReloadNotification(Action? beforeEnteringGate = null) + { + beforeEnteringGate?.Invoke(); + + if (!TryBeginSerializedOperation()) { - return false; + return; } - if (RuntimeConfig is not null) + bool gateEntered = false; + try { try { - _configFileWatcher = new(new FileSystemWatcherWrapper(_fileSystem), GetConfigDirectoryName(), GetConfigFileName()); - _configFileWatcher.NewFileContentsDetected += OnNewFileContentsDetected; + _hotReloadGate.Wait(_disposeCancellation.Token); + gateEntered = true; + } + catch (OperationCanceledException) when (IsDisposed) + { + return; + } + + try + { + if (RuntimeConfig is not null) + { + HotReloadConfig( + RuntimeConfig.IsDevelopmentMode(), + _disposeCancellation.Token); + } + } + catch (OperationCanceledException) when (IsDisposed) + { + // Host shutdown canceled this generation before it could finish publication. } catch (Exception ex) { - // Need to remove the dependencies in startup on the RuntimeConfigProvider - // before we can have an ILogger here. - Console.WriteLine($"Attempt to configure config file watcher for hot reload failed due to: {ex.Message}."); + SendLogToBufferOrLogger( + LogLevel.Error, + $"Unable to hot reload configuration file due to {ex.Message}"); + } + } + finally + { + if (gateEntered) + { + // Release before completing the tracked operation. The final operation can make + // the shutdown continuation dispose the semaphore immediately. + _hotReloadGate.Release(); } - return _configFileWatcher is not null; + EndSerializedOperation(); } + } - return false; + /// + /// Executes initial runtime dependency construction under the same per-loader gate used by + /// file-triggered hot reload. The operation may be asynchronous, and the gate remains held + /// until it completes so a configuration cannot change between metadata initialization and + /// dependent component publication. + /// + /// The complete initial configuration operation to serialize. + public Task ExecuteWithHotReloadSerializationAsync(Func operation) + { + ArgumentNullException.ThrowIfNull(operation); + return ExecuteWithHotReloadSerializationAsync(_ => operation()); } /// - /// When a change is detected in the Config file being watched this trigger - /// function is called and handles the hot reload logic when appropriate, - /// ie: in a local development scenario. + /// Executes initial runtime dependency construction under the same per-loader gate used by + /// file-triggered hot reload, with cooperative shutdown cancellation. /// - private void OnNewFileContentsDetected(object? sender, EventArgs e) + /// + /// The complete initial configuration operation to serialize. The supplied token is canceled + /// when loader shutdown begins. + /// + public async Task ExecuteWithHotReloadSerializationAsync( + Func operation) { + ArgumentNullException.ThrowIfNull(operation); + + if (IsDisposed) + { + throw new ObjectDisposedException(nameof(FileSystemRuntimeConfigLoader)); + } + + if (!TryBeginSerializedOperation()) + { + throw new ObjectDisposedException(nameof(FileSystemRuntimeConfigLoader)); + } + + bool gateEntered = false; try { - if (RuntimeConfig is not null) + try + { + await _hotReloadGate.WaitAsync(_disposeCancellation.Token).ConfigureAwait(false); + gateEntered = true; + } + catch (OperationCanceledException) when (IsDisposed) + { + throw new ObjectDisposedException(nameof(FileSystemRuntimeConfigLoader)); + } + + await operation(_disposeCancellation.Token).ConfigureAwait(false); + } + finally + { + if (gateEntered) + { + _hotReloadGate.Release(); + } + + EndSerializedOperation(); + } + } + + private bool IsDisposed => Volatile.Read(ref _disposed) != 0; + + private bool TryBeginSerializedOperation() + { + lock (_operationLock) + { + if (IsDisposed) + { + return false; + } + + if (_activeOperationCount++ == 0) + { + _activeOperationsDrained = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + } + + return true; + } + } + + private void EndSerializedOperation() + { + TaskCompletionSource? drainedSignal = null; + lock (_operationLock) + { + if (--_activeOperationCount == 0) + { + drainedSignal = _activeOperationsDrained; + } + } + + drainedSignal?.TrySetResult(); + } + + private static TaskCompletionSource CreateCompletedDrainSignal() + { + TaskCompletionSource signal = new( + TaskCreationOptions.RunContinuationsAsynchronously); + signal.SetResult(); + return signal; + } + + private void ScheduleConfigFileWatcherDisposal(IConfigFileWatcher configFileWatcher) + { + Action disposeWatcher = () => + { + try + { + configFileWatcher.Dispose(); + } + catch (Exception ex) { - HotReloadConfig(RuntimeConfig.IsDevelopmentMode()); + SendLogToBufferOrLogger( + LogLevel.Warning, + $"Unable to dispose the configuration file watcher due to {ex.Message}"); } + }; + + if (ThreadPool.QueueUserWorkItem( + static callback => callback(), + disposeWatcher, + preferLocal: false)) + { + return; + } + + try + { + Thread fallbackWorker = new( + static callback => ((Action)callback!).Invoke()) + { + IsBackground = true, + Name = "DAB configuration watcher disposal" + }; + fallbackWorker.Start(disposeWatcher); } catch (Exception ex) { - // Need to remove the dependencies in startup on the RuntimeConfigProvider - // before we can have an ILogger here. - Console.WriteLine("Unable to hot reload configuration file due to " + ex.Message); + SendLogToBufferOrLogger( + LogLevel.Warning, + $"Unable to schedule configuration file watcher disposal due to {ex.Message}"); } } @@ -229,6 +570,27 @@ public bool TryLoadConfig( bool? isDevMode = null, DeserializationVariableReplacementSettings? replacementSettings = null) { + return TryLoadConfig( + path, + out config, + logger, + isDevMode, + replacementSettings, + CancellationToken.None); + } + + /// + /// Loads runtime configuration with cooperative cancellation for retry waits. + /// + public bool TryLoadConfig( + string path, + [NotNullWhen(true)] out RuntimeConfig? config, + ILogger? logger, + bool? isDevMode, + DeserializationVariableReplacementSettings? replacementSettings, + CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); IsParseErrorEmitted = false; if (_fileSystem.File.Exists(path)) { @@ -244,6 +606,7 @@ public bool TryLoadConfig( string json = string.Empty; while (runCount <= FileUtilities.RunLimit) { + cancellationToken.ThrowIfCancellationRequested(); try { json = _fileSystem.File.ReadAllText(path); @@ -258,7 +621,13 @@ public bool TryLoadConfig( throw; } - Thread.Sleep(TimeSpan.FromSeconds(Math.Pow(FileUtilities.ExponentialRetryBase, runCount))); + TimeSpan retryDelay = TimeSpan.FromSeconds( + Math.Pow(FileUtilities.ExponentialRetryBase, runCount)); + if (cancellationToken.WaitHandle.WaitOne(retryDelay)) + { + cancellationToken.ThrowIfCancellationRequested(); + } + runCount++; } } @@ -341,14 +710,23 @@ public override bool TryLoadKnownConfig([NotNullWhen(true)] out RuntimeConfig? c /// Hot Reloads the runtime config when the file watcher /// is active and detects a change to the underlying config file. /// - private void HotReloadConfig(bool isDevMode, ILogger? logger = null) + private void HotReloadConfig(bool isDevMode, CancellationToken cancellationToken) { - logger?.LogInformation(message: "Starting hot-reload process for config: {ConfigFilePath}", ConfigFilePath); + cancellationToken.ThrowIfCancellationRequested(); + SendLogToBufferOrLogger( + LogLevel.Information, + $"Starting hot-reload process for config: {ConfigFilePath}"); // Use default replacement settings for hot reload DeserializationVariableReplacementSettings replacementSettings = new(azureKeyVaultOptions: null, doReplaceEnvVar: true, doReplaceAkvVar: true); - if (!TryLoadConfig(ConfigFilePath, out _, logger: logger, isDevMode: isDevMode, replacementSettings: replacementSettings)) + if (!TryLoadConfig( + ConfigFilePath, + out _, + logger: null, + isDevMode: isDevMode, + replacementSettings: replacementSettings, + cancellationToken: cancellationToken)) { throw new DataApiBuilderException( message: "Deserialization of the configuration file failed.", @@ -358,14 +736,14 @@ private void HotReloadConfig(bool isDevMode, ILogger? logger = null) IsNewConfigDetected = true; IsNewConfigValidated = false; - SignalConfigChanged(); + SignalConfigChanged(message: string.Empty, cancellationToken); // Telemetry (and any other) logs buffered during the reload parse are otherwise only // drained once at startup. Flush them now so hot-reload logs are actually emitted and the // shared static buffer does not accumulate entries across successive reloads. FlushLogBuffer(); - logger?.LogInformation("Hot-reload process finished."); + SendLogToBufferOrLogger(LogLevel.Information, "Hot-reload process finished."); } /// diff --git a/src/Config/HotReloadEventArgs.cs b/src/Config/HotReloadEventArgs.cs index 5fa20e8d8d..6c82efb2eb 100644 --- a/src/Config/HotReloadEventArgs.cs +++ b/src/Config/HotReloadEventArgs.cs @@ -9,9 +9,23 @@ public class HotReloadEventArgs : EventArgs public string Message { get; set; } + /// + /// Cancels the current ordered hot-reload generation during loader shutdown. + /// + public CancellationToken CancellationToken { get; } + public HotReloadEventArgs(string eventName, string message) + : this(eventName, message, CancellationToken.None) + { + } + + public HotReloadEventArgs( + string eventName, + string message, + CancellationToken cancellationToken) { EventName = eventName; Message = message; + CancellationToken = cancellationToken; } } diff --git a/src/Config/HotReloadEventHandler.cs b/src/Config/HotReloadEventHandler.cs index 666c3c227b..cf905bd202 100644 --- a/src/Config/HotReloadEventHandler.cs +++ b/src/Config/HotReloadEventHandler.cs @@ -31,6 +31,7 @@ public HotReloadEventHandler() { MUTATION_ENGINE_FACTORY_ON_CONFIG_CHANGED, null }, { DOCUMENTOR_ON_CONFIG_CHANGED, null }, { AUTHZ_RESOLVER_ON_CONFIG_CHANGED, null }, + { MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED, null }, { GRAPHQL_SCHEMA_CREATOR_ON_CONFIG_CHANGED, null }, { GRAPHQL_SCHEMA_REFRESH_ON_CONFIG_CHANGED, null }, { GRAPHQL_SCHEMA_EVICTION_ON_CONFIG_CHANGED, null }, @@ -46,6 +47,11 @@ public void OnConfigChangedEvent(object sender, TEventArgs args) } } + /// + /// Subscribes a synchronous ordered hot-reload callback. Handlers must observe + /// and must not block indefinitely; + /// host shutdown is allowed to stop waiting when its configured timeout expires. + /// public void Subscribe(string eventName, EventHandler handler) { if (_eventHandlers.ContainsKey(eventName)) diff --git a/src/Config/RuntimeConfigLoader.cs b/src/Config/RuntimeConfigLoader.cs index 1c0c9c9ac8..e2f9ff795f 100644 --- a/src/Config/RuntimeConfigLoader.cs +++ b/src/Config/RuntimeConfigLoader.cs @@ -109,32 +109,57 @@ protected virtual void OnConfigChangedEvent(HotReloadEventArgs args) /// protected void SignalConfigChanged(string message = "") { + SignalConfigChanged(message, CancellationToken.None); + } + + /// + /// Notifies subscribers of an ordered configuration change with cooperative cancellation. + /// + protected void SignalConfigChanged( + string message, + CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + // Signal that a change has occurred to all change token listeners. RaiseChanged(); // All the data inside of the if statement should only update when DAB is in development mode. if (RuntimeConfig!.IsDevelopmentMode()) { - OnConfigChangedEvent(new HotReloadEventArgs(QUERY_MANAGER_FACTORY_ON_CONFIG_CHANGED, message)); - OnConfigChangedEvent(new HotReloadEventArgs(METADATA_PROVIDER_FACTORY_ON_CONFIG_CHANGED, message)); - OnConfigChangedEvent(new HotReloadEventArgs(QUERY_ENGINE_FACTORY_ON_CONFIG_CHANGED, message)); - OnConfigChangedEvent(new HotReloadEventArgs(MUTATION_ENGINE_FACTORY_ON_CONFIG_CHANGED, message)); - OnConfigChangedEvent(new HotReloadEventArgs(DOCUMENTOR_ON_CONFIG_CHANGED, message)); + RaiseOrderedEvent(QUERY_MANAGER_FACTORY_ON_CONFIG_CHANGED); + RaiseOrderedEvent(METADATA_PROVIDER_FACTORY_ON_CONFIG_CHANGED); + RaiseOrderedEvent(QUERY_ENGINE_FACTORY_ON_CONFIG_CHANGED); + RaiseOrderedEvent(MUTATION_ENGINE_FACTORY_ON_CONFIG_CHANGED); + RaiseOrderedEvent(DOCUMENTOR_ON_CONFIG_CHANGED); // Order of event firing matters: Authorization rules can only be updated after the // MetadataProviderFactory has been updated with latest database object metadata. // RuntimeConfig must already be updated and is implied to have been updated by the time // this function is called. - OnConfigChangedEvent(new HotReloadEventArgs(AUTHZ_RESOLVER_ON_CONFIG_CHANGED, message)); + RaiseOrderedEvent(AUTHZ_RESOLVER_ON_CONFIG_CHANGED); + + // Custom MCP tool schemas depend on refreshed database metadata. Publish the new + // registry only after query, mutation, and authorization dependencies are ready. + RaiseOrderedEvent(MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED); // Order of event firing matters: Eviction must be done before creating a new schema and then updating the schema. - OnConfigChangedEvent(new HotReloadEventArgs(GRAPHQL_SCHEMA_EVICTION_ON_CONFIG_CHANGED, message)); - OnConfigChangedEvent(new HotReloadEventArgs(GRAPHQL_SCHEMA_CREATOR_ON_CONFIG_CHANGED, message)); - OnConfigChangedEvent(new HotReloadEventArgs(GRAPHQL_SCHEMA_REFRESH_ON_CONFIG_CHANGED, message)); + RaiseOrderedEvent(GRAPHQL_SCHEMA_EVICTION_ON_CONFIG_CHANGED); + RaiseOrderedEvent(GRAPHQL_SCHEMA_CREATOR_ON_CONFIG_CHANGED); + RaiseOrderedEvent(GRAPHQL_SCHEMA_REFRESH_ON_CONFIG_CHANGED); } // Log Level Initializer is outside of if statement as it can be updated on both development and production mode. - OnConfigChangedEvent(new HotReloadEventArgs(LOG_LEVEL_INITIALIZER_ON_CONFIG_CHANGE, message)); + RaiseOrderedEvent(LOG_LEVEL_INITIALIZER_ON_CONFIG_CHANGE); + + void RaiseOrderedEvent(string eventName) + { + cancellationToken.ThrowIfCancellationRequested(); + OnConfigChangedEvent(new HotReloadEventArgs( + eventName, + message, + cancellationToken)); + } } /// diff --git a/src/Core/Authorization/AuthorizationResolver.cs b/src/Core/Authorization/AuthorizationResolver.cs index fd0da59393..1b76401a0e 100644 --- a/src/Core/Authorization/AuthorizationResolver.cs +++ b/src/Core/Authorization/AuthorizationResolver.cs @@ -77,6 +77,7 @@ public AuthorizationResolver( /// protected void OnConfigChanged(object? sender, HotReloadEventArgs args) { + args.CancellationToken.ThrowIfCancellationRequested(); SetEntityPermissionMap(_runtimeConfigProvider.GetConfig()); } diff --git a/src/Core/Resolvers/Factories/MutationEngineFactory.cs b/src/Core/Resolvers/Factories/MutationEngineFactory.cs index 08a5fea2e3..7b039682d7 100644 --- a/src/Core/Resolvers/Factories/MutationEngineFactory.cs +++ b/src/Core/Resolvers/Factories/MutationEngineFactory.cs @@ -94,6 +94,7 @@ private void ConfigureMutationEngines() public void OnConfigChanged(object? sender, HotReloadEventArgs args) { + args.CancellationToken.ThrowIfCancellationRequested(); _mutationEngines = new Dictionary(); ConfigureMutationEngines(); } diff --git a/src/Core/Resolvers/Factories/QueryEngineFactory.cs b/src/Core/Resolvers/Factories/QueryEngineFactory.cs index 1d2ae2935d..5c0b4daa11 100644 --- a/src/Core/Resolvers/Factories/QueryEngineFactory.cs +++ b/src/Core/Resolvers/Factories/QueryEngineFactory.cs @@ -91,6 +91,7 @@ public void ConfigureQueryEngines() public void OnConfigChanged(object? sender, HotReloadEventArgs args) { + args.CancellationToken.ThrowIfCancellationRequested(); _queryEngines = new Dictionary(); ConfigureQueryEngines(); } diff --git a/src/Core/Resolvers/Factories/QueryManagerFactory.cs b/src/Core/Resolvers/Factories/QueryManagerFactory.cs index 68896318d5..7a03d8562b 100644 --- a/src/Core/Resolvers/Factories/QueryManagerFactory.cs +++ b/src/Core/Resolvers/Factories/QueryManagerFactory.cs @@ -106,6 +106,7 @@ private void ConfigureQueryManagerFactory() public void OnConfigChanged(object? sender, HotReloadEventArgs args) { + args.CancellationToken.ThrowIfCancellationRequested(); _queryBuilders = new Dictionary(); _queryExecutors = new Dictionary(); _dbExceptionsParsers = new Dictionary(); diff --git a/src/Core/Resolvers/IQueryExecutor.cs b/src/Core/Resolvers/IQueryExecutor.cs index 2eac7242de..6257388a75 100644 --- a/src/Core/Resolvers/IQueryExecutor.cs +++ b/src/Core/Resolvers/IQueryExecutor.cs @@ -34,6 +34,29 @@ public interface IQueryExecutor HttpContext? httpContext = null, List? args = null); + /// + /// Executes SQL text with cooperative cancellation. Implementations that do not override + /// this member retain their existing query execution behavior. + /// + public Task ExecuteQueryAsync( + string sqltext, + IDictionary parameters, + Func?, Task>? dataReaderHandler, + string dataSourceName, + CancellationToken cancellationToken, + HttpContext? httpContext = null, + List? args = null) + { + cancellationToken.ThrowIfCancellationRequested(); + return ExecuteQueryAsync( + sqltext, + parameters, + dataReaderHandler, + dataSourceName, + httpContext, + args); + } + /// /// Executes sql text with the given parameters and /// uses the function dataReaderHandler to process @@ -152,7 +175,23 @@ public Dictionary GetResultProperties( /// /// Modified the properties of the supplied connection to support managed identity access. /// - public Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection conn, string dataSourceName); + public Task SetManagedIdentityAccessTokenIfAnyAsync( + DbConnection conn, + string dataSourceName); + + /// + /// Modifies the supplied connection for managed identity access with cooperative + /// cancellation. Implementations that do not override this member retain their existing + /// access-token behavior. + /// + public Task SetManagedIdentityAccessTokenIfAnyAsync( + DbConnection conn, + string dataSourceName, + CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return SetManagedIdentityAccessTokenIfAnyAsync(conn, dataSourceName); + } /// /// Method to generate the query to send user data to the underlying database which might be used diff --git a/src/Core/Resolvers/MsSqlQueryExecutor.cs b/src/Core/Resolvers/MsSqlQueryExecutor.cs index 2078b4f1c5..26a4582c40 100644 --- a/src/Core/Resolvers/MsSqlQueryExecutor.cs +++ b/src/Core/Resolvers/MsSqlQueryExecutor.cs @@ -339,8 +339,12 @@ private void ConfigureMsSqlQueryExecutor() /// /// The supplied connection to modify for managed identity access. /// Name of datasource for which to set access token. Default dbName taken from config if null - public override async Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection conn, string dataSourceName) + public override async Task SetManagedIdentityAccessTokenIfAnyAsync( + DbConnection conn, + string dataSourceName, + CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); // using default datasource name for first db - maintaining backward compatibility for single db scenario. if (string.IsNullOrEmpty(dataSourceName)) { @@ -359,7 +363,9 @@ public override async Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection { // At runtime with an HTTP request - attempt OBO flow // Note: DatabaseAudience is validated at startup by RuntimeConfigValidator - string? oboToken = await GetOboAccessTokenAsync(userDelegatedAuth.DatabaseAudience!); + string? oboToken = await GetOboAccessTokenAsync( + userDelegatedAuth.DatabaseAudience!, + cancellationToken); if (oboToken is not null) { sqlConn.AccessToken = oboToken; @@ -392,7 +398,7 @@ public override async Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection string? accessToken = accessTokenFromController ?? (IsDefaultAccessTokenValid() ? ((AccessToken)_defaultAccessToken!).Token : - await GetAccessTokenAsync()); + await GetAccessTokenAsync(cancellationToken)); if (accessToken is not null) { @@ -406,7 +412,9 @@ public override async Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection /// /// The target database audience. /// The OBO access token, or null if OBO cannot be performed. - private async Task GetOboAccessTokenAsync(string databaseAudience) + private async Task GetOboAccessTokenAsync( + string databaseAudience, + CancellationToken cancellationToken) { if (_oboTokenProvider is null || HttpContextAccessor?.HttpContext is null) { @@ -429,7 +437,8 @@ public override async Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection return await _oboTokenProvider.GetAccessTokenOnBehalfOfAsync( principal!, incomingJwt, - databaseAudience); + databaseAudience, + cancellationToken); } /// @@ -466,11 +475,14 @@ private bool IsDefaultAccessTokenValid() /// /// The string representation of the access token if found, /// null otherwise. - private async Task GetAccessTokenAsync() + private async Task GetAccessTokenAsync( + CancellationToken cancellationToken) { try { - _defaultAccessToken = await AzureCredential.GetTokenAsync(new TokenRequestContext(new[] { DATABASE_SCOPE })); + _defaultAccessToken = await AzureCredential.GetTokenAsync( + new TokenRequestContext(new[] { DATABASE_SCOPE }), + cancellationToken); } catch (CredentialUnavailableException ex) { diff --git a/src/Core/Resolvers/MySqlQueryExecutor.cs b/src/Core/Resolvers/MySqlQueryExecutor.cs index 3c31a7de60..9ca1000e95 100644 --- a/src/Core/Resolvers/MySqlQueryExecutor.cs +++ b/src/Core/Resolvers/MySqlQueryExecutor.cs @@ -108,8 +108,12 @@ private void ConfigureMySqlQueryExecutor() /// /// The supplied connection to modify for managed identity access. /// Name of datasource for which to set access token. Default dbName taken from config if null - public override async Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection conn, string dataSourceName) + public override async Task SetManagedIdentityAccessTokenIfAnyAsync( + DbConnection conn, + string dataSourceName, + CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); // using default datasource name for first db - maintaining backward compatibility for single db scenario. if (string.IsNullOrEmpty(dataSourceName)) { @@ -128,7 +132,7 @@ public override async Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection string? accessToken = accessTokenFromController ?? (IsDefaultAccessTokenValid() ? ((AccessToken)_defaultAccessToken!).Token : - await GetAccessTokenAsync()); + await GetAccessTokenAsync(cancellationToken)); if (accessToken is not null) { @@ -172,11 +176,14 @@ private bool IsDefaultAccessTokenValid() /// /// The string representation of the access token if found, /// null otherwise. - private async Task GetAccessTokenAsync() + private async Task GetAccessTokenAsync( + CancellationToken cancellationToken) { try { - _defaultAccessToken = await AzureCredential.GetTokenAsync(new TokenRequestContext(new[] { DATABASE_SCOPE })); + _defaultAccessToken = await AzureCredential.GetTokenAsync( + new TokenRequestContext(new[] { DATABASE_SCOPE }), + cancellationToken); } catch (CredentialUnavailableException ex) { diff --git a/src/Core/Resolvers/PostgreSqlExecutor.cs b/src/Core/Resolvers/PostgreSqlExecutor.cs index 4130cd1378..919eaff5e1 100644 --- a/src/Core/Resolvers/PostgreSqlExecutor.cs +++ b/src/Core/Resolvers/PostgreSqlExecutor.cs @@ -104,8 +104,12 @@ private void ConfigurePostgreSqlQueryExecutor() /// /// The supplied connection to modify for managed identity access. /// Name of datasource for which to set access token. Default dbName taken from config if null - public override async Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection conn, string dataSourceName) + public override async Task SetManagedIdentityAccessTokenIfAnyAsync( + DbConnection conn, + string dataSourceName, + CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); // using default datasource name for first db - maintaining backward compatibility for single db scenario. if (string.IsNullOrEmpty(dataSourceName)) { @@ -126,7 +130,7 @@ public override async Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection string? accessToken = accessTokenFromController ?? (IsDefaultAccessTokenValid() ? ((AccessToken)_defaultAccessToken!).Token : - await GetAccessTokenAsync(dataSourceName)); + await GetAccessTokenAsync(dataSourceName, cancellationToken)); if (accessToken is not null) { @@ -246,7 +250,9 @@ private bool IsDefaultAccessTokenValid() /// /// The string representation of the access token if found, /// null otherwise. - private async Task GetAccessTokenAsync(string dataSourceName) + private async Task GetAccessTokenAsync( + string dataSourceName, + CancellationToken cancellationToken) { bool firstAttemptAtDefaultAccessToken = _defaultAccessToken is null; @@ -254,7 +260,12 @@ private bool IsDefaultAccessTokenValid() { _defaultAccessToken = await AzureCredential.GetTokenAsync( - new TokenRequestContext(new[] { DATABASE_SCOPE })); + new TokenRequestContext(new[] { DATABASE_SCOPE }), + cancellationToken); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + throw; } // because there can be scenarios where password is not specified but // default managed identity is not the intended method of authentication diff --git a/src/Core/Resolvers/QueryExecutor.cs b/src/Core/Resolvers/QueryExecutor.cs index 98917ed2c9..c9b9cb63b5 100644 --- a/src/Core/Resolvers/QueryExecutor.cs +++ b/src/Core/Resolvers/QueryExecutor.cs @@ -172,6 +172,59 @@ public QueryExecutor(DbExceptionParser dbExceptionParser, HttpContext? httpContext = null, List? args = null) { + return await ExecuteQueryAsyncCore( + sqltext, + parameters, + dataReaderHandler, + dataSourceName, + CancellationToken.None, + httpContext, + args); + } + + /// + public async Task ExecuteQueryAsync( + string sqltext, + IDictionary parameters, + Func?, Task>? dataReaderHandler, + string dataSourceName, + CancellationToken cancellationToken, + HttpContext? httpContext = null, + List? args = null) + { + return await ExecuteQueryAsyncCore( + sqltext, + parameters, + dataReaderHandler, + dataSourceName, + cancellationToken, + httpContext, + args); + } + + private async Task ExecuteQueryAsyncCore( + string sqltext, + IDictionary parameters, + Func?, Task>? dataReaderHandler, + string dataSourceName, + CancellationToken cancellationToken, + HttpContext? httpContext, + List? args) + { + CancellationToken requestAborted = + httpContext?.RequestAborted ?? CancellationToken.None; + using CancellationTokenSource? linkedCancellation = + cancellationToken.CanBeCanceled && + requestAborted.CanBeCanceled && + cancellationToken != requestAborted + ? CancellationTokenSource.CreateLinkedTokenSource( + cancellationToken, + requestAborted) + : null; + CancellationToken operationCancellationToken = linkedCancellation?.Token ?? + (cancellationToken.CanBeCanceled ? cancellationToken : requestAborted); + + operationCancellationToken.ThrowIfCancellationRequested(); int retryAttempt = 0; if (string.IsNullOrEmpty(dataSourceName)) @@ -190,12 +243,16 @@ public QueryExecutor(DbExceptionParser dbExceptionParser, DataApiBuilderException.SubStatusCodes.UnexpectedError); } - await SetManagedIdentityAccessTokenIfAnyAsync(conn, dataSourceName); + await SetManagedIdentityAccessTokenIfAnyAsync( + conn, + dataSourceName, + operationCancellationToken); TResult? result = default(TResult); - result = await _retryPolicyAsync.ExecuteAsync(async () => + result = await _retryPolicyAsync.ExecuteAsync(async retryCancellationToken => { + retryCancellationToken.ThrowIfCancellationRequested(); retryAttempt++; try { @@ -206,7 +263,28 @@ public QueryExecutor(DbExceptionParser dbExceptionParser, QueryExecutorLogger.LogDebug("{correlationId} Executing query: {queryText}", correlationId, sqltext); } - TResult? result = await ExecuteQueryAgainstDbAsync(conn, sqltext, parameters, dataReaderHandler, httpContext, dataSourceName, args); + // Preserve virtual dispatch to the established overload for legacy callers + // and test doubles. The token-aware overload is required when the caller + // supplied a token; retryCancellationToken then represents that token linked + // with HttpContext.RequestAborted when both are cancellable. + TResult? result = cancellationToken.CanBeCanceled + ? await ExecuteQueryAgainstDbAsync( + conn, + sqltext, + parameters, + dataReaderHandler, + httpContext, + dataSourceName, + args, + retryCancellationToken) + : await ExecuteQueryAgainstDbAsync( + conn, + sqltext, + parameters, + dataReaderHandler, + httpContext, + dataSourceName, + args); if (retryAttempt > 1) { @@ -236,7 +314,7 @@ public QueryExecutor(DbExceptionParser dbExceptionParser, throw DbExceptionParser.Parse(e); } } - }); + }, operationCancellationToken); return result; } @@ -284,19 +362,60 @@ public virtual TConnection CreateConnection(string dataSourceName) HttpContext? httpContext, string dataSourceName, List? args = null) + { + return await ExecuteQueryAgainstDbAsyncCore( + conn, + sqltext, + parameters, + dataReaderHandler, + httpContext, + dataSourceName, + args, + httpContext?.RequestAborted ?? CancellationToken.None); + } + + public virtual async Task ExecuteQueryAgainstDbAsync( + TConnection conn, + string sqltext, + IDictionary parameters, + Func?, Task>? dataReaderHandler, + HttpContext? httpContext, + string dataSourceName, + List? args, + CancellationToken cancellationToken) + { + return await ExecuteQueryAgainstDbAsyncCore( + conn, + sqltext, + parameters, + dataReaderHandler, + httpContext, + dataSourceName, + args, + cancellationToken); + } + + private async Task ExecuteQueryAgainstDbAsyncCore( + TConnection conn, + string sqltext, + IDictionary parameters, + Func?, Task>? dataReaderHandler, + HttpContext? httpContext, + string dataSourceName, + List? args, + CancellationToken cancellationToken) { Stopwatch queryExecutionTimer = new(); queryExecutionTimer.Start(); try { - await conn.OpenAsync(); + await conn.OpenAsync(cancellationToken); DbCommand cmd = PrepareDbCommand(conn, sqltext, parameters, httpContext, dataSourceName); TResult? result = default(TResult); try { CommandBehavior commandBehavior = ConfigProvider.GetConfig().MaxResponseSizeLogicEnabled() ? CommandBehavior.SequentialAccess : CommandBehavior.CloseConnection; // CancellationToken is passed to ExecuteReaderAsync to ensure that if the client times out while the query is executing, the execution will be cancelled and resources will be freed up. - CancellationToken cancellationToken = httpContext?.RequestAborted ?? CancellationToken.None; using DbDataReader dbDataReader = await cmd.ExecuteReaderAsync(commandBehavior, cancellationToken); if (dataReaderHandler is not null && dbDataReader is not null) @@ -429,8 +548,23 @@ public virtual void PopulateDbTypeForParameter(KeyValuePair - public virtual async Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection conn, string dataSourceName = "") + public virtual async Task SetManagedIdentityAccessTokenIfAnyAsync( + DbConnection conn, + string dataSourceName = "") + { + await SetManagedIdentityAccessTokenIfAnyAsync( + conn, + dataSourceName, + CancellationToken.None); + } + + /// + public virtual async Task SetManagedIdentityAccessTokenIfAnyAsync( + DbConnection conn, + string dataSourceName, + CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); // no-op in the base class. await Task.Yield(); } diff --git a/src/Core/Services/GraphQLSchemaCreator.cs b/src/Core/Services/GraphQLSchemaCreator.cs index 7c9b505e96..5bad1dbb01 100644 --- a/src/Core/Services/GraphQLSchemaCreator.cs +++ b/src/Core/Services/GraphQLSchemaCreator.cs @@ -85,6 +85,7 @@ public GraphQLSchemaCreator( /// protected void OnConfigChanged(object? sender, HotReloadEventArgs args) { + args.CancellationToken.ThrowIfCancellationRequested(); RuntimeConfig runtimeConfig = _runtimeConfigProvider.GetConfig(); _isMultipleCreateOperationEnabled = runtimeConfig.IsMultipleCreateOperationEnabled(); _isAggregationEnabled = runtimeConfig.EnableAggregation; diff --git a/src/Core/Services/MetadataProviders/CosmosSqlMetadataProvider.cs b/src/Core/Services/MetadataProviders/CosmosSqlMetadataProvider.cs index b41a1752b8..fcfd0a0caa 100644 --- a/src/Core/Services/MetadataProviders/CosmosSqlMetadataProvider.cs +++ b/src/Core/Services/MetadataProviders/CosmosSqlMetadataProvider.cs @@ -391,6 +391,12 @@ public Task InitializeAsync() return Task.CompletedTask; } + public Task InitializeAsync(CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.CompletedTask; + } + private string GraphQLSchema() { if (_cosmosDb.GraphQLSchema is not null) diff --git a/src/Core/Services/MetadataProviders/IMetadataProviderFactory.cs b/src/Core/Services/MetadataProviders/IMetadataProviderFactory.cs index 86fa6df4fd..dc0cf65f9b 100644 --- a/src/Core/Services/MetadataProviders/IMetadataProviderFactory.cs +++ b/src/Core/Services/MetadataProviders/IMetadataProviderFactory.cs @@ -28,6 +28,16 @@ public interface IMetadataProviderFactory /// public Task InitializeAsync(); + /// + /// Initializes the metadata providers with cooperative cancellation. Implementations that + /// do not override this member retain their existing initialization behavior. + /// + public Task InitializeAsync(CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return InitializeAsync(); + } + /// /// Initializes the metadata providers with parameters /// Note : this is used in GraphQL workload to call the parameterized initialize async method in providers diff --git a/src/Core/Services/MetadataProviders/ISqlMetadataProvider.cs b/src/Core/Services/MetadataProviders/ISqlMetadataProvider.cs index 892cd89013..524596fd72 100644 --- a/src/Core/Services/MetadataProviders/ISqlMetadataProvider.cs +++ b/src/Core/Services/MetadataProviders/ISqlMetadataProvider.cs @@ -22,6 +22,17 @@ public interface ISqlMetadataProvider /// Task InitializeAsync(); + /// + /// Initializes this metadata provider for the runtime with cooperative cancellation. + /// Implementations that do not override this member retain their existing initialization + /// behavior. + /// + Task InitializeAsync(CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return InitializeAsync(); + } + /// /// Obtains the underlying source object's schema name (SQL) or container name (Cosmos). /// diff --git a/src/Core/Services/MetadataProviders/MetadataProviderFactory.cs b/src/Core/Services/MetadataProviders/MetadataProviderFactory.cs index 6fe20969ed..635928ccc8 100644 --- a/src/Core/Services/MetadataProviders/MetadataProviderFactory.cs +++ b/src/Core/Services/MetadataProviders/MetadataProviderFactory.cs @@ -65,10 +65,11 @@ private void ConfigureMetadataProviders() public void OnConfigChanged(object? sender, HotReloadEventArgs args) { + args.CancellationToken.ThrowIfCancellationRequested(); _metadataProviders.Clear(); ConfigureMetadataProviders(); // Blocks the current thread until initialization is finished. - this.InitializeAsync().GetAwaiter().GetResult(); + this.InitializeAsync(args.CancellationToken).GetAwaiter().GetResult(); } /// @@ -87,12 +88,19 @@ public ISqlMetadataProvider GetMetadataProvider(string dataSourceName) /// public async Task InitializeAsync() + { + await InitializeAsync(CancellationToken.None); + } + + /// + public async Task InitializeAsync(CancellationToken cancellationToken) { foreach ((_, ISqlMetadataProvider provider) in _metadataProviders) { + cancellationToken.ThrowIfCancellationRequested(); if (provider is not null) { - await provider.InitializeAsync(); + await provider.InitializeAsync(cancellationToken); } } } diff --git a/src/Core/Services/MetadataProviders/MsSqlMetadataProvider.cs b/src/Core/Services/MetadataProviders/MsSqlMetadataProvider.cs index 20de74a96d..3162f81428 100644 --- a/src/Core/Services/MetadataProviders/MsSqlMetadataProvider.cs +++ b/src/Core/Services/MetadataProviders/MsSqlMetadataProvider.cs @@ -60,7 +60,27 @@ public override Type SqlToCLRType(string sqlType) } /// - public override async Task PopulateTriggerMetadataForTable(string entityName, string schemaName, string tableName, SourceDefinition sourceDefinition) + public override Task PopulateTriggerMetadataForTable( + string entityName, + string schemaName, + string tableName, + SourceDefinition sourceDefinition) + { + return PopulateTriggerMetadataForTable( + entityName, + schemaName, + tableName, + sourceDefinition, + CancellationToken.None); + } + + /// + public override async Task PopulateTriggerMetadataForTable( + string entityName, + string schemaName, + string tableName, + SourceDefinition sourceDefinition, + CancellationToken cancellationToken) { string enumerateEnabledTriggers = SqlQueryBuilder.BuildFetchEnabledTriggersQuery(); Dictionary parameters = new() @@ -73,7 +93,8 @@ public override async Task PopulateTriggerMetadataForTable(string entityName, st sqltext: enumerateEnabledTriggers, parameters: parameters, dataReaderHandler: QueryExecutor.GetJsonArrayAsync, - dataSourceName: _dataSourceName); + dataSourceName: _dataSourceName, + cancellationToken: cancellationToken); using JsonDocument sqlResult = JsonDocument.Parse(resultArray!.ToJsonString()); foreach (JsonElement element in sqlResult.RootElement.EnumerateArray()) @@ -158,12 +179,16 @@ protected override async Task FillSchemaForStoredProcedureAsync( string entityName, string schemaName, string storedProcedureSourceName, - StoredProcedureDefinition storedProcedureDefinition) + StoredProcedureDefinition storedProcedureDefinition, + CancellationToken cancellationToken) { using DbConnection conn = new SqlConnection(); conn.ConnectionString = ConnectionString; - await QueryExecutor.SetManagedIdentityAccessTokenIfAnyAsync(conn, _dataSourceName); - await conn.OpenAsync(); + await QueryExecutor.SetManagedIdentityAccessTokenIfAnyAsync( + conn, + _dataSourceName, + cancellationToken); + await conn.OpenAsync(cancellationToken); string[] procedureRestrictions = new string[NUMBER_OF_RESTRICTIONS]; @@ -172,7 +197,10 @@ protected override async Task FillSchemaForStoredProcedureAsync( procedureRestrictions[1] = schemaName; procedureRestrictions[2] = storedProcedureSourceName; - DataTable procedureMetadata = await conn.GetSchemaAsync(collectionName: "Procedures", restrictionValues: procedureRestrictions); + DataTable procedureMetadata = await conn.GetSchemaAsync( + collectionName: "Procedures", + restrictionValues: procedureRestrictions, + cancellationToken: cancellationToken); // Stored procedure does not exist in DB schema if (procedureMetadata.Rows.Count == 0) @@ -184,7 +212,10 @@ protected override async Task FillSchemaForStoredProcedureAsync( } // Each row in the procedureParams DataTable corresponds to a single parameter - DataTable parameterMetadata = await conn.GetSchemaAsync(collectionName: "ProcedureParameters", restrictionValues: procedureRestrictions); + DataTable parameterMetadata = await conn.GetSchemaAsync( + collectionName: "ProcedureParameters", + restrictionValues: procedureRestrictions, + cancellationToken: cancellationToken); // For each row/parameter, add an entry to StoredProcedureDefinition.Parameters dictionary foreach (DataRow row in parameterMetadata.Rows) @@ -309,7 +340,18 @@ private bool TryResolveDbType(string sqlDbTypeName, out DbType dbType) } /// - protected override async Task GenerateAutoentitiesIntoEntities(IReadOnlyDictionary? autoentities) + protected override Task GenerateAutoentitiesIntoEntities( + IReadOnlyDictionary? autoentities) + { + return GenerateAutoentitiesIntoEntities( + autoentities, + CancellationToken.None); + } + + /// + protected override async Task GenerateAutoentitiesIntoEntities( + IReadOnlyDictionary? autoentities, + CancellationToken cancellationToken) { if (autoentities is null) { @@ -321,8 +363,12 @@ protected override async Task GenerateAutoentitiesIntoEntities(IReadOnlyDictiona Dictionary entityNameToRawEntity = new(); foreach ((string autoentityName, Autoentity autoentity) in autoentities) { + cancellationToken.ThrowIfCancellationRequested(); int addedEntities = 0; - JsonArray? resultArray = await QueryAutoentitiesAsync(autoentityName, autoentity); + JsonArray? resultArray = await QueryAutoentitiesAsync( + autoentityName, + autoentity, + cancellationToken); if (resultArray is null) { continue; @@ -430,7 +476,23 @@ protected override async Task GenerateAutoentitiesIntoEntities(IReadOnlyDictiona /// The name of the autoentity definition. /// The autoentity definition containing patterns for inclusion, exclusion, and name. /// A JsonArray containing the queried autoentities, or an empty array if none are found. - public async Task QueryAutoentitiesAsync(string autoentityName, Autoentity autoentity) + public Task QueryAutoentitiesAsync( + string autoentityName, + Autoentity autoentity) + { + return QueryAutoentitiesAsync( + autoentityName, + autoentity, + CancellationToken.None); + } + + /// + /// Queries the database for autoentities with cooperative cancellation. + /// + public async Task QueryAutoentitiesAsync( + string autoentityName, + Autoentity autoentity, + CancellationToken cancellationToken) { string include = string.Join(",", autoentity.Patterns.Include); string exclude = string.Join(",", autoentity.Patterns.Exclude); @@ -452,7 +514,8 @@ protected override async Task GenerateAutoentitiesIntoEntities(IReadOnlyDictiona sqltext: getAutoentitiesQuery, parameters: parameters, dataReaderHandler: QueryExecutor.GetJsonArrayAsync, - dataSourceName: _dataSourceName); + dataSourceName: _dataSourceName, + cancellationToken: cancellationToken); return resultArray; } diff --git a/src/Core/Services/MetadataProviders/MySqlMetadataProvider.cs b/src/Core/Services/MetadataProviders/MySqlMetadataProvider.cs index 26098c2d15..c0ede2c347 100644 --- a/src/Core/Services/MetadataProviders/MySqlMetadataProvider.cs +++ b/src/Core/Services/MetadataProviders/MySqlMetadataProvider.cs @@ -46,16 +46,22 @@ public MySqlMetadataProvider( /// support 3 level naming of tables. protected override async Task GetColumnsAsync( string schemaName, - string tableName) + string tableName, + CancellationToken cancellationToken) { using MySqlConnection conn = new(ConnectionString); - await QueryExecutor.SetManagedIdentityAccessTokenIfAnyAsync(conn, _dataSourceName); - await conn.OpenAsync(); + await QueryExecutor.SetManagedIdentityAccessTokenIfAnyAsync( + conn, + _dataSourceName, + cancellationToken); + await conn.OpenAsync(cancellationToken); // Each row in the allColumns table corresponds to a single column. // Since column restrictions are ignored, this retrieves all the columns // in the engine irrespective of database and table name. - DataTable allColumns = await conn.GetSchemaAsync("Columns"); + DataTable allColumns = await conn.GetSchemaAsync( + "Columns", + cancellationToken); // Manually filter here to find out which columns need to be removed // by checking the database name and table name. diff --git a/src/Core/Services/MetadataProviders/SqlMetadataProvider.cs b/src/Core/Services/MetadataProviders/SqlMetadataProvider.cs index 9517c41781..d7999deb24 100644 --- a/src/Core/Services/MetadataProviders/SqlMetadataProvider.cs +++ b/src/Core/Services/MetadataProviders/SqlMetadataProvider.cs @@ -302,6 +302,13 @@ public string GetEntityName(string graphQLType) /// public async Task InitializeAsync() { + await InitializeAsync(CancellationToken.None); + } + + /// + public async Task InitializeAsync(CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); System.Diagnostics.Stopwatch timer = System.Diagnostics.Stopwatch.StartNew(); if (_isValidateOnly) @@ -311,7 +318,11 @@ public async Task InitializeAsync() // To enable to check for multiple data-sources just remove this validation and each entity will have its own connection check. try { - await ValidateDatabaseConnection(); + await ValidateDatabaseConnection(cancellationToken); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + throw; } catch (Exception e) { @@ -326,16 +337,18 @@ public async Task InitializeAsync() if (GetDatabaseType() == DatabaseType.MSSQL) { - await GenerateAutoentitiesIntoEntities(Autoentities); + await GenerateAutoentitiesIntoEntities(Autoentities, cancellationToken); } + cancellationToken.ThrowIfCancellationRequested(); // Running these entity validations only in development mode to ensure // fast startup of engine in production mode. RuntimeConfig runtimeConfig = _runtimeConfigProvider.GetConfig(); _runtimeConfigValidator.ValidateEntityAndAutoentityConfigurations(runtimeConfig); GenerateDatabaseObjectForEntities(); - await PopulateObjectDefinitionForEntities(); + await PopulateObjectDefinitionForEntities(cancellationToken); + cancellationToken.ThrowIfCancellationRequested(); GenerateExposedToBackingColumnMapsForEntities(); // When IsLateConfigured is true we are in a hosted scenario and do not reveal primary key information. @@ -460,12 +473,20 @@ private void LogPrimaryKeys() /// /// Verify that the stored procedure exists in the database schema, then populate its database object parameters accordingly /// + /// + /// The cancellation-token signature intentionally replaces the former tokenless protected + /// virtual slot. This method owns cancellable database schema I/O, and custom metadata + /// provider subclassing is not a documented provider plug-in contract. Direct subclasses + /// must update their override and propagate . Retaining + /// tokenless dispatch here would allow schema I/O to escape coordinated shutdown cancellation. + /// protected virtual async Task FillSchemaForStoredProcedureAsync( Entity procedureEntity, string entityName, string schemaName, string storedProcedureSourceName, - StoredProcedureDefinition storedProcedureDefinition) + StoredProcedureDefinition storedProcedureDefinition, + CancellationToken cancellationToken) { using ConnectionT conn = new(); conn.ConnectionString = ConnectionString; @@ -474,15 +495,25 @@ protected virtual async Task FillSchemaForStoredProcedureAsync( try { - await QueryExecutor.SetManagedIdentityAccessTokenIfAnyAsync(conn, _dataSourceName); - await conn.OpenAsync(); + await QueryExecutor.SetManagedIdentityAccessTokenIfAnyAsync( + conn, + _dataSourceName, + cancellationToken); + await conn.OpenAsync(cancellationToken); // To restrict the parameters for the current stored procedure, specify its name procedureRestrictions[0] = conn.Database; procedureRestrictions[1] = schemaName; procedureRestrictions[2] = storedProcedureSourceName; - procedureMetadata = await conn.GetSchemaAsync(collectionName: "Procedures", restrictionValues: procedureRestrictions); + procedureMetadata = await conn.GetSchemaAsync( + collectionName: "Procedures", + restrictionValues: procedureRestrictions, + cancellationToken); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + throw; } catch (Exception ex) { @@ -507,7 +538,10 @@ protected virtual async Task FillSchemaForStoredProcedureAsync( } // Each row in the procedureParams DataTable corresponds to a single parameter - DataTable parameterMetadata = await conn.GetSchemaAsync(collectionName: "ProcedureParameters", restrictionValues: procedureRestrictions); + DataTable parameterMetadata = await conn.GetSchemaAsync( + collectionName: "ProcedureParameters", + restrictionValues: procedureRestrictions, + cancellationToken); // For each row/parameter, add an entry to StoredProcedureDefinition.Parameters dictionary foreach (DataRow row in parameterMetadata.Rows) @@ -571,11 +605,34 @@ protected virtual async Task FillSchemaForStoredProcedureAsync( /// Name of the schema in which the table is present. /// Name of the table. /// Table definition to update. - public virtual Task PopulateTriggerMetadataForTable(string entityName, string schemaName, string tableName, SourceDefinition sourceDefinition) + public virtual Task PopulateTriggerMetadataForTable( + string entityName, + string schemaName, + string tableName, + SourceDefinition sourceDefinition) { throw new NotImplementedException(); } + /// + /// Updates trigger metadata with cooperative cancellation. Derived implementations that + /// only override the established member retain their existing behavior. + /// + public virtual Task PopulateTriggerMetadataForTable( + string entityName, + string schemaName, + string tableName, + SourceDefinition sourceDefinition, + CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return PopulateTriggerMetadataForTable( + entityName, + schemaName, + tableName, + sourceDefinition); + } + /// /// Generates the map used to find a given entity based /// on the path that will be used for that entity. @@ -727,11 +784,24 @@ private void GenerateDatabaseObjectForEntities() /// Creates entities for each table that is found, based on the autoentity configuration. /// This method is only called for tables in MsSql. /// - protected virtual Task GenerateAutoentitiesIntoEntities(IReadOnlyDictionary? autoentities) + protected virtual Task GenerateAutoentitiesIntoEntities( + IReadOnlyDictionary? autoentities) { throw new NotSupportedException($"{GetType().Name} does not support autoentities yet."); } + /// + /// Creates autoentities with cooperative cancellation. Derived implementations that only + /// override the established member retain their existing behavior. + /// + protected virtual Task GenerateAutoentitiesIntoEntities( + IReadOnlyDictionary? autoentities, + CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return GenerateAutoentitiesIntoEntities(autoentities); + } + /// /// Removes the entities that were generated from the autoentities property. /// This should only be done when we only want to validate the entities. @@ -1196,21 +1266,34 @@ public IReadOnlyDictionary GetLinkingEntities() /// Populates table definition for entities specified as tables or views /// Populates procedure definition for entities specified as stored procedures /// - private async Task PopulateObjectDefinitionForEntities() + private async Task PopulateObjectDefinitionForEntities( + CancellationToken cancellationToken) { foreach ((string entityName, Entity entity) in Entities) { - await PopulateObjectDefinitionForEntity(entityName, entity); + cancellationToken.ThrowIfCancellationRequested(); + await PopulateObjectDefinitionForEntity( + entityName, + entity, + cancellationToken); } foreach ((string entityName, Entity entity) in _linkingEntities) { - await PopulateObjectDefinitionForEntity(entityName, entity); + cancellationToken.ThrowIfCancellationRequested(); + await PopulateObjectDefinitionForEntity( + entityName, + entity, + cancellationToken); } try { - await PopulateForeignKeyDefinitionAsync(); + await PopulateForeignKeyDefinitionAsync(cancellationToken); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + throw; } catch (Exception e) { @@ -1218,10 +1301,14 @@ private async Task PopulateObjectDefinitionForEntities() } } - private async Task PopulateObjectDefinitionForEntity(string entityName, Entity entity) + private async Task PopulateObjectDefinitionForEntity( + string entityName, + Entity entity, + CancellationToken cancellationToken) { try { + cancellationToken.ThrowIfCancellationRequested(); EntitySourceType entitySourceType = GetEntitySourceType(entityName, entity); if (entitySourceType is EntitySourceType.StoredProcedure) { @@ -1230,14 +1317,16 @@ await FillSchemaForStoredProcedureAsync( entityName, GetSchemaName(entityName), GetDatabaseObjectName(entityName), - GetStoredProcedureDefinition(entityName)); + GetStoredProcedureDefinition(entityName), + cancellationToken); if (GetDatabaseType() == DatabaseType.MSSQL || GetDatabaseType() == DatabaseType.DWSQL) { await PopulateResultSetDefinitionsForStoredProcedureAsync( GetSchemaName(entityName), GetDatabaseObjectName(entityName), - GetStoredProcedureDefinition(entityName)); + GetStoredProcedureDefinition(entityName), + cancellationToken); } } else if (entitySourceType is EntitySourceType.Table) @@ -1265,7 +1354,8 @@ await PopulateResultSetDefinitionsForStoredProcedureAsync( DataTable dataTable = await GetTableWithSchemaFromDataSetAsync( entityName, GetSchemaName(entityName), - GetDatabaseObjectName(entityName)); + GetDatabaseObjectName(entityName), + cancellationToken); pkFields = dataTable.PrimaryKey.Select(pk => pk.ColumnName).ToList(); } @@ -1278,7 +1368,8 @@ await PopulateSourceDefinitionAsync( GetSchemaName(entityName), GetDatabaseObjectName(entityName), GetSourceDefinition(entityName), - pkFields); + pkFields, + cancellationToken); } else { @@ -1305,7 +1396,8 @@ await PopulateSourceDefinitionAsync( DataTable dataTable = await GetTableWithSchemaFromDataSetAsync( entityName, GetSchemaName(entityName), - GetDatabaseObjectName(entityName)); + GetDatabaseObjectName(entityName), + cancellationToken); pkFields = dataTable.PrimaryKey.Select(pk => pk.ColumnName).ToList(); } @@ -1316,9 +1408,14 @@ await PopulateSourceDefinitionAsync( GetSchemaName(entityName), GetDatabaseObjectName(entityName), viewDefinition, - pkFields); + pkFields, + cancellationToken); } } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + throw; + } catch (Exception e) { HandleOrRecordException(e); @@ -1332,7 +1429,8 @@ await PopulateSourceDefinitionAsync( private async Task PopulateResultSetDefinitionsForStoredProcedureAsync( string schemaName, string storedProcedureName, - SourceDefinition sourceDefinition) + SourceDefinition sourceDefinition, + CancellationToken cancellationToken) { StoredProcedureDefinition storedProcedureDefinition = (StoredProcedureDefinition)sourceDefinition; string dbStoredProcedureName = $"{schemaName}.{storedProcedureName}"; @@ -1346,7 +1444,8 @@ private async Task PopulateResultSetDefinitionsForStoredProcedureAsync( sqltext: queryForResultSetDetails, parameters: null!, dataReaderHandler: QueryExecutor.GetJsonArrayAsync, - dataSourceName: _dataSourceName); + dataSourceName: _dataSourceName, + cancellationToken: cancellationToken); using JsonDocument sqlResult = JsonDocument.Parse(resultArray!.ToJsonString()); @@ -1506,8 +1605,10 @@ private async Task PopulateSourceDefinitionAsync( string schemaName, string tableName, SourceDefinition sourceDefinition, - List pkFields) + List pkFields, + CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); sourceDefinition.PrimaryKey = [.. pkFields]; if (sourceDefinition.PrimaryKey.Count == 0) @@ -1521,10 +1622,19 @@ private async Task PopulateSourceDefinitionAsync( Entities.TryGetValue(entityName, out Entity? entity); if (GetDatabaseType() is DatabaseType.MSSQL && entity is not null && entity.Source.Type is EntitySourceType.Table) { - await PopulateTriggerMetadataForTable(entityName, schemaName, tableName, sourceDefinition); + await PopulateTriggerMetadataForTable( + entityName, + schemaName, + tableName, + sourceDefinition, + cancellationToken); } - DataTable dataTable = await GetTableWithSchemaFromDataSetAsync(entityName, schemaName, tableName); + DataTable dataTable = await GetTableWithSchemaFromDataSetAsync( + entityName, + schemaName, + tableName, + cancellationToken); using DataTableReader reader = new(dataTable); DataTable schemaTable = reader.GetSchemaTable(); RuntimeConfig runtimeConfig = _runtimeConfigProvider.GetConfig(); @@ -1569,7 +1679,10 @@ private async Task PopulateSourceDefinitionAsync( sourceDefinition.Columns.TryAdd(columnName, column); } - DataTable columnsInTable = await GetColumnsAsync(schemaName, tableName); + DataTable columnsInTable = await GetColumnsAsync( + schemaName, + tableName, + cancellationToken); PopulateColumnDefinitionWithHasDefaultAndDbType( sourceDefinition, @@ -1579,7 +1692,11 @@ private async Task PopulateSourceDefinitionAsync( { // For MySql, database name is equivalent to schema name. string schemaOrDatabaseName = GetDatabaseType() is DatabaseType.MySQL ? GetDatabaseName() : schemaName; - await PopulateColumnDefinitionsWithReadOnlyFlag(tableName, schemaOrDatabaseName, sourceDefinition); + await PopulateColumnDefinitionsWithReadOnlyFlag( + tableName, + schemaOrDatabaseName, + sourceDefinition, + cancellationToken); } } @@ -1590,7 +1707,11 @@ private async Task PopulateSourceDefinitionAsync( /// Name of the table. /// Name of the schema (for MsSql/PgSql)/database (for MySql) of the table. /// Table definition. - private async Task PopulateColumnDefinitionsWithReadOnlyFlag(string tableName, string schemaOrDatabaseName, SourceDefinition sourceDefinition) + private async Task PopulateColumnDefinitionsWithReadOnlyFlag( + string tableName, + string schemaOrDatabaseName, + SourceDefinition sourceDefinition, + CancellationToken cancellationToken) { string schemaOrDatabaseParamName = $"{BaseQueryStructure.PARAM_NAME_PREFIX}param0"; string quotedTableName = SqlQueryBuilder.QuoteTableNameAsDBConnectionParam(tableName); @@ -1606,7 +1727,8 @@ private async Task PopulateColumnDefinitionsWithReadOnlyFlag(string tableName, s sqltext: queryToGetReadOnlyColumns, parameters: parameters, dataReaderHandler: SummarizeReadOnlyFieldsMetadata, - dataSourceName: _dataSourceName); + dataSourceName: _dataSourceName, + cancellationToken: cancellationToken); if (readOnlyFields is not null && readOnlyFields.Count > 0) { @@ -1671,7 +1793,8 @@ public static bool IsGraphQLReservedName(Entity entity, string databaseColumnNam private async Task GetTableWithSchemaFromDataSetAsync( string entityName, string schemaName, - string tableName) + string tableName, + CancellationToken cancellationToken) { // Because we have an instance of SqlMetadataProvider for each individual database // (note: this means each actual database not each database type), we do not @@ -1685,7 +1808,14 @@ private async Task GetTableWithSchemaFromDataSetAsync( { try { - dataTable = await FillSchemaForTableAsync(schemaName, tableName); + dataTable = await FillSchemaForTableAsync( + schemaName, + tableName, + cancellationToken); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + throw; } catch (Exception ex) when (ex is not DataApiBuilderException) { @@ -1729,14 +1859,22 @@ private async Task GetTableWithSchemaFromDataSetAsync( /// It is specifically used to validate the connection string provided in the runtime configuration /// for single datasource. /// - private async Task ValidateDatabaseConnection() + private async Task ValidateDatabaseConnection( + CancellationToken cancellationToken) { using ConnectionT conn = new(); conn.ConnectionString = ConnectionString; - await QueryExecutor.SetManagedIdentityAccessTokenIfAnyAsync(conn, _dataSourceName); + await QueryExecutor.SetManagedIdentityAccessTokenIfAnyAsync( + conn, + _dataSourceName, + cancellationToken); try { - await conn.OpenAsync(); + await conn.OpenAsync(cancellationToken); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + throw; } catch (Exception ex) { @@ -1756,7 +1894,8 @@ private async Task ValidateDatabaseConnection() /// private async Task FillSchemaForTableAsync( string schemaName, - string tableName) + string tableName, + CancellationToken cancellationToken) { using ConnectionT conn = new(); // If connection string is set to empty string @@ -1780,7 +1919,14 @@ private async Task FillSchemaForTableAsync( // for non-MySql DB types, this will throw an exception // for malformed connection strings conn.ConnectionString = ConnectionString; - await QueryExecutor.SetManagedIdentityAccessTokenIfAnyAsync(conn, _dataSourceName); + await QueryExecutor.SetManagedIdentityAccessTokenIfAnyAsync( + conn, + _dataSourceName, + cancellationToken); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + throw; } catch (Exception ex) { @@ -1793,10 +1939,10 @@ private async Task FillSchemaForTableAsync( innerException: ex); } - await conn.OpenAsync(); + await conn.OpenAsync(cancellationToken); - DataAdapterT adapterForTable = new(); - CommandT selectCommand = new() + using DataAdapterT adapterForTable = new(); + using CommandT selectCommand = new() { Connection = conn }; @@ -1806,7 +1952,40 @@ private async Task FillSchemaForTableAsync( = $"SELECT * FROM {tableNameWithSchemaPrefix}"; adapterForTable.SelectCommand = selectCommand; - DataTable[] dataTable = adapterForTable.FillSchema(EntitiesDataSet, SchemaType.Source, tableNameWithSchemaPrefix); + cancellationToken.ThrowIfCancellationRequested(); + using CancellationTokenRegistration cancellationRegistration = + cancellationToken.Register( + static commandState => + { + try + { + ((DbCommand)commandState!).Cancel(); + } + catch (Exception) + { + // Cancellation is best effort. The provider operation reports its + // own completion or failure to the thread executing FillSchema. + } + }, + selectCommand); + + DataTable[] dataTable; + try + { + dataTable = adapterForTable.FillSchema( + EntitiesDataSet, + SchemaType.Source, + tableNameWithSchemaPrefix); + } + catch (Exception ex) when (cancellationToken.IsCancellationRequested) + { + throw new OperationCanceledException( + "Schema discovery was canceled during shutdown.", + ex, + cancellationToken); + } + + cancellationToken.ThrowIfCancellationRequested(); return dataTable[0]; } @@ -1843,14 +2022,25 @@ internal string GetTableNameWithSchemaPrefix(string schemaName, string tableName /// /// A data table where each row corresponds to a /// column of the table. + /// + /// The cancellation-token signature intentionally replaces the former tokenless protected + /// virtual slot. This method owns cancellable database schema I/O, and custom metadata + /// provider subclassing is not a documented provider plug-in contract. Direct subclasses + /// must update their override and propagate . Retaining + /// tokenless dispatch here would allow schema I/O to escape coordinated shutdown cancellation. + /// protected virtual async Task GetColumnsAsync( string schemaName, - string tableName) + string tableName, + CancellationToken cancellationToken) { using ConnectionT conn = new(); conn.ConnectionString = ConnectionString; - await QueryExecutor.SetManagedIdentityAccessTokenIfAnyAsync(conn, _dataSourceName); - await conn.OpenAsync(); + await QueryExecutor.SetManagedIdentityAccessTokenIfAnyAsync( + conn, + _dataSourceName, + cancellationToken); + await conn.OpenAsync(cancellationToken); // We can specify the Catalog, Schema, Table Name, Column Name to get // the specified column(s). // Hence, we should create a 4 members array. @@ -1864,7 +2054,10 @@ protected virtual async Task GetColumnsAsync( // Each row in the columnsInTable DataTable corresponds to // a single column of the table. - DataTable columnsInTable = await conn.GetSchemaAsync("Columns", columnRestrictions); + DataTable columnsInTable = await conn.GetSchemaAsync( + "Columns", + columnRestrictions, + cancellationToken); return columnsInTable; } @@ -1899,8 +2092,10 @@ protected virtual void PopulateColumnDefinitionWithHasDefaultAndDbType( /// Fills the table definition with information of the foreign keys /// for all the tables. /// - private async Task PopulateForeignKeyDefinitionAsync() + private async Task PopulateForeignKeyDefinitionAsync( + CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); // For each database object, that has a relationship metadata, // build the array storing all the schemaNames(for now the defaultSchemaName) // and the array for all tableNames @@ -1935,7 +2130,8 @@ private async Task PopulateForeignKeyDefinitionAsync() dataReaderHandler: SummarizeFkMetadata, dataSourceName: _dataSourceName, httpContext: null, - args: null); + args: null, + cancellationToken: cancellationToken); if (PairToFkDefinition is not null) { diff --git a/src/Core/Services/OpenAPI/OpenApiDocumentor.cs b/src/Core/Services/OpenAPI/OpenApiDocumentor.cs index 0b090f4335..fb4c382a71 100644 --- a/src/Core/Services/OpenAPI/OpenApiDocumentor.cs +++ b/src/Core/Services/OpenAPI/OpenApiDocumentor.cs @@ -86,6 +86,7 @@ public OpenApiDocumentor( public void OnConfigChanged(object? sender, HotReloadEventArgs args) { + args.CancellationToken.ThrowIfCancellationRequested(); CreateDocument(doOverrideExistingDocument: true); _roleSpecificDocuments.Clear(); // Clear role-specific document cache on config change } diff --git a/src/Service.Tests/Configuration/HotReload/ConfigurationHotReloadTests.cs b/src/Service.Tests/Configuration/HotReload/ConfigurationHotReloadTests.cs index d85c3ddf01..27940155a7 100644 --- a/src/Service.Tests/Configuration/HotReload/ConfigurationHotReloadTests.cs +++ b/src/Service.Tests/Configuration/HotReload/ConfigurationHotReloadTests.cs @@ -8,6 +8,7 @@ using System.Net.Http.Json; using System.Text.Json; using System.Threading.Tasks; +using Azure.DataApiBuilder.Config; using Azure.DataApiBuilder.Config.ObjectModel; using Azure.DataApiBuilder.Core.Configurations; using Azure.DataApiBuilder.Service.Tests.SqlTests; @@ -27,6 +28,7 @@ public class ConfigurationHotReloadTests private static RuntimeConfigProvider _configProvider; private static StringWriter _writer; private static readonly object _writerLock = new(); + private static HotReloadFailureObserver _hotReloadFailureObserver; private const string CONFIG_FILE_NAME = "hot-reload.dab-config.json"; private const string GQL_QUERY_NAME = "books"; private const string HOT_RELOAD_SUCCESS_MESSAGE = "Validated hot-reloaded configuration file"; @@ -229,6 +231,11 @@ public static async Task ClassInitializeAsync(TestContext context) { Console.WriteLine($"Initializing test server (attempt {attempt}/{maxRetries})..."); _testServer = new(Program.CreateWebHostBuilder(new string[] { "--ConfigFileName", CONFIG_FILE_NAME })); + _hotReloadFailureObserver = new( + _testServer.Services.GetRequiredService>()); + _testServer.Services + .GetRequiredService() + .SetLogger(_hotReloadFailureObserver); _testClient = _testServer.CreateClient(); _configProvider = _testServer.Services.GetService(); @@ -316,6 +323,79 @@ private static bool WriterContains(string message) } } + /// + /// Observes the loader's structured hot-reload failure log. The loader now owns and logs reload + /// failures inside its serialized pipeline, so they no longer escape to ConfigFileWatcher's + /// legacy Console.WriteLine fallback. + /// + private sealed class HotReloadFailureObserver( + ILogger innerLogger) : ILogger + { + private readonly object _syncRoot = new(); + private TaskCompletionSource _failureSource = CreateFailureSource(); + + public IDisposable? BeginScope(TState state) + where TState : notnull => innerLogger.BeginScope(state); + + public bool IsEnabled(LogLevel logLevel) => + logLevel == LogLevel.Error || innerLogger.IsEnabled(logLevel); + + public void Log( + LogLevel logLevel, + EventId eventId, + TState state, + Exception? exception, + Func formatter) + { + innerLogger.Log(logLevel, eventId, state, exception, formatter); + + if (logLevel != LogLevel.Error) + { + return; + } + + string message = formatter(state, exception); + if (!message.Contains( + HOT_RELOAD_FAILURE_MESSAGE, + StringComparison.Ordinal)) + { + return; + } + + RecordFailure(message); + } + + public void Reset() + { + lock (_syncRoot) + { + _failureSource = CreateFailureSource(); + } + } + + public async Task WaitForFailureAsync(TimeSpan timeout) + { + Task failureTask; + lock (_syncRoot) + { + failureTask = _failureSource.Task; + } + + return await failureTask.WaitAsync(timeout); + } + + private static TaskCompletionSource CreateFailureSource() => + new(TaskCreationOptions.RunContinuationsAsynchronously); + + private void RecordFailure(string message) + { + lock (_syncRoot) + { + _failureSource.TrySetResult(message); + } + } + } + /// /// Hot reload the configuration by saving a new file with different rest and graphQL paths. /// Validate that the response is correct when making a request with the newly hot-reloaded paths. @@ -754,18 +834,15 @@ public async Task HotReloadConfigConnectionString() // Act // Hot Reload should fail here + _hotReloadFailureObserver.Reset(); GenerateConfigFile( connectionString: $"WrongConnectionString"); - await WaitForConditionAsync( - () => WriterContains(HOT_RELOAD_FAILURE_MESSAGE), - TimeSpan.FromSeconds(HOT_RELOAD_TIMEOUT_SECONDS), - TimeSpan.FromMilliseconds(500)); + string failedConfigLog = await _hotReloadFailureObserver.WaitForFailureAsync( + TimeSpan.FromSeconds(HOT_RELOAD_TIMEOUT_SECONDS)); // Log that shows that hot-reload was not able to validate properly - string failedConfigLog; lock (_writerLock) { - failedConfigLog = _writer.ToString(); _writer.GetStringBuilder().Clear(); } @@ -851,19 +928,16 @@ public async Task HotReloadConfigDatabaseType() // Act // Hot Reload should fail here + _hotReloadFailureObserver.Reset(); GenerateConfigFile( databaseType: DatabaseType.PostgreSQL, connectionString: $"{ConfigurationTests.GetConnectionStringFromEnvironmentConfig(TestCategory.POSTGRESQL).Replace("\\", "\\\\")}"); - await WaitForConditionAsync( - () => WriterContains(HOT_RELOAD_FAILURE_MESSAGE), - TimeSpan.FromSeconds(HOT_RELOAD_TIMEOUT_SECONDS), - TimeSpan.FromMilliseconds(500)); + string failedConfigLog = await _hotReloadFailureObserver.WaitForFailureAsync( + TimeSpan.FromSeconds(HOT_RELOAD_TIMEOUT_SECONDS)); // Log that shows that hot-reload was not able to validate properly - string failedConfigLog; lock (_writerLock) { - failedConfigLog = _writer.ToString(); _writer.GetStringBuilder().Clear(); } @@ -919,6 +993,7 @@ public async Task HotReloadValidationFail() // Act // Generate a config that will fail validation by disabling REST, GraphQL, and MCP (which is not allowed) + _hotReloadFailureObserver.Reset(); GenerateConfigFile( connectionString: $"{ConfigurationTests.GetConnectionStringFromEnvironmentConfig(TestCategory.MSSQL).Replace("\\", "\\\\")}", restEnabled: "false", @@ -926,10 +1001,8 @@ public async Task HotReloadValidationFail() mcpEnabled: "false"); // Wait for hot-reload to fail - await WaitForConditionAsync( - () => WriterContains(HOT_RELOAD_FAILURE_MESSAGE), - TimeSpan.FromSeconds(HOT_RELOAD_TIMEOUT_SECONDS), - TimeSpan.FromMilliseconds(500)); + await _hotReloadFailureObserver.WaitForFailureAsync( + TimeSpan.FromSeconds(HOT_RELOAD_TIMEOUT_SECONDS)); RuntimeConfig newRuntimeConfig = _configProvider.GetConfig(); @@ -967,16 +1040,15 @@ public async Task HotReloadParsingFail() bool originalGraphQLEnabled = lkgRuntimeConfig.Runtime.GraphQL.Enabled; // Act + _hotReloadFailureObserver.Reset(); GenerateConfigFile( connectionString: $"{ConfigurationTests.GetConnectionStringFromEnvironmentConfig(TestCategory.MSSQL).Replace("\\", "\\\\")}", restEnabled: "invalid", gQLEnabled: "invalid"); // Wait for hot-reload to fail (parsing error should trigger failure message) - await WaitForConditionAsync( - () => WriterContains(HOT_RELOAD_FAILURE_MESSAGE), - TimeSpan.FromSeconds(HOT_RELOAD_TIMEOUT_SECONDS), - TimeSpan.FromMilliseconds(500)); + await _hotReloadFailureObserver.WaitForFailureAsync( + TimeSpan.FromSeconds(HOT_RELOAD_TIMEOUT_SECONDS)); RuntimeConfig newRuntimeConfig = _configProvider.GetConfig(); diff --git a/src/Service.Tests/Mcp/DynamicCustomToolMsSqlIntegrationTests.cs b/src/Service.Tests/Mcp/DynamicCustomToolMsSqlIntegrationTests.cs index d6cf1bf2b6..6ea4bb7f45 100644 --- a/src/Service.Tests/Mcp/DynamicCustomToolMsSqlIntegrationTests.cs +++ b/src/Service.Tests/Mcp/DynamicCustomToolMsSqlIntegrationTests.cs @@ -8,6 +8,7 @@ using System.Threading.Tasks; using Azure.DataApiBuilder.Config.ObjectModel; using Azure.DataApiBuilder.Core.Configurations; +using Azure.DataApiBuilder.Core.Services.MetadataProviders; using Azure.DataApiBuilder.Mcp.Core; using Microsoft.Extensions.DependencyInjection; using Microsoft.VisualStudio.TestTools.UnitTesting; @@ -121,7 +122,9 @@ public void InitializeMetadata_SchemaReflectsDbParameterTypes(string entityName, Entity entity = configProvider.GetConfig().Entities[entityName]; DynamicCustomTool tool = new(entityName, entity); - tool.InitializeMetadata(serviceProvider); + tool.InitializeMetadata( + configProvider.GetConfig(), + serviceProvider.GetRequiredService()); JsonElement properties = tool.GetToolMetadata().InputSchema.GetProperty("properties"); @@ -142,7 +145,9 @@ public void InitializeMetadata_ZeroParamSP_HasEmptyProperties() Entity entity = configProvider.GetConfig().Entities["GetBooks"]; DynamicCustomTool tool = new("GetBooks", entity); - tool.InitializeMetadata(serviceProvider); + tool.InitializeMetadata( + configProvider.GetConfig(), + serviceProvider.GetRequiredService()); JsonElement properties = tool.GetToolMetadata().InputSchema.GetProperty("properties"); @@ -168,7 +173,9 @@ public void InitializeMetadata_DescriptionIncludesConfigDefaults(string entityNa Entity entity = configProvider.GetConfig().Entities[entityName]; DynamicCustomTool tool = new(entityName, entity); - tool.InitializeMetadata(serviceProvider); + tool.InitializeMetadata( + configProvider.GetConfig(), + serviceProvider.GetRequiredService()); JsonElement properties = tool.GetToolMetadata().InputSchema.GetProperty("properties"); string description = properties.GetProperty(paramName).GetProperty("description").GetString()!; @@ -188,7 +195,9 @@ public void InitializeMetadata_RequiredArray_IncludesParamWithoutDefault() Entity entity = configProvider.GetConfig().Entities["GetBook"]; DynamicCustomTool tool = new("GetBook", entity); - tool.InitializeMetadata(serviceProvider); + tool.InitializeMetadata( + configProvider.GetConfig(), + serviceProvider.GetRequiredService()); JsonElement schema = tool.GetToolMetadata().InputSchema; @@ -216,7 +225,9 @@ public void InitializeMetadata_RequiredArray_ExcludesParamsWithConfigDefaults() Entity entity = configProvider.GetConfig().Entities["InsertBook"]; DynamicCustomTool tool = new("InsertBook", entity); - tool.InitializeMetadata(serviceProvider); + tool.InitializeMetadata( + configProvider.GetConfig(), + serviceProvider.GetRequiredService()); JsonElement schema = tool.GetToolMetadata().InputSchema; @@ -246,7 +257,9 @@ public void InitializeMetadata_ZeroParamSP_OmitsRequiredArray() Entity entity = configProvider.GetConfig().Entities["GetBooks"]; DynamicCustomTool tool = new("GetBooks", entity); - tool.InitializeMetadata(serviceProvider); + tool.InitializeMetadata( + configProvider.GetConfig(), + serviceProvider.GetRequiredService()); JsonElement schema = tool.GetToolMetadata().InputSchema; diff --git a/src/Service.Tests/Mcp/DynamicCustomToolTests.cs b/src/Service.Tests/Mcp/DynamicCustomToolTests.cs index 9d7c36f90c..aad83ed358 100644 --- a/src/Service.Tests/Mcp/DynamicCustomToolTests.cs +++ b/src/Service.Tests/Mcp/DynamicCustomToolTests.cs @@ -571,13 +571,19 @@ public void InitializeMetadata_UnresolvableOrNonStoredProcedureMetadata_FallsBac { Entity entity = CreateTestStoredProcedureEntity(parameters: new[] { new ParameterMetadata { Name = "id" } }); DynamicCustomTool missingMetadataTool = new("MissingMetadata", entity); - missingMetadataTool.InitializeMetadata(BuildMetadataServiceProvider("MissingMetadata", metadataObject: null)); + IServiceProvider missingMetadataServices = BuildMetadataServiceProvider("MissingMetadata", metadataObject: null); + Assert.IsFalse(missingMetadataTool.InitializeMetadata( + missingMetadataServices.GetRequiredService().GetConfig(), + missingMetadataServices.GetRequiredService())); Assert.AreEqual(JsonValueKind.Array, ParseSchemaProperties(missingMetadataTool.GetToolMetadata()).GetProperty("id").GetProperty("type").ValueKind); DynamicCustomTool tableMetadataTool = new("TableMetadata", entity); - tableMetadataTool.InitializeMetadata(BuildMetadataServiceProvider( + IServiceProvider tableMetadataServices = BuildMetadataServiceProvider( "TableMetadata", - new DatabaseTable("dbo", "books") { SourceType = EntitySourceType.Table })); + new DatabaseTable("dbo", "books") { SourceType = EntitySourceType.Table }); + Assert.IsFalse(tableMetadataTool.InitializeMetadata( + tableMetadataServices.GetRequiredService().GetConfig(), + tableMetadataServices.GetRequiredService())); Assert.AreEqual(JsonValueKind.Array, ParseSchemaProperties(tableMetadataTool.GetToolMetadata()).GetProperty("id").GetProperty("type").ValueKind); } @@ -883,19 +889,25 @@ public void GetToolMetadata_UsesDbMetadata_WhenInitialized() [TestMethod] public void GetToolMetadata_FallsBackToConfig_WhenDbMetadataUnavailable() { - // Arrange - use a service provider without metadata factory + // Arrange - metadata mapping does not contain the configured entity. ParameterMetadata[] parameters = new[] { new ParameterMetadata { Name = "userId", Description = "User ID" } }; Entity entity = CreateTestStoredProcedureEntity(parameters: parameters); DynamicCustomTool tool = new("GetUser", entity); - - ServiceCollection services = new(); - services.AddLogging(); + IServiceProvider serviceProvider = BuildServiceProviderForMetadata( + "GetUser", + new Dictionary(), + metadataAvailable: false); + RuntimeConfig config = serviceProvider + .GetRequiredService() + .GetConfig(); + IMetadataProviderFactory metadataProviderFactory = serviceProvider + .GetRequiredService(); // Act - tool.InitializeMetadata(services.BuildServiceProvider()); + tool.InitializeMetadata(config, metadataProviderFactory); JsonElement props = ParseSchemaProperties(tool.GetToolMetadata()); // Assert - should use config-based permissive type array @@ -1169,9 +1181,14 @@ private static JsonElement InitializeAndGetSchema( { Entity entity = CreateTestStoredProcedureEntity(); DynamicCustomTool tool = new(entityName, entity); - IServiceProvider sp = BuildServiceProviderForMetadata(entityName, dbParameters); - - tool.InitializeMetadata(sp); + IServiceProvider serviceProvider = BuildServiceProviderForMetadata(entityName, dbParameters); + RuntimeConfig config = serviceProvider + .GetRequiredService() + .GetConfig(); + IMetadataProviderFactory metadataProviderFactory = serviceProvider + .GetRequiredService(); + + tool.InitializeMetadata(config, metadataProviderFactory); return tool.GetToolMetadata().InputSchema; } @@ -1185,9 +1202,14 @@ private static JsonElement InitializeAndGetSchemaProperties( { Entity entity = CreateTestStoredProcedureEntity(); DynamicCustomTool tool = new(entityName, entity); - IServiceProvider sp = BuildServiceProviderForMetadata(entityName, dbParameters); - - tool.InitializeMetadata(sp); + IServiceProvider serviceProvider = BuildServiceProviderForMetadata(entityName, dbParameters); + RuntimeConfig config = serviceProvider + .GetRequiredService() + .GetConfig(); + IMetadataProviderFactory metadataProviderFactory = serviceProvider + .GetRequiredService(); + + tool.InitializeMetadata(config, metadataProviderFactory); return ParseSchemaProperties(tool.GetToolMetadata()); } @@ -1204,7 +1226,8 @@ private static JsonElement ParseSchemaProperties(ModelContextProtocol.Protocol.T /// private static IServiceProvider BuildServiceProviderForMetadata( string entityName, - Dictionary dbParameters) + Dictionary dbParameters, + bool metadataAvailable = true) { Entity entity = new( Source: new("test_procedure", EntitySourceType.StoredProcedure, Parameters: null, KeyFields: null), @@ -1252,7 +1275,9 @@ private static IServiceProvider BuildServiceProviderForMetadata( Mock mockSqlMetadataProvider = new(); mockSqlMetadataProvider .Setup(x => x.EntityToDatabaseObject) - .Returns(new Dictionary { [entityName] = dbObject }); + .Returns(metadataAvailable + ? new Dictionary { [entityName] = dbObject } + : new Dictionary()); Mock mockMetadataProviderFactory = new(); mockMetadataProviderFactory diff --git a/src/Service.Tests/Mcp/McpHttpToolRegistryHotReloadIntegrationTests.cs b/src/Service.Tests/Mcp/McpHttpToolRegistryHotReloadIntegrationTests.cs new file mode 100644 index 0000000000..e475f51252 --- /dev/null +++ b/src/Service.Tests/Mcp/McpHttpToolRegistryHotReloadIntegrationTests.cs @@ -0,0 +1,553 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; +using System.IO; +using System.Linq; +using System.Net; +using System.Net.Http; +using System.Net.Http.Json; +using System.Text.Json; +using System.Threading.Tasks; +using Azure.DataApiBuilder.Config; +using Azure.DataApiBuilder.Config.ObjectModel; +using Azure.DataApiBuilder.Core.Authorization; +using Azure.DataApiBuilder.Service.Tests.Configuration; +using Microsoft.AspNetCore.Hosting; +using Microsoft.AspNetCore.TestHost; +using Microsoft.Data.SqlClient; +using Microsoft.Extensions.Logging; +using Microsoft.VisualStudio.TestTools.UnitTesting; + +namespace Azure.DataApiBuilder.Service.Tests.Mcp +{ + [TestClass, TestCategory(TestCategory.MSSQL)] + public class McpHttpToolRegistryHotReloadIntegrationTests + { + private const string MCP_PATH = "/mcp"; + private const string TEST_CONNECTION_STRING_ENV = "DAB_TEST_MSSQL_CONNECTION_STRING"; + + [TestMethod] + public async Task HttpTransport_FileReload_UpdatesDiscoveryCallsAndRecoversFromFailure() + { + TestHelper.SetupDatabaseEnvironment(TestCategory.MSSQL); + SqlConnectionStringBuilder connectionString = new( + Environment.GetEnvironmentVariable(TEST_CONNECTION_STRING_ENV) ?? + ConfigurationTests.GetConnectionStringFromEnvironmentConfig(TestCategory.MSSQL)) + { + TrustServerCertificate = true + }; + string testDirectory = Path.Combine( + Path.GetTempPath(), + $"dab-mcp-hot-reload-{Guid.NewGuid():N}"); + Directory.CreateDirectory(testDirectory); + string configPath = Path.Combine(testDirectory, "dab-config.json"); + await WriteConfigAsync( + configPath, + CreateConfig(connectionString.ConnectionString, ("GetBook", "Initial description"))); + + try + { + string[] args = + { + $"--ConfigFileName={configPath}", + "--no-https-redirect" + }; + using RejectedCandidateLogObserver rejectedCandidateLogObserver = new(); + using TestServer server = new( + Program.CreateWebHostBuilder(args) + .ConfigureLogging(logging => + logging.AddProvider(rejectedCandidateLogObserver))); + using HttpClient client = server.CreateClient(); + + McpHttpResponse initialize = await SendMcpAsync( + client, + sessionId: null, + new + { + jsonrpc = "2.0", + id = 1, + method = "initialize", + @params = new + { + protocolVersion = "2025-11-25", + capabilities = new { }, + clientInfo = new { name = "hot-reload-test", version = "1.0" } + } + }, + HttpStatusCode.OK); + Assert.IsNotNull(initialize.SessionId); + JsonElement toolCapabilities = initialize.Payload!.Value + .GetProperty("result") + .GetProperty("capabilities") + .GetProperty("tools"); + Assert.IsTrue( + !toolCapabilities.TryGetProperty("listChanged", out JsonElement listChanged) || + !listChanged.GetBoolean(), + "HTTP must not advertise listChanged until session broadcast is implemented."); + + string sessionId = initialize.SessionId; + await SendMcpAsync( + client, + sessionId, + new + { + jsonrpc = "2.0", + method = "notifications/initialized", + @params = new { } + }, + HttpStatusCode.Accepted); + + JsonElement initialList = await ListToolsAsync(client, sessionId, requestId: 2); + AssertTool(initialList, "get_book", "Initial description"); + Assert.IsTrue( + GetTools(initialList) + .Single(tool => tool.GetProperty("name").GetString() == "get_book") + .GetProperty("inputSchema") + .GetProperty("properties") + .TryGetProperty("id", out JsonElement idSchema) && + idSchema.GetProperty("type").GetString() == "integer", + "Initial HTTP discovery should use database metadata."); + await AssertToolCallSucceedsAsync(client, sessionId, "get_book", requestId: 3); + + // Change only the backing stored procedure. The refreshed metadata provider must + // supply update_book_title's additional @title parameter to the new tool schema. + await WriteConfigAsync( + configPath, + CreateConfig( + connectionString.ConnectionString, + storedProcedure: "update_book_title", + dmlToolsEnabled: false, + ("GetBook", "Initial description"))); + JsonElement changedSchemaList = await WaitForToolSchemaPropertyAsync( + client, + sessionId, + toolName: "get_book", + propertyName: "title"); + JsonElement changedProperties = GetTools(changedSchemaList) + .Single(tool => tool.GetProperty("name").GetString() == "get_book") + .GetProperty("inputSchema") + .GetProperty("properties"); + Assert.AreEqual("integer", changedProperties.GetProperty("id").GetProperty("type").GetString()); + Assert.AreEqual("string", changedProperties.GetProperty("title").GetProperty("type").GetString()); + + // Global built-in DML visibility is also snapshot state. Toggle it through real + // file changes and observe the production HTTP tools/list handler in both directions. + await WriteConfigAsync( + configPath, + CreateConfig( + connectionString.ConnectionString, + storedProcedure: "update_book_title", + dmlToolsEnabled: true, + ("GetBook", "Initial description"))); + JsonElement dmlEnabledList = await WaitForToolSetAsync( + client, + sessionId, + expectedName: "create_record", + absentName: "not_a_tool"); + Assert.IsTrue(HasTool(dmlEnabledList, "get_book")); + + await WriteConfigAsync( + configPath, + CreateConfig( + connectionString.ConnectionString, + storedProcedure: "update_book_title", + dmlToolsEnabled: false, + ("GetBook", "Initial description"))); + JsonElement dmlDisabledList = await WaitForToolSetAsync( + client, + sessionId, + expectedName: "get_book", + absentName: "create_record"); + Assert.IsFalse(HasTool(dmlDisabledList, "create_record")); + + await WriteConfigAsync( + configPath, + CreateConfig(connectionString.ConnectionString, ("LookupBook", "Reloaded description"))); + JsonElement renamedList = await WaitForToolSetAsync( + client, + sessionId, + expectedName: "lookup_book", + absentName: "get_book"); + AssertTool(renamedList, "lookup_book", "Reloaded description"); + await AssertToolCallFailsAsync(client, sessionId, "get_book", requestId: 4); + await AssertToolCallSucceedsAsync(client, sessionId, "lookup_book", requestId: 5); + + // Two physical writes without waiting for the first reload to finish exercise + // coalesced/overlapping watcher notifications. The eventual snapshot must be the + // latest complete generation. + await WriteConfigAsync( + configPath, + CreateConfig(connectionString.ConnectionString, ("IntermediateBook", "Intermediate"))); + await Task.Delay(20); + await WriteConfigAsync( + configPath, + CreateConfig(connectionString.ConnectionString, ("LatestBook", "Latest"))); + JsonElement latestList = await WaitForToolSetAsync( + client, + sessionId, + expectedName: "latest_book", + absentName: "lookup_book"); + AssertTool(latestList, "latest_book", "Latest"); + + // Both entity names normalize to duplicate_tool. The rejected candidate must leave + // latest_book published until a later valid file change recovers. + await WriteConfigAsync( + configPath, + CreateConfig( + connectionString.ConnectionString, + ("DuplicateTool", "First duplicate"), + ("duplicate_tool", "Second duplicate"))); + await rejectedCandidateLogObserver.WaitForRejectionAsync( + TimeSpan.FromSeconds(10)); + + JsonElement afterFailure = await ListToolsAsync(client, sessionId, requestId: 6); + AssertTool(afterFailure, "latest_book", "Latest"); + Assert.IsFalse(HasTool(afterFailure, "duplicate_tool")); + + await WriteConfigAsync( + configPath, + CreateConfig(connectionString.ConnectionString, ("RecoveredBook", "Recovered"))); + JsonElement recoveredList = await WaitForToolSetAsync( + client, + sessionId, + expectedName: "recovered_book", + absentName: "latest_book"); + AssertTool(recoveredList, "recovered_book", "Recovered"); + } + finally + { + Directory.Delete(testDirectory, recursive: true); + } + } + + private static RuntimeConfig CreateConfig( + string connectionString, + params (string EntityName, string Description)[] tools) + { + return CreateConfig( + connectionString, + storedProcedure: "get_book_by_id", + dmlToolsEnabled: false, + tools); + } + + private static RuntimeConfig CreateConfig( + string connectionString, + string storedProcedure, + bool dmlToolsEnabled, + params (string EntityName, string Description)[] tools) + { + Dictionary entities = tools.ToDictionary( + tool => tool.EntityName, + tool => new Entity( + Source: new( + Object: storedProcedure, + Type: EntitySourceType.StoredProcedure, + Parameters: null, + KeyFields: null), + GraphQL: new( + Singular: tool.EntityName, + Plural: tool.EntityName, + Enabled: false, + Operation: GraphQLOperation.Mutation), + Rest: new(Enabled: true), + Fields: null, + Permissions: new[] + { + new EntityPermission( + Role: AuthorizationResolver.ROLE_ANONYMOUS, + Actions: new[] + { + new EntityAction( + Action: EntityActionOperation.Execute, + Fields: null, + Policy: null) + }) + }, + Relationships: null, + Mappings: null, + Description: tool.Description, + Mcp: new EntityMcpOptions(customToolEnabled: true, dmlToolsEnabled: false))); + + return new RuntimeConfig( + Schema: FileSystemRuntimeConfigLoader.SCHEMA, + DataSource: new DataSource(DatabaseType.MSSQL, connectionString, Options: null), + Runtime: new( + Rest: new(Enabled: true), + GraphQL: new(Enabled: false), + Mcp: new( + Enabled: true, + Path: MCP_PATH, + DmlTools: DmlToolsConfig.FromBoolean(dmlToolsEnabled)), + Host: new( + Cors: null, + Authentication: new( + Provider: AuthenticationOptions.UNAUTHENTICATED_AUTHENTICATION), + Mode: HostMode.Development)), + Entities: new(entities)); + } + + private static async Task WriteConfigAsync(string configPath, RuntimeConfig config) + { + const int MAX_ATTEMPTS = 20; + for (int attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) + { + try + { + await File.WriteAllTextAsync(configPath, config.ToJson()); + return; + } + catch (IOException) when (attempt < MAX_ATTEMPTS) + { + await Task.Delay(25); + } + } + } + + private static async Task SendMcpAsync( + HttpClient client, + string? sessionId, + object payload, + HttpStatusCode expectedStatus) + { + using HttpRequestMessage request = new(HttpMethod.Post, MCP_PATH) + { + Content = JsonContent.Create(payload) + }; + request.Headers.Add("Accept", "application/json, text/event-stream"); + if (sessionId is not null) + { + request.Headers.Add("Mcp-Session-Id", sessionId); + } + + using HttpResponseMessage response = await client.SendAsync(request); + string responseBody = await response.Content.ReadAsStringAsync(); + Assert.AreEqual(expectedStatus, response.StatusCode, responseBody); + + string? responseSessionId = response.Headers.TryGetValues( + "Mcp-Session-Id", + out IEnumerable? values) + ? values.Single() + : sessionId; + JsonElement? responsePayload = string.IsNullOrWhiteSpace(responseBody) + ? null + : ParseMcpPayload(responseBody); + return new McpHttpResponse(responseSessionId, responsePayload); + } + + private static async Task ListToolsAsync( + HttpClient client, + string sessionId, + int requestId) + { + McpHttpResponse response = await SendMcpAsync( + client, + sessionId, + new + { + jsonrpc = "2.0", + id = requestId, + method = "tools/list", + @params = new { } + }, + HttpStatusCode.OK); + return response.Payload!.Value; + } + + private static async Task WaitForToolSetAsync( + HttpClient client, + string sessionId, + string expectedName, + string absentName) + { + for (int attempt = 0; attempt < 100; attempt++) + { + JsonElement response = await ListToolsAsync(client, sessionId, 100 + attempt); + if (HasTool(response, expectedName) && !HasTool(response, absentName)) + { + return response; + } + + await Task.Delay(100); + } + + Assert.Fail($"Timed out waiting for MCP tool '{expectedName}' to replace '{absentName}'."); + return default; + } + + private static async Task WaitForToolSchemaPropertyAsync( + HttpClient client, + string sessionId, + string toolName, + string propertyName) + { + for (int attempt = 0; attempt < 100; attempt++) + { + JsonElement response = await ListToolsAsync(client, sessionId, 300 + attempt); + JsonElement? matchingTool = GetTools(response) + .Cast() + .SingleOrDefault(tool => + tool?.GetProperty("name").GetString() == toolName); + if (matchingTool.HasValue && + matchingTool.Value + .GetProperty("inputSchema") + .GetProperty("properties") + .TryGetProperty(propertyName, out _)) + { + return response; + } + + await Task.Delay(100); + } + + Assert.Fail( + $"Timed out waiting for MCP tool '{toolName}' schema property '{propertyName}'."); + return default; + } + + private static async Task AssertToolCallSucceedsAsync( + HttpClient client, + string sessionId, + string toolName, + int requestId) + { + McpHttpResponse response = await SendMcpAsync( + client, + sessionId, + new + { + jsonrpc = "2.0", + id = requestId, + method = "tools/call", + @params = new + { + name = toolName, + arguments = new { id = 1 } + } + }, + HttpStatusCode.OK); + + Assert.IsTrue(response.Payload!.Value.TryGetProperty("result", out JsonElement result)); + Assert.IsFalse(result.TryGetProperty("isError", out JsonElement isError) && isError.GetBoolean()); + } + + private static async Task AssertToolCallFailsAsync( + HttpClient client, + string sessionId, + string toolName, + int requestId) + { + McpHttpResponse response = await SendMcpAsync( + client, + sessionId, + new + { + jsonrpc = "2.0", + id = requestId, + method = "tools/call", + @params = new + { + name = toolName, + arguments = new { id = 1 } + } + }, + HttpStatusCode.OK); + + JsonElement payload = response.Payload!.Value; + bool hasJsonRpcError = payload.TryGetProperty("error", out _); + bool hasToolError = payload.TryGetProperty("result", out JsonElement result) && + result.TryGetProperty("isError", out JsonElement isError) && + isError.GetBoolean(); + Assert.IsTrue( + hasJsonRpcError || hasToolError, + $"Calling removed tool '{toolName}' should return an MCP error result."); + } + + private static JsonElement ParseMcpPayload(string responseBody) + { + string json = responseBody.TrimStart().StartsWith('{') + ? responseBody + : responseBody + .Split('\n') + .Select(line => line.TrimEnd('\r')) + .Where(line => line.StartsWith("data:", StringComparison.Ordinal)) + .Select(line => line["data:".Length..].TrimStart()) + .First(payload => payload.StartsWith('{')); + using JsonDocument document = JsonDocument.Parse(json); + return document.RootElement.Clone(); + } + + private static IEnumerable GetTools(JsonElement response) + { + return response + .GetProperty("result") + .GetProperty("tools") + .EnumerateArray(); + } + + private static bool HasTool(JsonElement response, string name) + { + return GetTools(response) + .Any(tool => string.Equals( + tool.GetProperty("name").GetString(), + name, + StringComparison.Ordinal)); + } + + private static void AssertTool(JsonElement response, string name, string description) + { + JsonElement tool = GetTools(response) + .Single(tool => tool.GetProperty("name").GetString() == name); + Assert.AreEqual(description, tool.GetProperty("description").GetString()); + } + + private sealed record McpHttpResponse(string? SessionId, JsonElement? Payload); + + private sealed class RejectedCandidateLogObserver : ILoggerProvider + { + private readonly TaskCompletionSource _rejectionObserved = new( + TaskCreationOptions.RunContinuationsAsynchronously); + + public ILogger CreateLogger(string categoryName) + { + return new RejectedCandidateLogger(_rejectionObserved); + } + + public async Task WaitForRejectionAsync(TimeSpan timeout) + { + await _rejectionObserved.Task.WaitAsync(timeout); + } + + public void Dispose() + { + } + + private sealed class RejectedCandidateLogger( + TaskCompletionSource rejectionObserved) : ILogger + { + public IDisposable? BeginScope(TState state) + where TState : notnull => null; + + public bool IsEnabled(LogLevel logLevel) => true; + + public void Log( + LogLevel logLevel, + EventId eventId, + TState state, + Exception? exception, + Func formatter) + { + if (logLevel == LogLevel.Error && + formatter(state, exception).Contains( + "Failed to refresh the MCP tool registry after a runtime configuration change.", + StringComparison.Ordinal)) + { + rejectionObserved.TrySetResult(); + } + } + } + } + } +} diff --git a/src/Service.Tests/Mcp/McpInitialHotReloadSerializationTests.cs b/src/Service.Tests/Mcp/McpInitialHotReloadSerializationTests.cs new file mode 100644 index 0000000000..917907c161 --- /dev/null +++ b/src/Service.Tests/Mcp/McpInitialHotReloadSerializationTests.cs @@ -0,0 +1,289 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; +using System.Data; +using System.IO; +using System.IO.Abstractions; +using System.Linq; +using System.Text.Json; +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.Mcp.Core; +using Azure.DataApiBuilder.Mcp.Model; +using Azure.DataApiBuilder.Service.Utilities; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging.Abstractions; +using Microsoft.VisualStudio.TestTools.UnitTesting; +using ModelContextProtocol.Protocol; +using Moq; +using static Azure.DataApiBuilder.Config.DabConfigEvents; + +namespace Azure.DataApiBuilder.Service.Tests.Mcp +{ + [TestClass] + public class McpInitialHotReloadSerializationTests + { + [TestMethod] + public async Task InitialConstructionAndReload_PublishLatestDatabaseMetadataGeneration() + { + string testDirectory = Path.Combine( + Path.GetTempPath(), + $"dab-mcp-initial-reload-serialization-{Guid.NewGuid():N}"); + Directory.CreateDirectory(testDirectory); + string configPath = Path.Combine(testDirectory, "dab-config.json"); + File.WriteAllText(configPath, CreateRuntimeConfig("Generation A").ToJson()); + + try + { + HotReloadEventHandler hotReloadEventHandler = new(); + FileSystem fileSystem = new(); + + // The OS watcher is disabled so synchronization barriers, rather than filesystem + // notification timing, deterministically control this startup-to-reload race. + using FileSystemRuntimeConfigLoader configLoader = new( + fileSystem, + hotReloadEventHandler, + configPath, + connectionString: null, + isCliLoader: true); + Assert.IsTrue(configLoader.TryLoadKnownConfig(out RuntimeConfig initialConfig)); + Assert.AreEqual( + "Generation A", + initialConfig.Entities["GetBook"].Description); + + // This focused test supplies database metadata directly. Use a provider backed by + // the real loader state without attaching live-database validation to its change + // token before the ordered handlers run. + Mock providerLoader = new(null, null); + Mock runtimeConfigProvider = new(providerLoader.Object); + runtimeConfigProvider + .Setup(provider => provider.GetConfig()) + .Returns(() => configLoader.RuntimeConfig!); + + Dictionary currentMetadata = + CreateStoredProcedureMetadata("a_database_parameter", typeof(string), DbType.String); + Mock sqlMetadataProvider = new(); + sqlMetadataProvider + .SetupGet(provider => provider.EntityToDatabaseObject) + .Returns(() => Volatile.Read(ref currentMetadata)); + + using ManualResetEventSlim initialMetadataInitializationEntered = new(); + TaskCompletionSource initialMetadataMayComplete = new( + TaskCreationOptions.RunContinuationsAsynchronously); + Mock metadataProviderFactory = new(); + metadataProviderFactory + .Setup(factory => factory.GetMetadataProvider(It.IsAny())) + .Returns(sqlMetadataProvider.Object); + metadataProviderFactory + .Setup(factory => factory.InitializeAsync(It.IsAny())) + .Callback(initialMetadataInitializationEntered.Set) + .Returns(initialMetadataMayComplete.Task); + + McpToolRegistry registry = new(); + McpToolRegistryRefreshService refreshService = new( + runtimeConfigProvider.Object, + Array.Empty(), + registry, + metadataProviderFactory.Object, + Array.Empty(), + NullLogger.Instance, + hotReloadEventHandler); + + RuntimeConfigValidator runtimeConfigValidator = new( + runtimeConfigProvider.Object, + fileSystem, + NullLogger.Instance); + using ServiceProvider serviceProvider = new ServiceCollection() + .AddSingleton(configLoader) + .AddSingleton(runtimeConfigProvider.Object) + .AddSingleton(runtimeConfigValidator) + .AddSingleton(metadataProviderFactory.Object) + .AddSingleton(refreshService) + .BuildServiceProvider(); + + using ManualResetEventSlim reloadPausedBeforeMetadata = new(); + using ManualResetEventSlim reloadReachedGate = new(); + using ManualResetEventSlim releaseReload = new(); + hotReloadEventHandler.Subscribe( + QUERY_MANAGER_FACTORY_ON_CONFIG_CHANGED, + (_, _) => + { + reloadPausedBeforeMetadata.Set(); + if (!releaseReload.Wait(TimeSpan.FromSeconds(10))) + { + throw new TimeoutException("Timed out waiting to resume reload B."); + } + }); + hotReloadEventHandler.Subscribe( + METADATA_PROVIDER_FACTORY_ON_CONFIG_CHANGED, + (_, _) => Volatile.Write( + ref currentMetadata, + CreateStoredProcedureMetadata( + "b_database_parameter", + typeof(int), + DbType.Int32))); + + Task initialConstruction = Task.Factory.StartNew( + () => RuntimeInitializationHelper.InitializeRuntimeDependenciesAsync( + serviceProvider), + CancellationToken.None, + TaskCreationOptions.LongRunning, + TaskScheduler.Default).Unwrap(); + Task reloadB = Task.CompletedTask; + + try + { + Assert.IsTrue( + initialMetadataInitializationEntered.Wait(TimeSpan.FromSeconds(10)), + "Initial metadata initialization for generation A did not start."); + + File.WriteAllText(configPath, CreateRuntimeConfig("Generation B").ToJson()); + reloadB = Task.Run(() => configLoader.ProcessHotReloadNotification( + beforeEnteringGate: reloadReachedGate.Set)); + Assert.IsTrue( + reloadReachedGate.Wait(TimeSpan.FromSeconds(10)), + "Reload B did not reach the shared serialization gate."); + + // Without startup serialization, B reaches this handler while A's metadata + // task is incomplete. Completing A then publishes a B/A candidate, and B's + // later MCP handler skips because B was incorrectly marked as applied. + bool reloadEnteredDuringInitialMetadata = + reloadPausedBeforeMetadata.Wait(TimeSpan.FromMilliseconds(500)); + Assert.AreEqual( + 0, + registry.GetAdvertisedTools().Count, + "No registry generation should publish before initial metadata completes."); + + initialMetadataMayComplete.SetResult(); + await initialConstruction.WaitAsync(TimeSpan.FromSeconds(10)); + + if (!reloadEnteredDuringInitialMetadata) + { + Assert.IsTrue( + reloadPausedBeforeMetadata.Wait(TimeSpan.FromSeconds(10)), + "Reload B did not pause before refreshing database metadata."); + } + + releaseReload.Set(); + await reloadB.WaitAsync(TimeSpan.FromSeconds(10)); + } + finally + { + initialMetadataMayComplete.TrySetResult(); + releaseReload.Set(); + await Task.WhenAll(reloadB, initialConstruction).WaitAsync(TimeSpan.FromSeconds(10)); + } + + Assert.IsTrue(initialMetadataInitializationEntered.IsSet); + metadataProviderFactory.Verify( + factory => factory.InitializeAsync(It.IsAny()), + Times.Once); + + Tool advertisedTool = registry.GetAdvertisedTools().Single(); + Assert.AreEqual("get_book", advertisedTool.Name); + Assert.AreEqual("Generation B", advertisedTool.Description); + JsonElement properties = advertisedTool.InputSchema.GetProperty("properties"); + Assert.IsTrue(properties.TryGetProperty("b_database_parameter", out JsonElement parameter)); + Assert.AreEqual("integer", parameter.GetProperty("type").GetString()); + Assert.IsFalse(properties.TryGetProperty("a_database_parameter", out _)); + Assert.IsFalse(properties.TryGetProperty("config_parameter", out _)); + } + finally + { + Directory.Delete(testDirectory, recursive: true); + } + } + + private static RuntimeConfig CreateRuntimeConfig(string description) + { + Entity entity = new( + Source: new( + Object: "test_procedure", + Type: EntitySourceType.StoredProcedure, + Parameters: new List + { + new() + { + Name = "config_parameter", + Description = "Configuration fallback parameter", + Required = true + } + }, + KeyFields: null), + GraphQL: new( + Singular: "GetBook", + Plural: "GetBooks", + Enabled: false, + Operation: GraphQLOperation.Mutation), + Rest: new(Enabled: true), + Fields: null, + Permissions: new[] + { + new EntityPermission( + Role: "anonymous", + Actions: new[] + { + new EntityAction( + Action: EntityActionOperation.Execute, + Fields: null, + Policy: null) + }) + }, + Relationships: null, + Mappings: null, + Description: description, + Mcp: new EntityMcpOptions(customToolEnabled: true, dmlToolsEnabled: false)); + + return new RuntimeConfig( + Schema: FileSystemRuntimeConfigLoader.SCHEMA, + DataSource: new DataSource( + DatabaseType.MSSQL, + "Server=test;Database=test;User ID=test;Password=test;TrustServerCertificate=true", + Options: null), + Runtime: new( + Rest: new(Enabled: true), + GraphQL: new(Enabled: false), + Mcp: new(Enabled: true, DmlTools: DmlToolsConfig.FromBoolean(false)), + Host: new( + Cors: null, + Authentication: new( + Provider: AuthenticationOptions.UNAUTHENTICATED_AUTHENTICATION), + Mode: HostMode.Development)), + Entities: new(new Dictionary { ["GetBook"] = entity })); + } + + private static Dictionary CreateStoredProcedureMetadata( + string parameterName, + Type systemType, + DbType dbType) + { + DatabaseStoredProcedure storedProcedure = new("dbo", "test_procedure") + { + SourceType = EntitySourceType.StoredProcedure, + StoredProcedureDefinition = new StoredProcedureDefinition + { + Parameters = new Dictionary + { + [parameterName] = new ParameterDefinition + { + Name = parameterName, + Required = true, + SystemType = systemType, + DbType = dbType + } + } + } + }; + + return new Dictionary { ["GetBook"] = storedProcedure }; + } + } +} diff --git a/src/Service.Tests/Mcp/McpMetadataHelperTests.cs b/src/Service.Tests/Mcp/McpMetadataHelperTests.cs index db2c90e909..7c1cb13d6c 100644 --- a/src/Service.Tests/Mcp/McpMetadataHelperTests.cs +++ b/src/Service.Tests/Mcp/McpMetadataHelperTests.cs @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +#nullable enable + using System; using System.Collections.Generic; using System.Net; @@ -26,17 +28,47 @@ public class McpMetadataHelperTests { private const string ENTITY_NAME = "Book"; + [DataTestMethod] + [DataRow(null, DisplayName = "Null entity name")] + [DataRow("", DisplayName = "Empty entity name")] + [DataRow(" ", DisplayName = "Whitespace entity name")] + public void TryResolveMetadata_ExplicitFactory_InvalidEntityNameReturnsFalse(string? entityName) + { + RuntimeConfig config = new( + Schema: "test-schema", + DataSource: null, + Runtime: null, + Entities: new RuntimeEntities(new Dictionary())); + Mock metadataProviderFactory = new(); + + bool resolved = McpMetadataHelper.TryResolveMetadata( + entityName!, + config, + metadataProviderFactory.Object, + out ISqlMetadataProvider _, + out DatabaseObject _, + out string dataSourceName, + out string error); + + Assert.IsFalse(resolved); + Assert.AreEqual(string.Empty, dataSourceName); + Assert.AreEqual("Entity name cannot be null or empty.", error); + metadataProviderFactory.Verify( + factory => factory.GetMetadataProvider(It.IsAny()), + Times.Never); + } + [DataTestMethod] [DataRow(null)] [DataRow("")] [DataRow(" ")] - public void TryResolveMetadata_NullOrEmptyEntityName_ReturnsFalse(string entityName) + public void TryResolveMetadata_NullOrEmptyEntityName_ReturnsFalse(string? entityName) { RuntimeConfig config = CreateConfig(includeBookEntity: true); IServiceProvider serviceProvider = CreateServiceProvider(registerFactory: true, includeBookMetadata: true); bool result = McpMetadataHelper.TryResolveMetadata( - entityName, config, serviceProvider, out _, out _, out _, out string error); + entityName!, config, serviceProvider, out _, out _, out _, out string error); Assert.IsFalse(result); Assert.AreEqual("Entity name cannot be null or empty.", error); @@ -79,7 +111,7 @@ public void TryResolveMetadata_EntityNotInMetadata_ReturnsFalse() ENTITY_NAME, config, serviceProvider, out _, out _, out _, out string error); Assert.IsFalse(result); - StringAssert.Contains(error, "is not defined in the configuration"); + StringAssert.Contains(error, "Database metadata for entity 'Book' was not available"); } [TestMethod] @@ -134,7 +166,7 @@ public void TryResolveDatabaseObject_Failure_ReturnsNull() ENTITY_NAME, config, serviceProvider, out string error); Assert.IsNull(dbObject); - StringAssert.Contains(error, "is not defined in the configuration"); + StringAssert.Contains(error, "Database metadata for entity 'Book' was not available"); } /// diff --git a/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs b/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs new file mode 100644 index 0000000000..96607551a4 --- /dev/null +++ b/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs @@ -0,0 +1,246 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#nullable enable + +using System; +using System.Collections.Generic; +using System.IO; +using System.IO.Abstractions; +using System.Linq; +using System.Text; +using System.Text.Json; +using System.Threading; +using System.Threading.Channels; +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.Mcp.Core; +using Azure.DataApiBuilder.Mcp.Model; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging.Abstractions; +using Microsoft.VisualStudio.TestTools.UnitTesting; +using Moq; + +namespace Azure.DataApiBuilder.Service.Tests.Mcp +{ + [TestClass] + public class McpStdioToolRegistryHotReloadIntegrationTests + { + [TestMethod] + public async Task InitializedClient_FileReload_EmitsOneNotificationAndReturnsUpdatedList() + { + string testDirectory = Path.Combine( + Path.GetTempPath(), + $"dab-mcp-stdio-hot-reload-{Guid.NewGuid():N}"); + Directory.CreateDirectory(testDirectory); + string configPath = Path.Combine(testDirectory, "dab-config.json"); + await WriteConfigAsync(configPath, CreateRuntimeConfig()); + + try + { + HotReloadEventHandler hotReloadEventHandler = new(); + using FileSystemRuntimeConfigLoader fileLoader = new( + new FileSystem(), + hotReloadEventHandler, + configPath); + Assert.IsTrue(fileLoader.TryLoadKnownConfig(out _)); + + // The refresh service reads the real loader's active generation. Keep the provider + // itself detached from the change token so this transport test does not invoke the + // separate live-database configuration validator. + Mock providerLoader = new(null, null); + Mock configProvider = new(providerLoader.Object); + configProvider + .Setup(provider => provider.GetConfig()) + .Returns(() => fileLoader.RuntimeConfig!); + + Mock sqlMetadataProvider = new(); + sqlMetadataProvider + .SetupGet(provider => provider.EntityToDatabaseObject) + .Returns(new Dictionary()); + Mock metadataProviderFactory = new(); + metadataProviderFactory + .Setup(factory => factory.GetMetadataProvider(It.IsAny())) + .Returns(sqlMetadataProvider.Object); + + McpToolRegistry registry = new(); + ChannelTextReader stdin = new(); + ChannelTextWriter stdout = new(); + using McpStdoutWriter stdoutWriter = new(stdout); + McpStdioToolListChangedNotifier notifier = new(stdoutWriter); + using ServiceProvider serviceProvider = new ServiceCollection() + .AddSingleton(stdoutWriter) + .AddSingleton(notifier) + .AddSingleton(configProvider.Object) + .BuildServiceProvider(); + + McpToolRegistryRefreshService refreshService = new( + configProvider.Object, + Array.Empty(), + registry, + metadataProviderFactory.Object, + new IMcpToolListChangedNotifier[] { notifier }, + NullLogger.Instance, + hotReloadEventHandler); + refreshService.EnsureInitialized(); + + McpStdioServer server = new(registry, serviceProvider, stdin); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + Task serverTask = server.RunAsync(timeout.Token); + + stdin.WriteLine( + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":\"2025-11-25\",\"capabilities\":{},\"clientInfo\":{\"name\":\"test\",\"version\":\"1.0\"}}}"); + stdin.WriteLine("{\"jsonrpc\":\"2.0\",\"method\":\"notifications/initialized\"}"); + stdin.WriteLine("{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"ping\"}"); + + using JsonDocument initializeResponse = await ReadJsonLineAsync(stdout, timeout.Token); + Assert.IsTrue( + initializeResponse.RootElement + .GetProperty("result") + .GetProperty("capabilities") + .GetProperty("tools") + .GetProperty("listChanged") + .GetBoolean()); + + using JsonDocument pingResponse = await ReadJsonLineAsync(stdout, timeout.Token); + Assert.AreEqual(2, pingResponse.RootElement.GetProperty("id").GetInt32(), + "The ping response is a barrier proving the initialized notification was processed."); + + await WriteConfigAsync( + configPath, + CreateRuntimeConfig(("GetBook", "Gets one book"))); + + using JsonDocument notification = await ReadJsonLineAsync(stdout, timeout.Token); + Assert.AreEqual( + "notifications/tools/list_changed", + notification.RootElement.GetProperty("method").GetString()); + Assert.IsFalse(notification.RootElement.TryGetProperty("id", out _)); + + stdin.WriteLine("{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"tools/list\"}"); + stdin.WriteLine("{\"jsonrpc\":\"2.0\",\"id\":4,\"method\":\"shutdown\"}"); + + using JsonDocument listResponse = await ReadJsonLineAsync(stdout, timeout.Token); + Assert.AreEqual( + 3, + listResponse.RootElement.GetProperty("id").GetInt32(), + "An extra notification would displace the list response and fail this barrier."); + JsonElement tool = listResponse.RootElement + .GetProperty("result") + .GetProperty("tools") + .EnumerateArray() + .Single(); + Assert.AreEqual("get_book", tool.GetProperty("name").GetString()); + Assert.AreEqual("Gets one book", tool.GetProperty("description").GetString()); + + using JsonDocument shutdownResponse = await ReadJsonLineAsync(stdout, timeout.Token); + Assert.AreEqual( + 4, + shutdownResponse.RootElement.GetProperty("id").GetInt32(), + "Exactly one notification should be emitted for one net-new file content."); + await serverTask; + } + finally + { + Directory.Delete(testDirectory, recursive: true); + } + } + + private static async Task ReadJsonLineAsync( + ChannelTextWriter output, + CancellationToken cancellationToken) + { + string line = await output.ReadLineAsync(cancellationToken); + return JsonDocument.Parse(line); + } + + private static async Task WriteConfigAsync(string configPath, RuntimeConfig config) + { + const int MAX_ATTEMPTS = 20; + for (int attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) + { + try + { + await File.WriteAllTextAsync(configPath, config.ToJson()); + return; + } + catch (IOException) when (attempt < MAX_ATTEMPTS) + { + await Task.Delay(25); + } + } + } + + private static RuntimeConfig CreateRuntimeConfig( + params (string EntityName, string Description)[] customTools) + { + Dictionary entities = customTools.ToDictionary( + item => item.EntityName, + item => new Entity( + Source: new("test_procedure", EntitySourceType.StoredProcedure, Parameters: null, KeyFields: null), + GraphQL: new(item.EntityName, item.EntityName), + Rest: new(Enabled: true), + Fields: null, + Permissions: new[] + { + new EntityPermission( + Role: "anonymous", + Actions: new[] + { + new EntityAction(Action: EntityActionOperation.Execute, Fields: null, Policy: null) + }) + }, + Relationships: null, + Mappings: null, + Description: item.Description, + Mcp: new EntityMcpOptions(customToolEnabled: true, dmlToolsEnabled: null))); + + return new RuntimeConfig( + Schema: "test-schema", + DataSource: new DataSource(DatabaseType.MSSQL, string.Empty, Options: null), + Runtime: new( + Rest: new(), + GraphQL: new(), + Mcp: new(Enabled: true), + Host: new(Cors: null, Authentication: null, Mode: HostMode.Development)), + Entities: new(entities)); + } + + private sealed class ChannelTextReader : TextReader + { + private readonly Channel _lines = Channel.CreateUnbounded(); + + public void WriteLine(string line) + { + Assert.IsTrue(_lines.Writer.TryWrite(line)); + } + + public override async ValueTask ReadLineAsync(CancellationToken cancellationToken) + { + return await _lines.Reader.ReadAsync(cancellationToken); + } + } + + private sealed class ChannelTextWriter : StringWriter + { + private readonly Channel _lines = Channel.CreateUnbounded(); + + public override Encoding Encoding => Encoding.UTF8; + + public override void WriteLine(string? value) + { + base.WriteLine(value); + Assert.IsTrue(_lines.Writer.TryWrite(value ?? string.Empty)); + } + + public async ValueTask ReadLineAsync(CancellationToken cancellationToken) + { + return await _lines.Reader.ReadAsync(cancellationToken); + } + } + } +} diff --git a/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs new file mode 100644 index 0000000000..c43271646c --- /dev/null +++ b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs @@ -0,0 +1,869 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.IO; +using System.Linq; +using System.Text.Json; +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.Mcp.Core; +using Azure.DataApiBuilder.Mcp.Model; +using Azure.DataApiBuilder.Service.Exceptions; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging; +using Microsoft.VisualStudio.TestTools.UnitTesting; +using ModelContextProtocol.Protocol; +using Moq; +using static Azure.DataApiBuilder.Config.DabConfigEvents; +using static Azure.DataApiBuilder.Mcp.Model.McpEnums; + +namespace Azure.DataApiBuilder.Service.Tests.Mcp +{ + [TestClass] + public class McpToolRegistryRefreshServiceTests + { + [TestMethod] + public void EnsureInitialized_IsIdempotentForSameConfig() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(); + TestContext context = CreateContext( + () => currentConfig, + new TestMcpTool("read_records", ToolType.BuiltIn)); + + context.Service.EnsureInitialized(); + IMcpTool? initialTool = GetRequiredTool(context.Registry, "read_records"); + context.Service.EnsureInitialized(); + + Assert.AreSame(initialTool, GetRequiredTool(context.Registry, "read_records")); + Assert.AreEqual(1, context.Registry.GetAdvertisedTools().Count); + context.Notifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Never); + } + + [TestMethod] + public void EnsureInitialized_WithCanceledToken_DoesNotPublish() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(); + TestContext context = CreateContext( + () => currentConfig, + new TestMcpTool("read_records", ToolType.BuiltIn)); + using CancellationTokenSource cancellation = new(); + cancellation.Cancel(); + + Assert.ThrowsException(() => + context.Service.EnsureInitialized(cancellation.Token)); + + Assert.AreEqual(0, context.Registry.GetAdvertisedTools().Count); + context.Notifier.Verify( + notifier => notifier.NotifyToolsListChanged(), + Times.Never); + } + + [TestMethod] + public void HotReload_AfterOutOfBandInitialization_RebuildsWithOrderedMetadataGeneration() + { + RuntimeConfig configA = CreateRuntimeConfig(("GetBook", "Stable description")); + RuntimeConfig configB = CreateRuntimeConfig(("GetBook", "Stable description")); + RuntimeConfig currentConfig = configA; + Dictionary currentMetadata = + CreateStoredProcedureMetadata("old_parameter", "Metadata generation A"); + TestContext context = CreateContextWithDatabaseMetadata( + () => currentConfig, + () => currentMetadata); + context.Service.EnsureInitialized(); + + // RuntimeConfigProvider exposes B before B's ordered metadata refresh runs. Simulate + // an out-of-band caller publishing B while the metadata provider still contains A. + currentConfig = configB; + context.Service.EnsureInitialized(); + CollectionAssert.AreEqual( + new[] { "old_parameter" }, + GetAdvertisedParameterNames(context.Registry)); + + // The metadata event installs B before the ordered MCP event. That event must rebuild + // even though the same RuntimeConfig reference was already applied above. + currentMetadata = CreateStoredProcedureMetadata( + "new_parameter", + "Metadata generation B"); + RaiseRegistryChanged(context.HotReloadEventHandler); + + CollectionAssert.AreEqual( + new[] { "new_parameter" }, + GetAdvertisedParameterNames(context.Registry)); + context.Notifier.Verify( + notifier => notifier.NotifyToolsListChanged(), + Times.Once); + } + + [TestMethod] + public async Task HostedStart_DefersInitializationToStartupOrchestrator() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(); + TestContext context = CreateContext( + () => currentConfig, + new TestMcpTool("read_records", ToolType.BuiltIn)); + + await context.Service.StartAsync(CancellationToken.None); + + Assert.AreEqual(0, context.Registry.GetAdvertisedTools().Count, + "Hosted service startup occurs before metadata initialization and must not publish."); + + context.Service.EnsureInitialized(); + + Assert.AreEqual(1, context.Registry.GetAdvertisedTools().Count, + "The startup orchestrator should publish after metadata initialization."); + } + + [TestMethod] + public void HotReload_AddsFreshCustomToolAndNotifiesClient() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(); + TestContext context = CreateContext( + () => currentConfig, + new TestMcpTool("read_records", ToolType.BuiltIn)); + context.Service.EnsureInitialized(); + + currentConfig = CreateRuntimeConfig(("GetBook", "Gets one book")); + RaiseRegistryChanged(context.HotReloadEventHandler); + + IMcpTool customTool = GetRequiredTool(context.Registry, "get_book"); + Assert.IsInstanceOfType(customTool); + Assert.AreEqual( + "Gets one book", + context.Registry.GetAdvertisedTools().Single(tool => tool.Name == "get_book").Description); + context.Notifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Once); + } + + [TestMethod] + public void HotReload_PreservesExplicitlyDiRegisteredCustomTool() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(); + Mock configLoader = new(null, null); + configLoader.Object.RuntimeConfig = currentConfig; + Mock configProvider = new(configLoader.Object); + configProvider.Setup(provider => provider.GetConfig()).Returns(() => currentConfig); + + Mock sqlMetadataProvider = new(); + sqlMetadataProvider + .SetupGet(provider => provider.EntityToDatabaseObject) + .Returns(new Dictionary()); + Mock metadataProviderFactory = new(); + metadataProviderFactory + .Setup(factory => factory.GetMetadataProvider(It.IsAny())) + .Returns(sqlMetadataProvider.Object); + + TestMcpTool registeredCustomTool = new("extension_tool", ToolType.Custom); + HotReloadEventHandler hotReloadEventHandler = new(); + ServiceCollection services = new(); + services.AddLogging(); + services.AddSingleton(configProvider.Object); + services.AddSingleton(metadataProviderFactory.Object); + services.AddSingleton(hotReloadEventHandler); + services.AddSingleton(registeredCustomTool); + services.AddDabMcpServer(configProvider.Object); + using ServiceProvider serviceProvider = services.BuildServiceProvider(); + McpToolRegistryRefreshService refreshService = serviceProvider + .GetRequiredService(); + McpToolRegistry registry = serviceProvider.GetRequiredService(); + + refreshService.EnsureInitialized(); + Assert.IsTrue(registry.TryGetTool("extension_tool", out IMcpTool? initialTool)); + Assert.AreSame(registeredCustomTool, initialTool); + Assert.IsTrue(registry.GetAdvertisedTools().Any(tool => tool.Name == "extension_tool")); + + currentConfig = CreateRuntimeConfig(); + RaiseRegistryChanged(hotReloadEventHandler); + + Assert.IsTrue(registry.TryGetTool("extension_tool", out IMcpTool? refreshedTool)); + Assert.AreSame( + registeredCustomTool, + refreshedTool, + "Independent DI-owned tools must remain published across configuration generations."); + Assert.IsTrue(registry.GetAdvertisedTools().Any(tool => tool.Name == "extension_tool")); + } + + [TestMethod] + public void HotReload_ReplacesCustomToolInstanceAndMetadata() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(("GetBook", "Old description")); + TestContext context = CreateContext(() => currentConfig); + context.Service.EnsureInitialized(); + IMcpTool oldTool = GetRequiredTool(context.Registry, "get_book"); + + currentConfig = CreateRuntimeConfig(("GetBook", "New description")); + RaiseRegistryChanged(context.HotReloadEventHandler); + + IMcpTool newTool = GetRequiredTool(context.Registry, "get_book"); + Assert.AreNotSame(oldTool, newTool); + Assert.AreEqual( + "New description", + context.Registry.GetAdvertisedTools().Single().Description); + context.Notifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Once); + } + + [TestMethod] + public void HotReload_WithDuplicateToolName_PreservesPreviousRegistry() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(); + TestMcpTool builtIn = new("read_records", ToolType.BuiltIn); + TestContext context = CreateContext(() => currentConfig, builtIn); + context.Service.EnsureInitialized(); + + currentConfig = CreateRuntimeConfig(("ReadRecords", "Conflicting custom tool")); + RaiseRegistryChanged(context.HotReloadEventHandler); + + Assert.AreSame(builtIn, GetRequiredTool(context.Registry, "read_records")); + Assert.AreEqual(1, context.Registry.GetAdvertisedTools().Count); + context.Notifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Never); + } + + [TestMethod] + public void EnsureInitialized_WithDuplicateToolName_Throws() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(("ReadRecords", "Conflicting custom tool")); + TestContext context = CreateContext( + () => currentConfig, + new TestMcpTool("read_records", ToolType.BuiltIn)); + + Assert.ThrowsException(context.Service.EnsureInitialized); + Assert.AreEqual(0, context.Registry.GetAdvertisedTools().Count); + } + + [TestMethod] + public void HotReload_WithEquivalentDiscoveryMetadata_DoesNotNotify() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(); + TestContext context = CreateContext( + () => currentConfig, + new TestMcpTool("read_records", ToolType.BuiltIn)); + context.Service.EnsureInitialized(); + + currentConfig = CreateRuntimeConfig(); + RaiseRegistryChanged(context.HotReloadEventHandler); + + context.Notifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Never); + } + + [DataTestMethod] + [DataRow("default", "enum")] + [DataRow("default", "type")] + [DataRow("default", "required")] + [DataRow("const", "enum")] + [DataRow("const", "type")] + [DataRow("const", "required")] + [DataRow("examples", "enum")] + [DataRow("examples", "type")] + [DataRow("examples", "required")] + public void HotReload_WithReorderedArrayInSchemaInstanceData_NotifiesClient(string keyword, string propertyName) + { + string CreateSchema(string values) + { + string instance = "{\"" + propertyName + "\":" + values + "}"; + return "{\"type\":\"object\",\"" + keyword + "\":" + + (keyword == "examples" ? "[" + instance + "]" : instance) + "}"; + } + + RuntimeConfig currentConfig = CreateRuntimeConfig(); + string currentSchema = CreateSchema("[\"a\",\"b\"]"); + TestContext context = CreateContext( + () => currentConfig, + new TestMcpTool("extension_tool", ToolType.Custom, metadataFactory: () => new Tool + { + Name = "extension_tool", + InputSchema = JsonSerializer.Deserialize(currentSchema) + })); + context.Service.EnsureInitialized(); + context.Notifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Never); + + currentSchema = CreateSchema("[\"b\",\"a\"]"); + currentConfig = CreateRuntimeConfig(); + RaiseRegistryChanged(context.HotReloadEventHandler); + + Assert.AreEqual(currentSchema, context.Registry.GetAdvertisedTools().Single().InputSchema.GetRawText()); + context.Notifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Once); + + currentConfig = CreateRuntimeConfig(); + RaiseRegistryChanged(context.HotReloadEventHandler); + context.Notifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Once); + } + + [TestMethod] + public void EnsureInitialized_WhenDatabaseMetadataUnavailable_PublishesConfigFallbackSchema() + { + RuntimeConfig currentConfig = CreateRuntimeConfigWithParameter( + parameterDescription: "Configured identifier"); + TestContext context = CreateContext(() => currentConfig); + + context.Service.EnsureInitialized(); + + Tool customTool = context.Registry.GetAdvertisedTools().Single(); + JsonElement properties = customTool.InputSchema.GetProperty("properties"); + JsonElement idSchema = properties.GetProperty("id"); + CollectionAssert.AreEqual( + new[] { "string", "number", "boolean", "null" }, + idSchema.GetProperty("type").EnumerateArray().Select(value => value.GetString()).ToArray()); + Assert.AreEqual("Configured identifier", idSchema.GetProperty("description").GetString()); + CollectionAssert.AreEqual( + new[] { "id" }, + customTool.InputSchema.GetProperty("required") + .EnumerateArray() + .Select(value => value.GetString()) + .ToArray()); + VerifyLogContains( + context.Logger, + LogLevel.Warning, + "Reason: Database metadata for entity 'GetBook' was not available from data source"); + VerifyLogContains( + context.Logger, + LogLevel.Information, + "with 0 built-in tools, 0 DI-registered custom tools, " + + "1 configuration-generated custom tools, 1 registered tools, and 1 advertised tools. " + + "Discovery changed: True."); + } + + [TestMethod] + public void HotReload_WithInputSchemaOnlyChange_NotifiesClient() + { + RuntimeConfig currentConfig = CreateRuntimeConfigWithParameter("Old parameter description"); + TestContext context = CreateContext(() => currentConfig); + context.Service.EnsureInitialized(); + + currentConfig = CreateRuntimeConfigWithParameter("New parameter description"); + RaiseRegistryChanged(context.HotReloadEventHandler); + + Assert.AreEqual( + "New parameter description", + context.Registry.GetAdvertisedTools() + .Single() + .InputSchema + .GetProperty("properties") + .GetProperty("id") + .GetProperty("description") + .GetString()); + context.Notifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Once); + } + + [TestMethod] + public void HotReload_WhenNotifierThrows_PreservesPublicationAndContinuesNotifying() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(); + ThrowingNotifier throwingNotifier = new(); + Mock healthyNotifier = new(); + TestContext context = CreateContextWithNotifiers( + () => currentConfig, + healthyNotifier, + new IMcpToolListChangedNotifier[] { throwingNotifier, healthyNotifier.Object }); + context.Service.EnsureInitialized(); + + currentConfig = CreateRuntimeConfig(("GetBook", "New tool")); + RaiseRegistryChanged(context.HotReloadEventHandler); + + Assert.IsTrue(context.Registry.TryGetTool("get_book", out _)); + Assert.AreEqual(1, throwingNotifier.CallCount); + healthyNotifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Once); + } + + [TestMethod] + public async Task HotReload_WhenNotifierBlocks_DoesNotBlockPublicationOrLaterHandlers() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(); + BlockingStringWriter output = new(); + using McpStdoutWriter stdoutWriter = new(output); + McpStdioToolListChangedNotifier notifier = new(stdoutWriter); + notifier.MarkInitialized(); + ManualResetEventSlim laterHandlerCalled = new(); + TestContext context = CreateContextWithNotifiers( + () => currentConfig, + new Mock(), + new IMcpToolListChangedNotifier[] { notifier }); + context.HotReloadEventHandler.Subscribe( + GRAPHQL_SCHEMA_EVICTION_ON_CONFIG_CHANGED, + (_, _) => laterHandlerCalled.Set()); + context.Service.EnsureInitialized(); + + currentConfig = CreateRuntimeConfig(("FirstTool", "First generation")); + TestRuntimeConfigLoader loader = new(context.HotReloadEventHandler) + { + RuntimeConfig = currentConfig + }; + Task firstRefresh = Task.Run(loader.RaiseConfigChanged); + + try + { + Assert.IsTrue( + output.WriteEntered.Wait(TimeSpan.FromSeconds(5)), + "The stdio notification worker did not reach the blocking writer."); + Assert.IsTrue( + laterHandlerCalled.Wait(TimeSpan.FromSeconds(5)), + "A blocked transport must not prevent later ordered hot-reload handlers."); + Assert.IsTrue(context.Registry.TryGetTool("first_tool", out _)); + } + finally + { + output.ReleaseWrite.Set(); + await firstRefresh.WaitAsync(TimeSpan.FromSeconds(5)); + Assert.IsTrue( + output.LineWritten.Wait(TimeSpan.FromSeconds(5)), + "The queued notification did not finish after stdout resumed."); + } + } + + [TestMethod] + public void HotReload_AfterRejectedCandidate_RecoversOnNextConfig() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(); + TestMcpTool builtIn = new("read_records", ToolType.BuiltIn); + TestContext context = CreateContext(() => currentConfig, builtIn); + context.Service.EnsureInitialized(); + + currentConfig = CreateRuntimeConfig(("ReadRecords", "Conflicting custom tool")); + RaiseRegistryChanged(context.HotReloadEventHandler); + Assert.AreSame(builtIn, GetRequiredTool(context.Registry, "read_records")); + Assert.IsFalse(context.Registry.TryGetTool("get_book", out _)); + + currentConfig = CreateRuntimeConfig(("GetBook", "Recovered tool")); + RaiseRegistryChanged(context.HotReloadEventHandler); + + Assert.IsTrue(context.Registry.TryGetTool("get_book", out _)); + Assert.AreEqual( + "Recovered tool", + context.Registry.GetAdvertisedTools().Single(tool => tool.Name == "get_book").Description); + context.Notifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Once); + } + + [TestMethod] + public void HotReload_WithSuccessiveConfigurations_PublishesLatestGeneration() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(); + TestContext context = CreateContext(() => currentConfig); + context.Service.EnsureInitialized(); + + currentConfig = CreateRuntimeConfig(("FirstTool", "First generation")); + RaiseRegistryChanged(context.HotReloadEventHandler); + currentConfig = CreateRuntimeConfig(("LatestTool", "Latest generation")); + RaiseRegistryChanged(context.HotReloadEventHandler); + + Assert.IsFalse(context.Registry.TryGetTool("first_tool", out _)); + Assert.IsTrue(context.Registry.TryGetTool("latest_tool", out _)); + CollectionAssert.AreEqual( + new[] { "latest_tool" }, + context.Registry.GetAdvertisedTools().Select(tool => tool.Name).ToArray()); + context.Notifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Exactly(2)); + } + + [TestMethod] + public void HotReload_DiscardsCandidateWhenNewerConfigBecomesActive() + { + RuntimeConfig initialConfig = CreateRuntimeConfig(); + RuntimeConfig candidateConfig = CreateRuntimeConfig(); + RuntimeConfig newerConfig = CreateRuntimeConfig(("GetBook", "Newer config")); + RuntimeConfig currentConfig = initialConfig; + int metadataReadCount = 0; + TestMcpTool builtIn = new( + "read_records", + ToolType.BuiltIn, + metadataFactory: () => + { + metadataReadCount++; + if (metadataReadCount == 2) + { + currentConfig = newerConfig; + } + + return CreateMetadata("read_records", "Built-in tool"); + }); + TestContext context = CreateContext(() => currentConfig, builtIn); + context.Service.EnsureInitialized(); + + currentConfig = candidateConfig; + RaiseRegistryChanged(context.HotReloadEventHandler); + + Assert.AreSame(builtIn, GetRequiredTool(context.Registry, "read_records")); + Assert.IsFalse(context.Registry.TryGetTool("get_book", out _)); + context.Notifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Never); + } + + [TestMethod] + public void RuntimeConfigLoader_RaisesMcpEventAfterDependenciesAndBeforeGraphQL() + { + List events = new(); + HotReloadEventHandler hotReloadEventHandler = new(); + hotReloadEventHandler.Subscribe( + METADATA_PROVIDER_FACTORY_ON_CONFIG_CHANGED, + (_, _) => events.Add(METADATA_PROVIDER_FACTORY_ON_CONFIG_CHANGED)); + hotReloadEventHandler.Subscribe( + AUTHZ_RESOLVER_ON_CONFIG_CHANGED, + (_, _) => events.Add(AUTHZ_RESOLVER_ON_CONFIG_CHANGED)); + hotReloadEventHandler.Subscribe( + MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED, + (_, _) => events.Add(MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED)); + hotReloadEventHandler.Subscribe( + GRAPHQL_SCHEMA_EVICTION_ON_CONFIG_CHANGED, + (_, _) => events.Add(GRAPHQL_SCHEMA_EVICTION_ON_CONFIG_CHANGED)); + + TestRuntimeConfigLoader loader = new(hotReloadEventHandler) + { + RuntimeConfig = CreateRuntimeConfig() + }; + + loader.RaiseConfigChanged(); + + CollectionAssert.AreEqual( + new[] + { + METADATA_PROVIDER_FACTORY_ON_CONFIG_CHANGED, + AUTHZ_RESOLVER_ON_CONFIG_CHANGED, + MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED, + GRAPHQL_SCHEMA_EVICTION_ON_CONFIG_CHANGED + }, + events); + } + + private static TestContext CreateContext( + Func getConfig, + params IMcpTool[] builtInTools) + { + Mock notifier = new(); + return CreateContextWithNotifiers( + getConfig, + notifier, + new[] { notifier.Object }, + builtInTools); + } + + private static TestContext CreateContextWithNotifiers( + Func getConfig, + Mock primaryNotifier, + IEnumerable notifiers, + params IMcpTool[] builtInTools) + { + return CreateContextCore( + getConfig, + primaryNotifier, + notifiers, + () => new Dictionary(), + builtInTools); + } + + private static TestContext CreateContextWithDatabaseMetadata( + Func getConfig, + Func> getMetadata) + { + Mock notifier = new(); + return CreateContextCore( + getConfig, + notifier, + new[] { notifier.Object }, + getMetadata, + Array.Empty()); + } + + private static TestContext CreateContextCore( + Func getConfig, + Mock primaryNotifier, + IEnumerable notifiers, + Func> getMetadata, + IEnumerable registeredTools) + { + Mock configLoader = new(null, null); + Mock configProvider = new(configLoader.Object); + configProvider.Setup(provider => provider.GetConfig()).Returns(getConfig); + + Mock sqlMetadataProvider = new(); + sqlMetadataProvider + .SetupGet(provider => provider.EntityToDatabaseObject) + .Returns(getMetadata); + Mock metadataProviderFactory = new(); + metadataProviderFactory + .Setup(factory => factory.GetMetadataProvider(It.IsAny())) + .Returns(sqlMetadataProvider.Object); + + McpToolRegistry registry = new(); + HotReloadEventHandler hotReloadEventHandler = new(); + Mock> logger = new(); + McpToolRegistryRefreshService service = new( + configProvider.Object, + registeredTools, + registry, + metadataProviderFactory.Object, + notifiers, + logger.Object, + hotReloadEventHandler); + + return new TestContext( + service, + registry, + primaryNotifier, + hotReloadEventHandler, + logger); + } + + private static string[] GetAdvertisedParameterNames(McpToolRegistry registry) + { + return registry.GetAdvertisedTools() + .Single() + .InputSchema + .GetProperty("properties") + .EnumerateObject() + .Select(property => property.Name) + .ToArray(); + } + + private static Dictionary CreateStoredProcedureMetadata( + string parameterName, + string description) + { + DatabaseStoredProcedure storedProcedure = new("dbo", "test_procedure") + { + SourceType = EntitySourceType.StoredProcedure, + StoredProcedureDefinition = new StoredProcedureDefinition + { + Parameters = new Dictionary + { + [parameterName] = new ParameterDefinition + { + Name = parameterName, + Description = description, + Required = true, + SystemType = typeof(string) + } + } + } + }; + + return new Dictionary + { + ["GetBook"] = storedProcedure + }; + } + + private static RuntimeConfig CreateRuntimeConfig( + params (string EntityName, string Description)[] customTools) + { + Dictionary entities = customTools.ToDictionary( + item => item.EntityName, + item => new Entity( + Source: new("test_procedure", EntitySourceType.StoredProcedure, Parameters: null, KeyFields: null), + GraphQL: new(item.EntityName, item.EntityName), + Rest: new(Enabled: true), + Fields: null, + Permissions: new[] + { + new EntityPermission( + Role: "anonymous", + Actions: new[] + { + new EntityAction(Action: EntityActionOperation.Execute, Fields: null, Policy: null) + }) + }, + Relationships: null, + Mappings: null, + Description: item.Description, + Mcp: new EntityMcpOptions(customToolEnabled: true, dmlToolsEnabled: null))); + + return new RuntimeConfig( + Schema: "test-schema", + DataSource: new DataSource(DatabaseType.MSSQL, "", Options: null), + Runtime: new( + Rest: new(), + GraphQL: new(), + Mcp: new(Enabled: true), + Host: new(Cors: null, Authentication: null, Mode: HostMode.Development)), + Entities: new(entities)); + } + + private static RuntimeConfig CreateRuntimeConfigWithParameter(string parameterDescription) + { + Entity entity = new( + Source: new( + "test_procedure", + EntitySourceType.StoredProcedure, + Parameters: new List + { + new() + { + Name = "id", + Description = parameterDescription, + Required = true + } + }, + KeyFields: null), + GraphQL: new("GetBook", "GetBooks"), + Rest: new(Enabled: true), + Fields: null, + Permissions: new[] + { + new EntityPermission( + Role: "anonymous", + Actions: new[] + { + new EntityAction( + Action: EntityActionOperation.Execute, + Fields: null, + Policy: null) + }) + }, + Relationships: null, + Mappings: null, + Description: "Stable tool description", + Mcp: new EntityMcpOptions(customToolEnabled: true, dmlToolsEnabled: null)); + + return new RuntimeConfig( + Schema: "test-schema", + DataSource: new DataSource(DatabaseType.MSSQL, string.Empty, Options: null), + Runtime: new( + Rest: new(), + GraphQL: new(), + Mcp: new(Enabled: true), + Host: new(Cors: null, Authentication: null, Mode: HostMode.Development)), + Entities: new(new Dictionary { ["GetBook"] = entity })); + } + + private static IMcpTool GetRequiredTool(McpToolRegistry registry, string name) + { + Assert.IsTrue(registry.TryGetTool(name, out IMcpTool? tool)); + Assert.IsNotNull(tool); + return tool; + } + + private static void RaiseRegistryChanged( + HotReloadEventHandler hotReloadEventHandler) + { + hotReloadEventHandler.OnConfigChangedEvent( + hotReloadEventHandler, + new HotReloadEventArgs(MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED, string.Empty)); + } + + private static Tool CreateMetadata(string name, string description) + { + return new Tool + { + Name = name, + Description = description, + InputSchema = JsonSerializer.Deserialize("{\"type\":\"object\"}") + }; + } + + private static void VerifyLogContains( + Mock> logger, + LogLevel logLevel, + string expectedMessage) + { + logger.Verify( + value => value.Log( + logLevel, + It.IsAny(), + It.Is((state, _) => + state.ToString()!.Contains(expectedMessage, StringComparison.Ordinal)), + It.IsAny(), + (Func)It.IsAny()), + Times.Once); + } + + private sealed record TestContext( + McpToolRegistryRefreshService Service, + McpToolRegistry Registry, + Mock Notifier, + HotReloadEventHandler HotReloadEventHandler, + Mock> Logger); + + private sealed class TestMcpTool : IMcpTool + { + private readonly string _name; + private readonly Func? _metadataFactory; + + public TestMcpTool( + string name, + ToolType toolType, + Func? metadataFactory = null) + { + _name = name; + ToolType = toolType; + _metadataFactory = metadataFactory; + } + + public ToolType ToolType { get; } + + public Tool GetToolMetadata() + { + return _metadataFactory?.Invoke() ?? CreateMetadata(_name, "Test tool"); + } + + public bool IsEnabled(RuntimeConfig config) => true; + + public Task ExecuteAsync( + JsonDocument? arguments, + IServiceProvider serviceProvider, + CancellationToken cancellationToken = default) + { + throw new NotImplementedException(); + } + } + + private sealed class ThrowingNotifier : IMcpToolListChangedNotifier + { + public int CallCount { get; private set; } + + public void NotifyToolsListChanged() + { + CallCount++; + throw new InvalidOperationException("Expected notification failure."); + } + } + + private sealed class BlockingStringWriter : StringWriter + { + public ManualResetEventSlim WriteEntered { get; } = new(); + + public ManualResetEventSlim ReleaseWrite { get; } = new(); + + public ManualResetEventSlim LineWritten { get; } = new(); + + public override void WriteLine(string? value) + { + WriteEntered.Set(); + if (!ReleaseWrite.Wait(TimeSpan.FromSeconds(10))) + { + throw new TimeoutException("Timed out waiting to release the stdout write."); + } + + base.WriteLine(value); + LineWritten.Set(); + } + } + + private sealed class TestRuntimeConfigLoader : RuntimeConfigLoader + { + public TestRuntimeConfigLoader(HotReloadEventHandler handler) + : base(handler) + { + } + + public void RaiseConfigChanged() + { + SignalConfigChanged(); + } + + public override bool TryLoadKnownConfig( + [NotNullWhen(true)] out RuntimeConfig? config, + bool replaceEnvVar = false) + { + config = RuntimeConfig; + return config is not null; + } + + public override string GetPublishedDraftSchemaLink() + { + return string.Empty; + } + } + } +} diff --git a/src/Service.Tests/Mcp/McpToolRegistryTests.cs b/src/Service.Tests/Mcp/McpToolRegistryTests.cs index d8c5dc0b59..e35f64ad3f 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryTests.cs @@ -2,10 +2,11 @@ // Licensed under the MIT License. using System; +using System.Collections.Concurrent; using System.Collections.Generic; using System.Linq; -using System.Net; using System.Text.Json; +using System.Text.Json.Nodes; using System.Threading; using System.Threading.Tasks; using Azure.DataApiBuilder.Config.ObjectModel; @@ -25,343 +26,545 @@ namespace Azure.DataApiBuilder.Service.Tests.Mcp public class McpToolRegistryTests { /// - /// Test that registering multiple tools with unique names succeeds. + /// Test that TryGetTool returns false for non-existent tool. /// [TestMethod] - public void RegisterTool_WithMultipleUniqueNames_Succeeds() + public void TryGetTool_WithNonExistentName_ReturnsFalse() { // Arrange McpToolRegistry registry = new(); - IMcpTool tool1 = new MockMcpTool("tool_one", ToolType.BuiltIn); - IMcpTool tool2 = new MockMcpTool("tool_two", ToolType.Custom); - IMcpTool tool3 = new MockMcpTool("tool_three", ToolType.BuiltIn); - - // Act & Assert - should not throw - registry.RegisterTool(tool1); - registry.RegisterTool(tool2); - registry.RegisterTool(tool3); - - // Verify all tools were registered - Assert.IsTrue(registry.TryGetTool("tool_one", out _)); - Assert.IsTrue(registry.TryGetTool("tool_two", out _)); - Assert.IsTrue(registry.TryGetTool("tool_three", out _)); + + // Act + bool found = registry.TryGetTool("non_existent_tool", out IMcpTool? tool); + + // Assert + Assert.IsFalse(found); + Assert.IsNull(tool); } /// - /// Test that registering duplicate tools of the same type throws an exception. - /// Validates that both built-in and custom tools enforce name uniqueness within their own type. + /// Test edge case: empty tool name should throw exception. /// - [DataTestMethod] - [DataRow(ToolType.BuiltIn, "duplicate_tool", "built-in", DisplayName = "Duplicate Built-In Tools")] - [DataRow(ToolType.Custom, "my_custom_tool", "custom", DisplayName = "Duplicate Custom Tools")] - public void RegisterTool_WithDuplicateSameType_ThrowsException( - ToolType toolType, - string toolName, - string expectedToolTypeText) + [TestMethod] + public void ReplaceAll_WithEmptyToolName_ThrowsException() { // Arrange McpToolRegistry registry = new(); - IMcpTool tool1 = new MockMcpTool(toolName, toolType); - IMcpTool tool2 = new MockMcpTool(toolName, toolType); - - // Act - Register first tool - registry.RegisterTool(tool1); + IMcpTool tool = new MockMcpTool("", ToolType.BuiltIn); - // Assert - Second registration should throw + // Assert - Empty tool names should be rejected DataApiBuilderException exception = Assert.ThrowsException( - () => registry.RegisterTool(tool2) + () => registry.ReplaceAll(new[] { tool }, CreateRuntimeConfig()) ); - // Verify exception details - Assert.IsTrue(exception.Message.Contains($"Duplicate MCP tool name '{toolName}' detected")); - Assert.IsTrue(exception.Message.Contains($"{expectedToolTypeText} tool with this name is already registered")); - Assert.IsTrue(exception.Message.Contains($"Cannot register {expectedToolTypeText} tool with the same name")); + Assert.IsTrue(exception.Message.Contains("cannot be null, empty, or whitespace")); Assert.AreEqual(DataApiBuilderException.SubStatusCodes.ErrorInInitialization, exception.SubStatusCode); - Assert.AreEqual(HttpStatusCode.ServiceUnavailable, exception.StatusCode); } /// - /// Test that registering tools with conflicting names across different types throws an exception. - /// Validates that tool names must be unique across all tool types (built-in and custom). + /// Test that leading/trailing whitespace is rejected rather than producing a lookup key + /// that differs from the advertised tool name. /// - [DataTestMethod] - [DataRow("create_record", ToolType.BuiltIn, ToolType.Custom, "built-in", "custom", DisplayName = "Built-In then Custom conflict")] - [DataRow("read_records", ToolType.BuiltIn, ToolType.Custom, "built-in", "custom", DisplayName = "Built-In then Custom conflict (read_records)")] - [DataRow("my_stored_proc", ToolType.Custom, ToolType.BuiltIn, "custom", "built-in", DisplayName = "Custom then Built-In conflict")] - public void RegisterTool_WithCrossTypeConflict_ThrowsException( - string toolName, - ToolType firstToolType, - ToolType secondToolType, - string expectedExistingType, - string expectedNewType) + [TestMethod] + public void ReplaceAll_WithLeadingTrailingWhitespace_ThrowsException() { - // Arrange McpToolRegistry registry = new(); - IMcpTool existingTool = new MockMcpTool(toolName, firstToolType); - IMcpTool conflictingTool = new MockMcpTool(toolName, secondToolType); + IMcpTool tool = new MockMcpTool(" my_tool ", ToolType.Custom); - // Act - Register first tool - registry.RegisterTool(existingTool); - - // Assert - Second tool registration should throw DataApiBuilderException exception = Assert.ThrowsException( - () => registry.RegisterTool(conflictingTool) - ); + () => registry.ReplaceAll(new[] { tool }, CreateRuntimeConfig())); - // Verify exception details - Assert.IsTrue(exception.Message.Contains($"Duplicate MCP tool name '{toolName}' detected")); - Assert.IsTrue(exception.Message.Contains($"{expectedExistingType} tool with this name is already registered")); - Assert.IsTrue(exception.Message.Contains($"Cannot register {expectedNewType} tool with the same name")); - Assert.IsTrue(exception.Message.Contains("Tool names must be unique across all tool types")); - Assert.AreEqual(DataApiBuilderException.SubStatusCodes.ErrorInInitialization, exception.SubStatusCode); - Assert.AreEqual(HttpStatusCode.ServiceUnavailable, exception.StatusCode); + StringAssert.Contains(exception.Message, "leading or trailing whitespace"); + Assert.IsFalse(registry.TryGetTool("my_tool", out _)); } /// - /// Test that tool name comparison is case-sensitive. - /// Tools with different casing should not be allowed. + /// Replacing the registry publishes a complete, deterministically ordered snapshot and + /// removes tools that belonged only to the previous generation. /// [TestMethod] - public void RegisterTool_WithDifferentCasing_ThrowsException() + public void ReplaceAll_PublishesCompleteOrderedSnapshot() { - // Arrange McpToolRegistry registry = new(); - IMcpTool tool1 = new MockMcpTool("my_tool", ToolType.BuiltIn); - IMcpTool tool2 = new MockMcpTool("My_Tool", ToolType.Custom); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll( + new[] { new MockMcpTool("old_tool", ToolType.Custom) }, + config); - // Act - Register first tool - registry.RegisterTool(tool1); + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new IMcpTool[] + { + new MockMcpTool("z_tool", ToolType.Custom), + new MockMcpTool("A_tool", ToolType.BuiltIn) + }, + config); + + Assert.IsFalse(registry.TryGetTool("old_tool", out _)); + Assert.IsTrue(registry.TryGetTool("a_TOOL", out _)); + Assert.IsTrue(registry.TryGetTool("z_tool", out _)); + CollectionAssert.AreEqual( + new[] { "A_tool", "z_tool" }, + registry.GetAdvertisedTools().Select(tool => tool.Name).ToArray()); + Assert.AreEqual(2, result.Version); + Assert.IsTrue(result.DiscoveryChanged); + Assert.AreEqual(2, result.RegisteredToolCount); + Assert.AreEqual(2, result.AdvertisedToolCount); + } - // Assert - Case-insensitive duplicate should throw - DataApiBuilderException exception = Assert.ThrowsException( - () => registry.RegisterTool(tool2) - ); + /// + /// Every name returned by discovery resolves against the exact same registry generation. + /// + [TestMethod] + public void ReplaceAll_EveryAdvertisedNameIsCallable() + { + McpToolRegistry registry = new(); + registry.ReplaceAll( + new IMcpTool[] + { + new MockMcpTool("A_tool", ToolType.BuiltIn), + new MockMcpTool("z_tool", ToolType.Custom) + }, + CreateRuntimeConfig()); - Assert.IsTrue(exception.Message.Contains("Duplicate MCP tool name")); - Assert.AreEqual(DataApiBuilderException.SubStatusCodes.ErrorInInitialization, exception.SubStatusCode); + foreach (Tool advertisedTool in registry.GetAdvertisedTools()) + { + Assert.IsTrue( + registry.TryGetTool(advertisedTool.Name, out IMcpTool? callableTool), + $"Advertised MCP tool '{advertisedTool.Name}' must be callable by that exact name."); + Assert.IsNotNull(callableTool); + } } /// - /// Test that registering the same tool instance twice is silently ignored (idempotent). - /// This supports stdio mode where both McpToolRegistryInitializer and McpStdioHelper may register the same tools. + /// A candidate containing a duplicate name is rejected before publication, leaving the + /// complete previous snapshot active. /// [TestMethod] - public void RegisterTool_SameInstanceTwice_IsIdempotent() + public void ReplaceAll_WithDuplicateName_PreservesPreviousSnapshot() { - // Arrange McpToolRegistry registry = new(); - IMcpTool tool = new MockMcpTool("my_tool", ToolType.BuiltIn); - - // Act - Register the same instance twice - registry.RegisterTool(tool); - registry.RegisterTool(tool); + RuntimeConfig config = CreateRuntimeConfig(); + IMcpTool previousTool = new MockMcpTool("previous_tool", ToolType.BuiltIn); + registry.ReplaceAll(new[] { previousTool }, config); - // Assert - Tool should be registered only once - Assert.IsTrue(registry.TryGetTool("my_tool", out _)); + Assert.ThrowsException(() => registry.ReplaceAll( + new IMcpTool[] + { + new MockMcpTool("duplicate", ToolType.BuiltIn), + new MockMcpTool("DUPLICATE", ToolType.Custom) + }, + config)); + + Assert.IsTrue(registry.TryGetTool("previous_tool", out IMcpTool? actualTool)); + Assert.AreSame(previousTool, actualTool); + Assert.IsFalse(registry.TryGetTool("duplicate", out _)); + CollectionAssert.AreEqual( + new[] { "previous_tool" }, + registry.GetAdvertisedTools().Select(tool => tool.Name).ToArray()); } /// - /// Test that registering a different instance with the same name throws an exception, - /// even though a same-instance re-registration would be allowed. + /// Replacing tool instances with semantically identical discovery metadata advances the + /// registry generation without reporting a client-visible discovery change. /// [TestMethod] - public void RegisterTool_DifferentInstanceSameName_ThrowsException() + public void ReplaceAll_WithEquivalentMetadata_DoesNotReportDiscoveryChange() { - // Arrange McpToolRegistry registry = new(); - IMcpTool tool1 = new MockMcpTool("my_tool", ToolType.BuiltIn); - IMcpTool tool2 = new MockMcpTool("my_tool", ToolType.BuiltIn); - - // Act - Register first instance - registry.RegisterTool(tool1); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, description: "Same description") }, + config); - // Assert - Different instance with same name should throw - DataApiBuilderException exception = Assert.ThrowsException( - () => registry.RegisterTool(tool2) - ); + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, description: "Same description") }, + config); - Assert.IsTrue(exception.Message.Contains("Duplicate MCP tool name 'my_tool' detected")); - Assert.AreEqual(DataApiBuilderException.SubStatusCodes.ErrorInInitialization, exception.SubStatusCode); + Assert.AreEqual(2, result.Version); + Assert.IsFalse(result.DiscoveryChanged); } /// - /// Test that TryGetTool returns false for non-existent tool. + /// Object property order is not semantically meaningful and must not trigger discovery + /// invalidation when equivalent metadata is rebuilt in a different insertion order. /// [TestMethod] - public void TryGetTool_WithNonExistentName_ReturnsFalse() + public void ReplaceAll_WithEquivalentSchemaPropertyOrder_DoesNotReportDiscoveryChange() { - // Arrange + const string SCHEMA_AB = + "{\"type\":\"object\",\"properties\":{\"a\":{\"type\":\"string\"},\"b\":{\"type\":\"integer\"}}}"; + const string SCHEMA_BA = + "{\"properties\":{\"b\":{\"type\":\"integer\"},\"a\":{\"type\":\"string\"}},\"type\":\"object\"}"; McpToolRegistry registry = new(); - registry.RegisterTool(new MockMcpTool("existing_tool", ToolType.BuiltIn)); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, inputSchemaJson: SCHEMA_AB) }, + config); - // Act - bool found = registry.TryGetTool("non_existent_tool", out IMcpTool? tool); + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, inputSchemaJson: SCHEMA_BA) }, + config); - // Assert - Assert.IsFalse(found); - Assert.IsNull(tool); + Assert.IsFalse(result.DiscoveryChanged); } /// - /// Test edge case: empty tool name should throw exception. + /// JSON Schema string arrays used as sets do not change schema semantics when rebuilt in a + /// different order and therefore must not invalidate client discovery. /// [TestMethod] - public void RegisterTool_WithEmptyToolName_ThrowsException() + public void ReplaceAll_WithEquivalentSchemaSetArrayOrder_DoesNotReportDiscoveryChange() { - // Arrange + const string SCHEMA_AB = + "{\"type\":\"object\",\"properties\":{" + + "\"a\":{\"type\":[\"string\",\"null\"],\"enum\":[\"alpha\",\"beta\"]}," + + "\"b\":{\"type\":\"integer\"}},\"required\":[\"a\",\"b\"]}"; + const string SCHEMA_BA = + "{\"required\":[\"b\",\"a\"],\"properties\":{" + + "\"b\":{\"type\":\"integer\"}," + + "\"a\":{\"enum\":[\"beta\",\"alpha\"],\"type\":[\"null\",\"string\"]}}," + + "\"type\":\"object\"}"; McpToolRegistry registry = new(); - IMcpTool tool = new MockMcpTool("", ToolType.BuiltIn); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, inputSchemaJson: SCHEMA_AB) }, + config); - // Assert - Empty tool names should be rejected - DataApiBuilderException exception = Assert.ThrowsException( - () => registry.RegisterTool(tool) - ); + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, inputSchemaJson: SCHEMA_BA) }, + config); - Assert.IsTrue(exception.Message.Contains("cannot be null, empty, or whitespace")); - Assert.AreEqual(DataApiBuilderException.SubStatusCodes.ErrorInInitialization, exception.SubStatusCode); + Assert.IsFalse(result.DiscoveryChanged); } /// - /// Test realistic scenario with actual built-in tool names. + /// Primitive arrays are not universally sets. Reordering an array-valued default changes + /// the advertised default instance and must remain a discovery change. /// [TestMethod] - public void RegisterTool_WithRealisticBuiltInToolNames_DetectsDuplicates() + public void ReplaceAll_WithReorderedArrayDefault_ReportsDiscoveryChange() { - // Arrange + const string SCHEMA_AB = + "{\"type\":\"object\",\"properties\":{\"values\":{" + + "\"type\":\"array\",\"items\":{\"type\":\"string\"}," + + "\"default\":[\"a\",\"b\"]}}}"; + const string SCHEMA_BA = + "{\"type\":\"object\",\"properties\":{\"values\":{" + + "\"type\":\"array\",\"items\":{\"type\":\"string\"}," + + "\"default\":[\"b\",\"a\"]}}}"; McpToolRegistry registry = new(); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, inputSchemaJson: SCHEMA_AB) }, + config); - // Simulate registering built-in tools - registry.RegisterTool(new MockMcpTool("create_record", ToolType.BuiltIn)); - registry.RegisterTool(new MockMcpTool("read_records", ToolType.BuiltIn)); - registry.RegisterTool(new MockMcpTool("update_record", ToolType.BuiltIn)); - registry.RegisterTool(new MockMcpTool("delete_record", ToolType.BuiltIn)); - registry.RegisterTool(new MockMcpTool("describe_entities", ToolType.BuiltIn)); + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, inputSchemaJson: SCHEMA_BA) }, + config); - // Try to register a custom tool with a conflicting name - IMcpTool customTool = new MockMcpTool("read_records", ToolType.Custom); + Assert.IsTrue(result.DiscoveryChanged); + } - // Assert - Should throw - DataApiBuilderException exception = Assert.ThrowsException( - () => registry.RegisterTool(customTool) - ); + /// + /// Schema instance values and unknown keywords contain data, not schemas. Even schema-like + /// names nested in those values must preserve array order when comparing discovery metadata. + /// + [DataTestMethod] + [DynamicData(nameof(SchemaInstanceDataCases), DynamicDataSourceType.Property)] + public void ReplaceAll_WithReorderedArrayInSchemaInstanceData_ReportsDiscoveryChange( + string instanceKeyword, + string dataPropertyName, + bool useOutputSchema, + bool nested) + { + string schemaAB = CreateSchemaWithInstanceData(instanceKeyword, dataPropertyName, nested, reverse: false); + string schemaBA = CreateSchemaWithInstanceData(instanceKeyword, dataPropertyName, nested, reverse: true); + McpToolRegistry registry = new(); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll(new[] { CreateSchemaTool(schemaAB, useOutputSchema) }, config); + + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { CreateSchemaTool(schemaBA, useOutputSchema) }, config); - Assert.IsTrue(exception.Message.Contains("read_records")); - Assert.IsTrue(exception.Message.Contains("built-in tool")); + Tool advertised = registry.GetAdvertisedTools().Single(); + JsonElement advertisedSchema = useOutputSchema ? advertised.OutputSchema!.Value : advertised.InputSchema; + Assert.AreEqual(schemaBA, advertisedSchema.GetRawText(), "Serving must preserve the updated data array order."); + Assert.IsTrue(result.DiscoveryChanged, "Changed instance data must invalidate cached discovery metadata."); } /// - /// Test that registering a tool with leading/trailing whitespace in the name is treated as a duplicate of the trimmed name. - /// Note: during tool registration, the registry should trim whitespace and detect duplicates accordingly. + /// Genuine subschemas still canonicalize set-like keywords, including when a property or + /// definition happens to be named default, const, or examples. /// - [TestMethod] - public void RegisterTool_WithLeadingTrailingWhitespace_DetectsDuplicate() + [DataTestMethod] + [DynamicData(nameof(NestedSchemaCases), DynamicDataSourceType.Property)] + public void ReplaceAll_WithEquivalentNestedSchemaSets_DoesNotReportDiscoveryChange( + string schemaTemplate, + bool useOutputSchema) { - // Arrange + const string SCHEMA_AB = + "{\"type\":\"object\",\"properties\":{" + + "\"a\":{\"type\":[\"string\",\"null\"],\"enum\":[\"alpha\",\"beta\"]}," + + "\"b\":{\"type\":\"integer\"}},\"required\":[\"a\",\"b\"]}"; + const string SCHEMA_BA = + "{\"required\":[\"b\",\"a\"],\"properties\":{" + + "\"b\":{\"type\":\"integer\"}," + + "\"a\":{\"enum\":[\"beta\",\"alpha\"],\"type\":[\"null\",\"string\"]}}," + + "\"type\":\"object\"}"; + string schemaAB = schemaTemplate.Replace("$SCHEMA", SCHEMA_AB, StringComparison.Ordinal); + string schemaBA = schemaTemplate.Replace("$SCHEMA", SCHEMA_BA, StringComparison.Ordinal); McpToolRegistry registry = new(); - IMcpTool tool1 = new MockMcpTool("my_tool", ToolType.BuiltIn); - IMcpTool tool2 = new MockMcpTool(" my_tool ", ToolType.Custom); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll(new[] { CreateSchemaTool(schemaAB, useOutputSchema) }, config); - // Act - registry.RegisterTool(tool1); + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { CreateSchemaTool(schemaBA, useOutputSchema) }, config); - // Assert - trimmed name should collide - Assert.ThrowsException( - () => registry.RegisterTool(tool2) - ); + Assert.IsFalse(result.DiscoveryChanged); + Tool advertised = registry.GetAdvertisedTools().Single(); + JsonElement advertisedSchema = useOutputSchema ? advertised.OutputSchema!.Value : advertised.InputSchema; + Assert.AreEqual(schemaBA, advertisedSchema.GetRawText(), "Canonicalization must not reorder the served schema."); + } + + [DataTestMethod] + [DataRow("items")] + [DataRow("prefixItems")] + public void ReplaceAll_WithReorderedSchemaTuple_ReportsDiscoveryChange(string keyword) + { + string schemaAB = "{\"type\":\"object\",\"properties\":{\"values\":{\"type\":\"array\",\"" + + keyword + "\":[{\"type\":\"string\"},{\"type\":\"integer\"}]}}}"; + string schemaBA = "{\"type\":\"object\",\"properties\":{\"values\":{\"type\":\"array\",\"" + + keyword + "\":[{\"type\":\"integer\"},{\"type\":\"string\"}]}}}"; + McpToolRegistry registry = new(); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll(new[] { CreateSchemaTool(schemaAB, useOutputSchema: false) }, config); + + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { CreateSchemaTool(schemaBA, useOutputSchema: false) }, config); + + Assert.IsTrue(result.DiscoveryChanged, "Tuple positions are significant even though their elements are schemas."); + } + + [DataTestMethod] + [DataRow("inputSchema")] + [DataRow("outputSchema")] + public void ReplaceAll_WithSchemaNamedPropertyInMetadata_ReportsDiscoveryChange(string propertyName) + { + string metadataAB = "{\"name\":\"same_tool\",\"inputSchema\":{\"type\":\"object\"},\"_meta\":{\"" + + propertyName + "\":{\"enum\":[\"a\",\"b\"]}}}"; + string metadataBA = "{\"name\":\"same_tool\",\"inputSchema\":{\"type\":\"object\"},\"_meta\":{\"" + + propertyName + "\":{\"enum\":[\"b\",\"a\"]}}}"; + McpToolRegistry registry = new(); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll( + new[] { new RetainedMetadataMcpTool(JsonSerializer.Deserialize(metadataAB)!) }, config); + + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { new RetainedMetadataMcpTool(JsonSerializer.Deserialize(metadataBA)!) }, config); + + JsonElement advertised = JsonSerializer.SerializeToElement(registry.GetAdvertisedTools().Single()); + Assert.AreEqual("b", advertised.GetProperty("_meta").GetProperty(propertyName).GetProperty("enum")[0].GetString()); + Assert.IsTrue(result.DiscoveryChanged, "Only the tool's actual input/output schema properties introduce schemas."); } /// - /// Parameterized test verifying GetEnabledTools returns only enabled tools. + /// Canonical property sorting is used only for change detection. The discovery payload + /// preserves schema-property insertion order for clients that render parameters in wire + /// order even though JSON Schema does not assign that order semantic meaning. /// - [DataTestMethod] - [DataRow(1, 1, DisplayName = "Mixed: 1 enabled, 1 disabled → returns 1")] - [DataRow(3, 0, DisplayName = "All enabled → returns all")] - [DataRow(0, 2, DisplayName = "All disabled → returns 0")] - public void GetEnabledTools_ReturnsCorrectCount(int enabledCount, int disabledCount) + [TestMethod] + public void GetAdvertisedTools_PreservesInputSchemaPropertyOrder() { - // Arrange + const string SCHEMA = + "{\"type\":\"object\",\"properties\":{" + + "\"second\":{\"type\":\"string\"}," + + "\"first\":{\"type\":\"integer\"}}," + + "\"required\":[\"second\",\"first\"]}"; McpToolRegistry registry = new(); - for (int i = 0; i < enabledCount; i++) - { - registry.RegisterTool(new MockMcpTool($"enabled_{i}", ToolType.BuiltIn, isEnabledFunc: _ => true)); - } + registry.ReplaceAll( + new[] { new MockMcpTool("ordered_tool", ToolType.Custom, inputSchemaJson: SCHEMA) }, + CreateRuntimeConfig()); + + string[] propertyNames = registry.GetAdvertisedTools() + .Single() + .InputSchema + .GetProperty("properties") + .EnumerateObject() + .Select(property => property.Name) + .ToArray(); + + CollectionAssert.AreEqual(new[] { "second", "first" }, propertyNames); + + string[] requiredNames = registry.GetAdvertisedTools() + .Single() + .InputSchema + .GetProperty("required") + .EnumerateArray() + .Select(item => item.GetString()!) + .ToArray(); + + CollectionAssert.AreEqual(new[] { "second", "first" }, requiredNames); + } - for (int i = 0; i < disabledCount; i++) + /// + /// A real input-schema change remains client-visible after canonicalization. + /// + [TestMethod] + public void ReplaceAll_WithChangedInputSchema_ReportsDiscoveryChange() + { + McpToolRegistry registry = new(); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll( + new[] + { + new MockMcpTool( + "same_tool", + ToolType.Custom, + inputSchemaJson: "{\"type\":\"object\",\"properties\":{\"id\":{\"type\":\"string\"}}}") + }, + config); + + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] + { + new MockMcpTool( + "same_tool", + ToolType.Custom, + inputSchemaJson: "{\"type\":\"object\",\"properties\":{\"id\":{\"type\":\"integer\"}}}") + }, + config); + + Assert.IsTrue(result.DiscoveryChanged); + } + + /// + /// Published metadata is isolated both from the tool-owned source object and from callers + /// mutating a value returned by the public snapshot accessor. + /// + [TestMethod] + public void ReplaceAll_DefensivelyClonesPublishedMetadata() + { + Tool retainedMetadata = new() { - registry.RegisterTool(new MockMcpTool($"disabled_{i}", ToolType.BuiltIn, isEnabledFunc: _ => false)); - } + Name = "isolated_tool", + Description = "Original description", + InputSchema = JsonSerializer.Deserialize("{\"type\":\"object\"}") + }; + McpToolRegistry registry = new(); + registry.ReplaceAll( + new[] { new RetainedMetadataMcpTool(retainedMetadata) }, + CreateRuntimeConfig()); + + retainedMetadata.Description = "Mutated by tool"; + Tool returnedMetadata = registry.GetAdvertisedTools().Single(); + Assert.AreEqual("Original description", returnedMetadata.Description); + + returnedMetadata.Description = "Mutated by caller"; + Assert.AreEqual( + "Original description", + registry.GetAdvertisedTools().Single().Description); + } + /// + /// A metadata-only change is reported so connected clients can refresh their cached list. + /// + [TestMethod] + public void ReplaceAll_WithChangedDescription_ReportsDiscoveryChange() + { + McpToolRegistry registry = new(); RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, description: "Old description") }, + config); - // Act - List result = registry.GetEnabledTools(config).ToList(); + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, description: "New description") }, + config); - // Assert - Assert.AreEqual(enabledCount, result.Count); + Assert.IsTrue(result.DiscoveryChanged); + Assert.AreEqual("New description", registry.GetAdvertisedTools().Single().Description); } /// - /// Test that GetEnabledTools passes the RuntimeConfig to IsEnabled so tools - /// can check DmlToolsConfig flags. + /// Advertised metadata and callable lookup state are built from one candidate generation. + /// Disabled built-ins remain callable so execution can return the existing structured + /// tool-disabled response, but they are absent from discovery. /// [TestMethod] - public void GetEnabledTools_PassesConfigToIsEnabled() + public void ReplaceAll_CapturesVisibilityFromCandidateConfig() { - // Arrange McpToolRegistry registry = new(); - - // This tool checks config.McpDmlTools?.CreateRecord IMcpTool configAwareTool = new MockMcpTool( - "create_record", ToolType.BuiltIn, + "create_record", + ToolType.BuiltIn, isEnabledFunc: config => config.McpDmlTools?.CreateRecord == true); - registry.RegisterTool(configAwareTool); + RuntimeConfig disabledConfig = CreateRuntimeConfig(new DmlToolsConfig(createRecord: false)); + registry.ReplaceAll(new[] { configAwareTool }, disabledConfig); - // Config with create-record disabled - DmlToolsConfig disabledConfig = new(createRecord: false); - RuntimeConfig configDisabled = CreateRuntimeConfig(disabledConfig); + Assert.AreEqual(0, registry.GetAdvertisedTools().Count); + Assert.IsTrue(registry.TryGetTool("create_record", out _)); - // Config with create-record enabled - DmlToolsConfig enabledConfig = new(createRecord: true); - RuntimeConfig configEnabled = CreateRuntimeConfig(enabledConfig); + RuntimeConfig enabledConfig = CreateRuntimeConfig(new DmlToolsConfig(createRecord: true)); + registry.ReplaceAll(new[] { configAwareTool }, enabledConfig); - // Act & Assert - disabled - List disabledTools = registry.GetEnabledTools(configDisabled).ToList(); - Assert.AreEqual(0, disabledTools.Count); - - // Act & Assert - enabled - List enabledTools = registry.GetEnabledTools(configEnabled).ToList(); - Assert.AreEqual(1, enabledTools.Count); - Assert.AreEqual("create_record", enabledTools[0].Name); + Assert.AreEqual(1, registry.GetAdvertisedTools().Count); } /// - /// Test that GetEnabledTools correctly filters a mix of built-in and custom tools. - /// Custom tools (always enabled) should remain while disabled built-in tools are excluded. + /// Concurrent readers see only a complete old or complete new advertised snapshot while + /// registry generations are repeatedly replaced. /// [TestMethod] - public void GetEnabledTools_MixedBuiltInAndCustomTools() + public void ReplaceAll_WithConcurrentReaders_NeverExposesPartialSnapshot() { - // Arrange McpToolRegistry registry = new(); - registry.RegisterTool(new MockMcpTool("describe_entities", ToolType.BuiltIn, isEnabledFunc: _ => true)); - registry.RegisterTool(new MockMcpTool("create_record", ToolType.BuiltIn, isEnabledFunc: _ => false)); - registry.RegisterTool(new MockMcpTool("delete_record", ToolType.BuiltIn, isEnabledFunc: _ => false)); - registry.RegisterTool(new MockMcpTool("read_records", ToolType.BuiltIn, isEnabledFunc: _ => true)); - registry.RegisterTool(new MockMcpTool("get_books", ToolType.Custom, isEnabledFunc: _ => true)); - RuntimeConfig config = CreateRuntimeConfig(); + IMcpTool[] generationA = + { + new MockMcpTool("a_one", ToolType.BuiltIn), + new MockMcpTool("a_two", ToolType.Custom) + }; + IMcpTool[] generationB = + { + new MockMcpTool("b_one", ToolType.BuiltIn), + new MockMcpTool("b_two", ToolType.Custom) + }; + registry.ReplaceAll(generationA, config); - // Act - List enabledTools = registry.GetEnabledTools(config).ToList(); - - // Assert - create_record and delete_record should be filtered out - Assert.AreEqual(3, enabledTools.Count); - Assert.IsTrue(enabledTools.Any(t => t.Name == "describe_entities")); - Assert.IsTrue(enabledTools.Any(t => t.Name == "read_records")); - Assert.IsTrue(enabledTools.Any(t => t.Name == "get_books")); - Assert.IsFalse(enabledTools.Any(t => t.Name == "create_record")); - Assert.IsFalse(enabledTools.Any(t => t.Name == "delete_record")); + ConcurrentQueue invalidSnapshots = new(); + Task writer = Task.Run(() => + { + for (int i = 0; i < 500; i++) + { + registry.ReplaceAll(i % 2 == 0 ? generationB : generationA, config); + } + }); + + Task[] readers = Enumerable.Range(0, 4) + .Select(_ => Task.Run(() => + { + for (int i = 0; i < 2_000; i++) + { + string[] names = registry.GetAdvertisedTools() + .Select(tool => tool.Name) + .ToArray(); + bool isGenerationA = names.SequenceEqual(new[] { "a_one", "a_two" }); + bool isGenerationB = names.SequenceEqual(new[] { "b_one", "b_two" }); + if (!isGenerationA && !isGenerationB) + { + invalidSnapshots.Enqueue(string.Join(",", names)); + } + } + })) + .ToArray(); + + Task.WaitAll(readers.Append(writer).ToArray()); + + Assert.AreEqual( + 0, + invalidSnapshots.Count, + $"Observed partial snapshots: {string.Join(" | ", invalidSnapshots.Take(5))}"); } /// @@ -454,6 +657,107 @@ public void BuiltInTools_IsEnabled_DefaultsToTrueWhenMcpNotConfigured() #region Private helpers + public static IEnumerable SchemaInstanceDataCases + { + get + { + foreach (string keyword in new[] { "default", "const", "examples", "enum", "x-extension" }) + { + foreach (string propertyName in new[] { "enum", "type", "required" }) + { + foreach (bool useOutputSchema in new[] { false, true }) + { + foreach (bool nested in new[] { false, true }) + { + yield return new object[] { keyword, propertyName, useOutputSchema, nested }; + } + } + } + } + } + } + + public static IEnumerable NestedSchemaCases + { + get + { + List templates = new() { "$SCHEMA" }; + foreach (string keyword in new[] { "properties", "patternProperties", "$defs", "definitions", "dependentSchemas", "dependencies" }) + { + templates.Add("{\"" + keyword + "\":{\"default\":$SCHEMA,\"const\":$SCHEMA,\"examples\":$SCHEMA}}"); + } + + foreach (string keyword in new[] { "items", "additionalItems", "additionalProperties", "unevaluatedItems", "unevaluatedProperties", "contains", "propertyNames", "not", "if", "then", "else", "contentSchema" }) + { + templates.Add("{\"" + keyword + "\":$SCHEMA}"); + } + + foreach (string keyword in new[] { "items", "prefixItems", "allOf", "anyOf", "oneOf" }) + { + templates.Add("{\"" + keyword + "\":[$SCHEMA]}"); + } + + foreach (string template in templates) + { + string toolSchemaTemplate = template == "$SCHEMA" + ? template + : "{\"type\":\"object\",\"properties\":{\"value\":" + template + "}}"; + foreach (bool useOutputSchema in new[] { false, true }) + { + yield return new object[] { toolSchemaTemplate, useOutputSchema }; + } + } + } + } + + private static string CreateSchemaWithInstanceData(string keyword, string propertyName, bool nested, bool reverse) + { + JsonObject instanceData = new() + { + [propertyName] = reverse ? new JsonArray("b", "a") : new JsonArray("a", "b") + }; + if (nested) + { + instanceData = new JsonObject + { + ["inputSchema"] = new JsonObject + { + ["outputSchema"] = new JsonObject + { + ["properties"] = new JsonObject { ["value"] = instanceData } + } + } + }; + } + + return new JsonObject + { + ["type"] = "object", + ["properties"] = new JsonObject + { + ["value"] = new JsonObject + { + [keyword] = keyword is "examples" or "enum" ? new JsonArray(instanceData) : instanceData + } + } + }.ToJsonString(); + } + + private static RetainedMetadataMcpTool CreateSchemaTool(string schemaJson, bool useOutputSchema) + { + Tool metadata = new() + { + Name = "same_tool", + InputSchema = JsonSerializer.Deserialize(useOutputSchema ? "{\"type\":\"object\"}" : schemaJson) + }; + if (useOutputSchema) + { + metadata.OutputSchema = JsonSerializer.Deserialize(schemaJson); + } + + return new RetainedMetadataMcpTool(metadata); + } + /// /// Mock implementation of IMcpTool for testing purposes. /// @@ -461,12 +765,21 @@ private class MockMcpTool : IMcpTool { private readonly string _toolName; private readonly Func? _isEnabledFunc; - - public MockMcpTool(string toolName, ToolType toolType, Func? isEnabledFunc = null) + private readonly string _description; + private readonly string _inputSchemaJson; + + public MockMcpTool( + string toolName, + ToolType toolType, + Func? isEnabledFunc = null, + string? description = null, + string? inputSchemaJson = null) { _toolName = toolName; ToolType = toolType; _isEnabledFunc = isEnabledFunc; + _description = description ?? $"Mock {toolType} tool"; + _inputSchemaJson = inputSchemaJson ?? "{\"type\":\"object\"}"; } public ToolType ToolType { get; } @@ -478,12 +791,11 @@ public bool IsEnabled(RuntimeConfig config) public Tool GetToolMetadata() { - // Create a simple JSON object for the input schema - using JsonDocument doc = JsonDocument.Parse("{\"type\": \"object\"}"); + using JsonDocument doc = JsonDocument.Parse(_inputSchemaJson); return new Tool { Name = _toolName, - Description = $"Mock {ToolType} tool", + Description = _description, InputSchema = doc.RootElement.Clone() }; } @@ -498,6 +810,30 @@ public Task ExecuteAsync( } } + private sealed class RetainedMetadataMcpTool : IMcpTool + { + private readonly Tool _metadata; + + public RetainedMetadataMcpTool(Tool metadata) + { + _metadata = metadata; + } + + public ToolType ToolType => ToolType.Custom; + + public bool IsEnabled(RuntimeConfig config) => true; + + public Tool GetToolMetadata() => _metadata; + + public Task ExecuteAsync( + JsonDocument? arguments, + IServiceProvider serviceProvider, + CancellationToken cancellationToken = default) + { + throw new NotImplementedException(); + } + } + /// /// Creates a RuntimeConfig with the specified DmlToolsConfig for testing. /// diff --git a/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs b/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs index b0ae580828..15a7660cec 100644 --- a/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs +++ b/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs @@ -1,15 +1,23 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +using System; +using System.Collections.Generic; using System.IO; using System.IO.Abstractions; +using System.Linq; using System.Text; using System.Threading; +using System.Threading.Tasks; using Azure.DataApiBuilder.Config; using Azure.DataApiBuilder.Config.ObjectModel; using Azure.DataApiBuilder.Core.Configurations; +using Azure.DataApiBuilder.Service.Utilities; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Hosting; using Microsoft.VisualStudio.TestTools.UnitTesting; using Moq; +using static Azure.DataApiBuilder.Config.DabConfigEvents; namespace Azure.DataApiBuilder.Service.Tests.UnitTests; @@ -159,6 +167,617 @@ public void HotReloadConfigRestRuntimeOptions() } } + [TestMethod] + public async Task ConcurrentHotReloadNotifications_SerializeCompletePipelines() + { + string testDirectory = Path.Combine( + Path.GetTempPath(), + $"dab-config-reload-serialization-{Guid.NewGuid():N}"); + Directory.CreateDirectory(testDirectory); + string configPath = Path.Combine(testDirectory, "dab-config.json"); + FileSystem fileSystem = new(); + fileSystem.File.WriteAllText( + configPath, + GenerateRuntimeSectionStringFromParams( + restPath: "/initial", + gqlPath: "/graphql", + restEnabled: true, + gqlEnabled: true, + gqlIntrospection: true, + mode: HostMode.Development)); + + HotReloadEventHandler hotReloadEventHandler = new(); + using FileSystemRuntimeConfigLoader configLoader = new( + fileSystem, + hotReloadEventHandler, + configPath, + connectionString: string.Empty, + isCliLoader: true); + Assert.IsTrue(configLoader.TryLoadKnownConfig(out _)); + + string[] orderedEvents = + { + QUERY_MANAGER_FACTORY_ON_CONFIG_CHANGED, + METADATA_PROVIDER_FACTORY_ON_CONFIG_CHANGED, + QUERY_ENGINE_FACTORY_ON_CONFIG_CHANGED, + MUTATION_ENGINE_FACTORY_ON_CONFIG_CHANGED, + DOCUMENTOR_ON_CONFIG_CHANGED, + AUTHZ_RESOLVER_ON_CONFIG_CHANGED, + MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED, + GRAPHQL_SCHEMA_EVICTION_ON_CONFIG_CHANGED, + GRAPHQL_SCHEMA_CREATOR_ON_CONFIG_CHANGED, + GRAPHQL_SCHEMA_REFRESH_ON_CONFIG_CHANGED, + LOG_LEVEL_INITIALIZER_ON_CONFIG_CHANGE + }; + List observedEvents = new(); + object observedEventsLock = new(); + using ManualResetEventSlim generationAEnteredPipeline = new(); + using ManualResetEventSlim releaseGenerationA = new(); + using ManualResetEventSlim generationBReachedGate = new(); + + foreach (string eventName in orderedEvents) + { + hotReloadEventHandler.Subscribe(eventName, (_, args) => + { + string generation = configLoader.RuntimeConfig!.Runtime!.Rest!.Path; + lock (observedEventsLock) + { + observedEvents.Add($"{generation}:{args.EventName}"); + } + + if (string.Equals(generation, "/generation-a", StringComparison.Ordinal) && + string.Equals(args.EventName, QUERY_MANAGER_FACTORY_ON_CONFIG_CHANGED, StringComparison.Ordinal)) + { + generationAEnteredPipeline.Set(); + if (!releaseGenerationA.Wait(TimeSpan.FromSeconds(10))) + { + throw new TimeoutException("Timed out waiting to release generation A."); + } + } + }); + } + + Task reloadA = Task.CompletedTask; + Task reloadB = Task.CompletedTask; + try + { + fileSystem.File.WriteAllText( + configPath, + GenerateRuntimeSectionStringFromParams( + restPath: "/generation-a", + gqlPath: "/graphql", + restEnabled: true, + gqlEnabled: true, + gqlIntrospection: true, + mode: HostMode.Development)); + reloadA = Task.Run(() => configLoader.ProcessHotReloadNotification()); + Assert.IsTrue( + generationAEnteredPipeline.Wait(TimeSpan.FromSeconds(10)), + "Generation A did not reach its first ordered handler."); + + fileSystem.File.WriteAllText( + configPath, + GenerateRuntimeSectionStringFromParams( + restPath: "/generation-b", + gqlPath: "/graphql", + restEnabled: true, + gqlEnabled: true, + gqlIntrospection: true, + mode: HostMode.Development)); + reloadB = Task.Run(() => configLoader.ProcessHotReloadNotification( + beforeEnteringGate: generationBReachedGate.Set)); + Assert.IsTrue( + generationBReachedGate.Wait(TimeSpan.FromSeconds(10)), + "Generation B did not reach the loader serialization gate."); + + Assert.AreEqual( + "/generation-a", + configLoader.RuntimeConfig!.Runtime!.Rest!.Path, + "Generation B must not replace RuntimeConfig while generation A handlers are running."); + + releaseGenerationA.Set(); + await Task.WhenAll(reloadA, reloadB).WaitAsync(TimeSpan.FromSeconds(10)); + + string[] expectedEvents = orderedEvents + .Select(eventName => $"/generation-a:{eventName}") + .Concat(orderedEvents.Select(eventName => $"/generation-b:{eventName}")) + .ToArray(); + string[] actualEvents; + lock (observedEventsLock) + { + actualEvents = observedEvents.ToArray(); + } + + CollectionAssert.AreEqual( + expectedEvents, + actualEvents, + "Every generation A handler must finish before generation B starts its pipeline."); + Assert.AreEqual("/generation-b", configLoader.RuntimeConfig!.Runtime!.Rest!.Path); + } + finally + { + releaseGenerationA.Set(); + await Task.WhenAll(reloadA, reloadB).WaitAsync(TimeSpan.FromSeconds(10)); + Directory.Delete(testDirectory, recursive: true); + } + } + + [TestMethod] + public async Task StopAsync_CancelsAndDrainsActiveReloadBeforeReturning() + { + string testDirectory = Path.Combine( + Path.GetTempPath(), + $"dab-config-dispose-during-reload-{Guid.NewGuid():N}"); + Directory.CreateDirectory(testDirectory); + string configPath = Path.Combine(testDirectory, "dab-config.json"); + FileSystem fileSystem = new(); + fileSystem.File.WriteAllText( + configPath, + GenerateRuntimeSectionStringFromParams( + restPath: "/initial", + gqlPath: "/graphql", + restEnabled: true, + gqlEnabled: true, + gqlIntrospection: true, + mode: HostMode.Development)); + + HotReloadEventHandler hotReloadEventHandler = new(); + BlockingDisposeConfigFileWatcher configFileWatcher = new(); + FileSystemRuntimeConfigLoader configLoader = new( + fileSystem, + hotReloadEventHandler, + configPath, + connectionString: string.Empty, + isCliLoader: false, + logger: null, + configFileWatcherFactory: (_, _, _) => configFileWatcher); + Assert.IsTrue(configLoader.TryLoadKnownConfig(out _)); + + using ManualResetEventSlim reloadHandlerEntered = new(); + using ManualResetEventSlim reloadCancellationObserved = new(); + using ManualResetEventSlim releaseReloadHandler = new(); + using ManualResetEventSlim queuedReloadReachedGate = new(); + int laterHandlerInvocationCount = 0; + hotReloadEventHandler.Subscribe( + METADATA_PROVIDER_FACTORY_ON_CONFIG_CHANGED, + (_, args) => + { + reloadHandlerEntered.Set(); + if (!args.CancellationToken.WaitHandle.WaitOne(TimeSpan.FromSeconds(10))) + { + throw new TimeoutException("Timed out waiting for reload cancellation."); + } + + reloadCancellationObserved.Set(); + if (!releaseReloadHandler.Wait(TimeSpan.FromSeconds(10))) + { + throw new TimeoutException("Timed out waiting to release the simulated metadata refresh."); + } + }); + hotReloadEventHandler.Subscribe( + MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED, + (_, _) => Interlocked.Increment(ref laterHandlerInvocationCount)); + + Task activeReload = Task.CompletedTask; + Task queuedReload = Task.CompletedTask; + Task stopTask = Task.CompletedTask; + try + { + fileSystem.File.WriteAllText( + configPath, + GenerateRuntimeSectionStringFromParams( + restPath: "/blocked", + gqlPath: "/graphql", + restEnabled: true, + gqlEnabled: true, + gqlIntrospection: true, + mode: HostMode.Development)); + activeReload = Task.Run(() => configLoader.ProcessHotReloadNotification()); + Assert.IsTrue( + reloadHandlerEntered.Wait(TimeSpan.FromSeconds(5)), + "The hot-reload pipeline did not reach the blocking metadata handler."); + + queuedReload = Task.Run(() => configLoader.ProcessHotReloadNotification( + beforeEnteringGate: queuedReloadReachedGate.Set)); + Assert.IsTrue( + queuedReloadReachedGate.Wait(TimeSpan.FromSeconds(5)), + "The queued callback did not reach the serialization gate."); + + stopTask = configLoader.StopAsync(CancellationToken.None); + Assert.IsTrue( + reloadCancellationObserved.Wait(TimeSpan.FromSeconds(5)), + "Shutdown cancellation did not reach the active metadata handler."); + Assert.IsTrue( + configFileWatcher.StopWatchingCalled.Wait(TimeSpan.FromSeconds(1)), + "Shutdown must synchronously disable the watcher."); + Assert.IsTrue( + configFileWatcher.DisposeEntered.Wait(TimeSpan.FromSeconds(5)), + "Watcher resource disposal was not scheduled."); + Assert.IsFalse( + configFileWatcher.DisposeCompleted.IsSet, + "Loader disposal must not wait for potentially blocking watcher resource cleanup."); + Assert.AreSame( + queuedReload, + await Task.WhenAny(queuedReload, Task.Delay(TimeSpan.FromSeconds(1))), + "A callback waiting on the serialization gate must be canceled during shutdown."); + Assert.IsFalse( + stopTask.IsCompleted, + "Shutdown must drain the active reload before host-owned dependencies can be disposed."); + Assert.IsFalse( + configLoader.ShutdownResourcesDisposed, + "Synchronization resources cannot be disposed while a gate owner is still active."); + Assert.AreEqual( + 0, + Volatile.Read(ref laterHandlerInvocationCount), + "Cancellation must prevent later ordered handlers from running."); + + releaseReloadHandler.Set(); + await Task.WhenAll(activeReload, queuedReload, stopTask).WaitAsync(TimeSpan.FromSeconds(5)); + Assert.IsTrue( + configLoader.ShutdownResourcesDisposed, + "StopAsync must dispose loader-owned synchronization resources after the drain."); + Assert.AreEqual( + 0, + Volatile.Read(ref laterHandlerInvocationCount), + "No later ordered handler may run after shutdown cancellation."); + Assert.AreEqual( + "/blocked", + configLoader.RuntimeConfig!.Runtime!.Rest!.Path, + "A callback queued before shutdown must exit without loading another generation."); + } + finally + { + releaseReloadHandler.Set(); + configFileWatcher.ReleaseDispose.Set(); + await configLoader.StopAsync(CancellationToken.None).WaitAsync(TimeSpan.FromSeconds(5)); + await Task.WhenAll(activeReload, queuedReload, stopTask).WaitAsync(TimeSpan.FromSeconds(5)); + Assert.IsTrue( + configFileWatcher.DisposeCompleted.Wait(TimeSpan.FromSeconds(5)), + "The watcher disposal worker did not finish after it was released."); + Directory.Delete(testDirectory, recursive: true); + } + } + + [TestMethod] + public async Task Dispose_IsNonBlockingAndEventuallyDisposesOwnedResources() + { + FileSystemRuntimeConfigLoader configLoader = new( + new FileSystem(), + isCliLoader: true); + TaskCompletionSource operationEntered = new( + TaskCreationOptions.RunContinuationsAsynchronously); + TaskCompletionSource releaseOperation = new( + TaskCreationOptions.RunContinuationsAsynchronously); + Task activeOperation = configLoader.ExecuteWithHotReloadSerializationAsync( + async _ => + { + operationEntered.TrySetResult(); + await releaseOperation.Task.ConfigureAwait(false); + }); + + try + { + await operationEntered.Task.WaitAsync(TimeSpan.FromSeconds(5)); + Task disposeCall = Task.Run(configLoader.Dispose); + await disposeCall.WaitAsync(TimeSpan.FromSeconds(1)); + + Assert.IsFalse( + configLoader.ShutdownResourcesDisposed, + "Dispose must not tear down synchronization resources while admitted work remains."); + } + finally + { + releaseOperation.TrySetResult(); + await activeOperation.WaitAsync(TimeSpan.FromSeconds(5)); + await configLoader.StopAsync(CancellationToken.None).WaitAsync(TimeSpan.FromSeconds(5)); + } + + Assert.IsTrue( + configLoader.ShutdownResourcesDisposed, + "Dispose must eventually release loader-owned synchronization resources."); + } + + [TestMethod] + public async Task HostShutdown_DrainsReloadBeforeHostedServicesAndDependenciesStop() + { + string testDirectory = Path.Combine( + Path.GetTempPath(), + $"dab-config-host-shutdown-{Guid.NewGuid():N}"); + Directory.CreateDirectory(testDirectory); + string configPath = Path.Combine(testDirectory, "dab-config.json"); + FileSystem fileSystem = new(); + fileSystem.File.WriteAllText( + configPath, + GenerateRuntimeSectionStringFromParams( + restPath: "/initial", + gqlPath: "/graphql", + restEnabled: true, + gqlEnabled: true, + gqlIntrospection: true, + mode: HostMode.Development)); + + HotReloadEventHandler hotReloadEventHandler = new(); + FileSystemRuntimeConfigLoader configLoader = new( + fileSystem, + hotReloadEventHandler, + configPath, + connectionString: string.Empty, + isCliLoader: true); + Assert.IsTrue(configLoader.TryLoadKnownConfig(out _)); + + ReloadDependency dependency = new(); + using ManualResetEventSlim reloadHandlerEntered = new(); + using ManualResetEventSlim reloadCancellationObserved = new(); + hotReloadEventHandler.Subscribe( + METADATA_PROVIDER_FACTORY_ON_CONFIG_CHANGED, + (_, args) => dependency.RunUntilCanceled( + args.CancellationToken, + reloadHandlerEntered, + reloadCancellationObserved)); + + Task activeReload = Task.CompletedTask; + HostedServiceStopObserver stopObserver = new(() => activeReload.IsCompleted); + IHost host = new HostBuilder() + .ConfigureServices(services => + { + services.AddSingleton(configLoader); + services.AddSingleton(_ => dependency); + services.AddSingleton(stopObserver); + services.AddSingleton(); + services.AddSingleton(serviceProvider => + serviceProvider.GetRequiredService()); + }) + .Build(); + + try + { + Assert.AreSame(dependency, host.Services.GetRequiredService()); + await host.StartAsync(); + fileSystem.File.WriteAllText( + configPath, + GenerateRuntimeSectionStringFromParams( + restPath: "/active", + gqlPath: "/graphql", + restEnabled: true, + gqlEnabled: true, + gqlIntrospection: true, + mode: HostMode.Development)); + activeReload = Task.Run(() => configLoader.ProcessHotReloadNotification()); + Assert.IsTrue( + reloadHandlerEntered.Wait(TimeSpan.FromSeconds(5)), + "The active reload did not reach its dependency."); + + await host.StopAsync().WaitAsync(TimeSpan.FromSeconds(5)); + + Assert.IsTrue( + reloadCancellationObserved.IsSet, + "Hosted shutdown did not cancel the active reload."); + Assert.IsTrue(activeReload.IsCompleted, "Hosted shutdown returned before reload drain."); + Assert.IsTrue( + stopObserver.ReloadWasDrained, + "The loader drain must run before earlier hosted services stop."); + Assert.IsFalse( + dependency.IsDisposed, + "Host stopping must drain reload work before singleton disposal begins."); + } + finally + { + await configLoader.StopAsync(CancellationToken.None).WaitAsync(TimeSpan.FromSeconds(5)); + await activeReload.WaitAsync(TimeSpan.FromSeconds(5)); + host.Dispose(); + Assert.IsTrue(dependency.IsDisposed, "Root provider disposal did not dispose the dependency."); + Assert.IsFalse( + dependency.WasUsedAfterDisposal, + "An active reload accessed a dependency after root provider disposal."); + Directory.Delete(testDirectory, recursive: true); + } + } + + [TestMethod] + public async Task ShutdownService_HonorsHostCancellationForUncooperativeHandler() + { + string testDirectory = Path.Combine( + Path.GetTempPath(), + $"dab-config-host-timeout-{Guid.NewGuid():N}"); + Directory.CreateDirectory(testDirectory); + string configPath = Path.Combine(testDirectory, "dab-config.json"); + FileSystem fileSystem = new(); + fileSystem.File.WriteAllText( + configPath, + GenerateRuntimeSectionStringFromParams( + restPath: "/initial", + gqlPath: "/graphql", + restEnabled: true, + gqlEnabled: true, + gqlIntrospection: true, + mode: HostMode.Development)); + + HotReloadEventHandler hotReloadEventHandler = new(); + FileSystemRuntimeConfigLoader configLoader = new( + fileSystem, + hotReloadEventHandler, + configPath, + connectionString: string.Empty, + isCliLoader: true); + Assert.IsTrue(configLoader.TryLoadKnownConfig(out _)); + + using ManualResetEventSlim handlerEntered = new(); + using ManualResetEventSlim releaseHandler = new(); + hotReloadEventHandler.Subscribe( + METADATA_PROVIDER_FACTORY_ON_CONFIG_CHANGED, + (_, _) => + { + handlerEntered.Set(); + if (!releaseHandler.Wait(TimeSpan.FromSeconds(10))) + { + throw new TimeoutException("Timed out waiting to release the handler."); + } + }); + + Task activeReload = Task.CompletedTask; + try + { + fileSystem.File.WriteAllText( + configPath, + GenerateRuntimeSectionStringFromParams( + restPath: "/active", + gqlPath: "/graphql", + restEnabled: true, + gqlEnabled: true, + gqlIntrospection: true, + mode: HostMode.Development)); + activeReload = Task.Run(() => configLoader.ProcessHotReloadNotification()); + Assert.IsTrue( + handlerEntered.Wait(TimeSpan.FromSeconds(5)), + "The active reload did not reach the uncooperative handler."); + + RuntimeConfigLoaderShutdownService shutdownService = new(configLoader); + using CancellationTokenSource hostCancellation = new(); + Task stopTask = shutdownService.StopAsync(hostCancellation.Token); + Assert.IsFalse( + stopTask.IsCompleted, + "The drain should still be waiting before the host timeout expires."); + + hostCancellation.Cancel(); + try + { + await stopTask.WaitAsync(TimeSpan.FromSeconds(5)); + Assert.Fail("Hosted shutdown should observe the host cancellation token."); + } + catch (OperationCanceledException) + { + // Expected: the configured host shutdown bound expired. + } + + Assert.IsFalse( + activeReload.IsCompleted, + "Host timeout must not falsely report that an uncooperative handler was drained."); + } + finally + { + releaseHandler.Set(); + await activeReload.WaitAsync(TimeSpan.FromSeconds(5)); + await configLoader.StopAsync(CancellationToken.None).WaitAsync(TimeSpan.FromSeconds(5)); + configLoader.Dispose(); + Directory.Delete(testDirectory, recursive: true); + } + } + + private sealed class ReloadDependency : IDisposable + { + private int _disposed; + private int _usedAfterDisposal; + + public bool IsDisposed => Volatile.Read(ref _disposed) != 0; + + public bool WasUsedAfterDisposal => Volatile.Read(ref _usedAfterDisposal) != 0; + + public void RunUntilCanceled( + CancellationToken cancellationToken, + ManualResetEventSlim entered, + ManualResetEventSlim cancellationObserved) + { + if (IsDisposed) + { + Interlocked.Exchange(ref _usedAfterDisposal, 1); + } + + entered.Set(); + cancellationToken.WaitHandle.WaitOne(); + cancellationObserved.Set(); + + if (IsDisposed) + { + Interlocked.Exchange(ref _usedAfterDisposal, 1); + } + } + + public void Dispose() + { + Interlocked.Exchange(ref _disposed, 1); + } + } + + private sealed class HostedServiceStopObserver(Func isReloadDrained) : IHostedService + { + public bool ReloadWasDrained { get; private set; } + + public Task StartAsync(CancellationToken cancellationToken) + { + return Task.CompletedTask; + } + + public Task StopAsync(CancellationToken cancellationToken) + { + ReloadWasDrained = isReloadDrained(); + return Task.CompletedTask; + } + } + + private sealed class BlockingDisposeConfigFileWatcher : IConfigFileWatcher + { + public event EventHandler? NewFileContentsDetected + { + add { } + remove { } + } + + public ManualResetEventSlim StopWatchingCalled { get; } = new(); + + public ManualResetEventSlim DisposeEntered { get; } = new(); + + public ManualResetEventSlim ReleaseDispose { get; } = new(); + + public ManualResetEventSlim DisposeCompleted { get; } = new(); + + public void StopWatching() + { + StopWatchingCalled.Set(); + } + + public void Dispose() + { + DisposeEntered.Set(); + if (!ReleaseDispose.Wait(TimeSpan.FromSeconds(10))) + { + throw new TimeoutException("Timed out waiting to release watcher disposal."); + } + + DisposeCompleted.Set(); + } + } + + [TestMethod] + public void ConfigFileWatcher_StopWatching_DisablesAndDetachesUnderlyingWatcher() + { + IFileSystem fileSystem = Mock.Of(); + Mock.Get(fileSystem) + .Setup(fs => fs.File.Exists(It.IsAny())) + .Returns(true); + Mock.Get(fileSystem) + .Setup(fs => fs.File.ReadAllBytes(It.IsAny())) + .Returns(Encoding.UTF8.GetBytes("InitialValue")); + Mock fileSystemWatcher = new(); + fileSystemWatcher + .Setup(watcher => watcher.FileSystem) + .Returns(fileSystem); + IConfigFileWatcher configFileWatcher = new ConfigFileWatcher( + fileSystemWatcher.Object, + Directory.GetCurrentDirectory(), + "dab-config.json"); + + configFileWatcher.StopWatching(); + configFileWatcher.Dispose(); + + fileSystemWatcher.VerifySet( + watcher => watcher.EnableRaisingEvents = false, + Times.Once); + fileSystemWatcher.VerifyRemove( + watcher => watcher.Changed -= It.IsAny(), + Times.Once); + fileSystemWatcher.Verify(watcher => watcher.Dispose(), Times.Once); + } + #region ConfigFileWatcher NewFileContentsDetected event invocation tests private const string UNEXPECTED_INVOCATION_COUNT_ERR = "Unexpected number of invocations of the NewFileContentsDetected event."; diff --git a/src/Service.Tests/UnitTests/McpServerConfigurationTests.cs b/src/Service.Tests/UnitTests/McpServerConfigurationTests.cs new file mode 100644 index 0000000000..199279b001 --- /dev/null +++ b/src/Service.Tests/UnitTests/McpServerConfigurationTests.cs @@ -0,0 +1,199 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#nullable enable + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Net; +using System.Net.Http; +using System.Net.Http.Json; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Azure.DataApiBuilder.Config.ObjectModel; +using Azure.DataApiBuilder.Mcp.Core; +using Azure.DataApiBuilder.Mcp.Model; +using Microsoft.AspNetCore.Builder; +using Microsoft.AspNetCore.Hosting; +using Microsoft.AspNetCore.TestHost; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Options; +using Microsoft.VisualStudio.TestTools.UnitTesting; +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; +using static Azure.DataApiBuilder.Mcp.Model.McpEnums; + +namespace Azure.DataApiBuilder.Service.Tests.UnitTests +{ + [TestClass] + public class McpServerConfigurationTests + { + [TestMethod] + public void ConfigureMcpServer_HttpDoesNotAdvertiseToolListChanges() + { + ServiceCollection services = new(); + services.AddLogging(); + services.ConfigureMcpServer(instructions: null); + + using ServiceProvider serviceProvider = services.BuildServiceProvider(); + McpServerOptions options = serviceProvider + .GetRequiredService>() + .Value; + + Assert.IsNotNull(options.Capabilities); + Assert.IsNotNull(options.Capabilities.Tools); + Assert.IsFalse( + options.Capabilities.Tools.ListChanged, + "HTTP must not promise tool-list notifications until session broadcast is implemented."); + } + + [TestMethod] + public async Task ListToolsHandler_RegistrySnapshot_OmitsDisabledTool() + { + McpToolRegistry registry = new(); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll(new[] { new DisabledMcpTool() }, config); + +#pragma warning disable ASPDEPR004 // TestServer uses the legacy in-memory web-host builder. + IWebHostBuilder hostBuilder = new WebHostBuilder() +#pragma warning restore ASPDEPR004 + .ConfigureServices(services => + { + services.AddRouting(); + services.AddLogging(); + services.AddSingleton(registry); + services.ConfigureMcpServer(instructions: null); + }) + .Configure(app => + { + app.UseRouting(); + app.UseEndpoints(endpoints => endpoints.MapMcp("/mcp")); + }); + using TestServer server = new(hostBuilder); + using HttpClient client = server.CreateClient(); + + using HttpRequestMessage initializeRequest = CreateRequest( + sessionId: null, + new + { + jsonrpc = "2.0", + id = 1, + method = "initialize", + @params = new + { + protocolVersion = "2025-11-25", + capabilities = new { }, + clientInfo = new { name = "registration-test", version = "1.0" } + } + }); + using HttpResponseMessage initializeResponse = await client.SendAsync(initializeRequest); + Assert.AreEqual(HttpStatusCode.OK, initializeResponse.StatusCode); + string sessionId = initializeResponse.Headers + .GetValues("Mcp-Session-Id") + .Single(); + + using HttpRequestMessage initializedRequest = CreateRequest( + sessionId, + new + { + jsonrpc = "2.0", + method = "notifications/initialized", + @params = new { } + }); + using HttpResponseMessage initializedResponse = await client.SendAsync(initializedRequest); + Assert.AreEqual(HttpStatusCode.Accepted, initializedResponse.StatusCode); + + using HttpRequestMessage listRequest = CreateRequest( + sessionId, + new + { + jsonrpc = "2.0", + id = 2, + method = "tools/list", + @params = new { } + }); + using HttpResponseMessage listResponse = await client.SendAsync(listRequest); + string responseBody = await listResponse.Content.ReadAsStringAsync(); + Assert.AreEqual(HttpStatusCode.OK, listResponse.StatusCode, responseBody); + using JsonDocument payload = JsonDocument.Parse(GetJsonPayload(responseBody)); + + Assert.AreEqual( + 0, + payload.RootElement + .GetProperty("result") + .GetProperty("tools") + .GetArrayLength(), + "The HTTP handler must serve the configuration-aware advertised snapshot."); + Assert.IsTrue( + registry.TryGetTool("disabled_tool", out _), + "Disabled tools remain registered for structured execution-time errors."); + } + + private static HttpRequestMessage CreateRequest(string? sessionId, object payload) + { + HttpRequestMessage request = new(HttpMethod.Post, "/mcp") + { + Content = JsonContent.Create(payload) + }; + request.Headers.Add("Accept", "application/json, text/event-stream"); + if (sessionId is not null) + { + request.Headers.Add("Mcp-Session-Id", sessionId); + } + + return request; + } + + private static string GetJsonPayload(string responseBody) + { + return responseBody.TrimStart().StartsWith('{') + ? responseBody + : responseBody + .Split('\n') + .Select(line => line.TrimEnd('\r')) + .Where(line => line.StartsWith("data:", StringComparison.Ordinal)) + .Select(line => line["data:".Length..].TrimStart()) + .First(payload => payload.StartsWith('{')); + } + + private static RuntimeConfig CreateRuntimeConfig() + { + return new RuntimeConfig( + Schema: "test-schema", + DataSource: new DataSource(DatabaseType.MSSQL, string.Empty, Options: null), + Runtime: new( + Rest: new(), + GraphQL: new(), + Mcp: new(Enabled: true), + Host: new(Cors: null, Authentication: null, Mode: HostMode.Development)), + Entities: new(new Dictionary())); + } + + private sealed class DisabledMcpTool : IMcpTool + { + public ToolType ToolType => ToolType.Custom; + + public Tool GetToolMetadata() + { + return new Tool + { + Name = "disabled_tool", + Description = "Disabled test tool", + InputSchema = JsonSerializer.Deserialize("{\"type\":\"object\"}") + }; + } + + public bool IsEnabled(RuntimeConfig config) => false; + + public Task ExecuteAsync( + JsonDocument? arguments, + IServiceProvider serviceProvider, + CancellationToken cancellationToken = default) + { + throw new NotImplementedException(); + } + } + } +} diff --git a/src/Service.Tests/UnitTests/McpStdioHelperTests.cs b/src/Service.Tests/UnitTests/McpStdioHelperTests.cs index ecea06b236..0d237bd99c 100644 --- a/src/Service.Tests/UnitTests/McpStdioHelperTests.cs +++ b/src/Service.Tests/UnitTests/McpStdioHelperTests.cs @@ -6,10 +6,13 @@ using System; using System.Collections.Generic; using System.IO; +using System.IO.Abstractions.TestingHelpers; using System.Net; using System.Threading; using System.Threading.Tasks; +using Azure.DataApiBuilder.Config; using Azure.DataApiBuilder.Config.DatabasePrimitives; +using Azure.DataApiBuilder.Core.Configurations; using Azure.DataApiBuilder.Core.Services; using Azure.DataApiBuilder.Core.Services.MetadataProviders; using Azure.DataApiBuilder.Mcp.Core; @@ -17,6 +20,7 @@ using Azure.DataApiBuilder.Service.Utilities; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Hosting; +using Microsoft.Extensions.Logging.Abstractions; using Microsoft.VisualStudio.TestTools.UnitTesting; namespace Azure.DataApiBuilder.Service.Tests.UnitTests @@ -31,6 +35,8 @@ public void RunMcpStdioHost_DoesNotStartWebHost() TestMetadataProviderFactory metadataProviderFactory = new(); using ServiceProvider serviceProvider = BuildServices(stdioServer, metadataProviderFactory, out TestApplicationLifetime lifetime); + TestMcpToolRegistryRefreshService refreshService = + (TestMcpToolRegistryRefreshService)serviceProvider.GetRequiredService(); TestHost host = new(serviceProvider); bool result = McpStdioHelper.RunMcpStdioHost(host); @@ -42,29 +48,45 @@ 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."); + CollectionAssert.AreEqual( + new[] { "metadata", "registry" }, + 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."); 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 the metadata providers itself: it never calls " + - "host.Run(), so Startup.Configure -- the only caller of PerformOnConfigChangeAsync " + - "-- never runs, and without this every tool call fails with " + - "\"Database object for entity '' has not been inferred.\""); + "MCP stdio mode must initialize metadata exactly once through the shared runtime initialization path."); + Assert.IsTrue(serviceProvider.GetRequiredService().ShutdownResourcesDisposed, + "MCP stdio shutdown must drain the loader before disposing the host."); } /// - /// A startup failure must not escape RunMcpStdioHost, whose contract is a bool, and must stop - /// the server rather than let it serve entities that have no database object -- the failure - /// this initialization exists to prevent. stdout carries JSON-RPC, so it reports on stderr. + /// Startup and loop failures are reported through the bool contract and stderr without + /// corrupting stdout. A startup failure must not let the server serve uninitialized tools. /// - [TestMethod] - public void RunMcpStdioHost_StartupFails_ReportsOnStandardErrorAndDoesNotServeTools() + [DataTestMethod] + [DataRow(false)] + [DataRow(true)] + public void RunMcpStdioHost_Fails_ReportsOnStandardErrorAndDisposesHost(bool failDuringStdio) { - TestMcpStdioServer stdioServer = new(); + Exception failure = failDuringStdio + ? new InvalidOperationException("The stdio loop failed.") + : InferenceFailure(); + TestMcpStdioServer stdioServer = new() + { + RunAsyncException = failDuringStdio ? failure : null + }; TestMetadataProviderFactory metadataProviderFactory = new() { - InitializeAsyncException = InferenceFailure() + InitializeAsyncException = failDuringStdio ? null : failure }; using ServiceProvider serviceProvider = BuildServices(stdioServer, metadataProviderFactory, out _); @@ -90,14 +112,20 @@ public void RunMcpStdioHost_StartupFails_ReportsOnStandardErrorAndDoesNotServeTo string reported = capturedError.ToString(); - Assert.IsFalse(result, "A startup failure should be reported through the bool contract."); - Assert.AreEqual(0, stdioServer.RunAsyncCallCount, - "The stdio loop must not run: it would advertise entities that have no database object."); + 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."); + TestMcpToolRegistryRefreshService refreshService = + (TestMcpToolRegistryRefreshService)serviceProvider.GetRequiredService(); + Assert.AreEqual(failDuringStdio ? 1 : 0, refreshService.EnsureInitializedCallCount, + "The registry must not publish tools after metadata initialization fails."); Assert.AreEqual(1, host.DisposeCallCount, - "The host must still be disposed when startup fails."); + "The host must still be disposed when initialization or the loop fails."); + Assert.IsTrue(serviceProvider.GetRequiredService().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, "has not been inferred", + StringAssert.Contains(reported, failure.Message, "GetAwaiter().GetResult() rethrows the original exception, so the cause must survive."); Assert.AreEqual(string.Empty, capturedOut.ToString(), "stdout is the JSON-RPC channel; a stray byte on it corrupts the protocol."); @@ -144,6 +172,39 @@ public void RunMcpStdioHost_StartupFails_WhenStandardErrorSuppressed_LeavesConso "The stdio loop must not run after startup failed."); } + [DataTestMethod] + [DataRow(false)] + [DataRow(true)] + public void RunMcpStdioHost_Canceled_PropagatesCancellationAndDrainsLoader(bool cancelDuringStdio) + { + using CancellationTokenSource cancellation = new(); + cancellation.Cancel(); + OperationCanceledException failure = new(cancellation.Token); + TestMcpStdioServer stdioServer = new() + { + RunAsyncException = cancelDuringStdio ? failure : null + }; + TestMetadataProviderFactory metadataProviderFactory = new() + { + InitializeAsyncException = cancelDuringStdio ? null : failure + }; + using ServiceProvider serviceProvider = + BuildServices(stdioServer, metadataProviderFactory, out _); + TestHost host = new(serviceProvider); + + OperationCanceledException actual = Assert.ThrowsException( + () => McpStdioHelper.RunMcpStdioHost(host)); + + Assert.AreSame(failure, actual, "Cancellation must propagate to Program.StartEngine, not become a startup failure."); + Assert.AreEqual(cancelDuringStdio ? 1 : 0, stdioServer.RunAsyncCallCount); + TestMcpToolRegistryRefreshService refreshService = + (TestMcpToolRegistryRefreshService)serviceProvider.GetRequiredService(); + Assert.AreEqual(cancelDuringStdio ? 1 : 0, refreshService.EnsureInitializedCallCount); + Assert.AreEqual(1, host.DisposeCallCount); + Assert.IsTrue(serviceProvider.GetRequiredService().ShutdownResourcesDisposed, + "Cancellation must still drain the loader before disposing the host."); + } + private static DataApiBuilderException InferenceFailure() => new( message: "Database object for entity 'Book' has not been inferred.", statusCode: HttpStatusCode.ServiceUnavailable, @@ -156,8 +217,24 @@ private static ServiceProvider BuildServices( { lifetime = new TestApplicationLifetime(); + MockFileSystem fileSystem = new(new Dictionary + { + [FileSystemRuntimeConfigLoader.DEFAULT_CONFIG_FILE_NAME] = + new MockFileData(TestHelper.INITIAL_CONFIG) + }); + FileSystemRuntimeConfigLoader configLoader = new(fileSystem, isCliLoader: true); + RuntimeConfigProvider runtimeConfigProvider = new(configLoader); + RuntimeConfigValidator runtimeConfigValidator = new( + runtimeConfigProvider, + fileSystem, + NullLogger.Instance); + TestMcpToolRegistryRefreshService refreshService = new(metadataProviderFactory.InitializationOrder); + ServiceCollection services = new(); - services.AddSingleton(); + services.AddSingleton(configLoader); + services.AddSingleton(runtimeConfigProvider); + services.AddSingleton(runtimeConfigValidator); + services.AddSingleton(refreshService); services.AddSingleton(lifetime); services.AddSingleton(stdioServer); services.AddSingleton(metadataProviderFactory); @@ -169,6 +246,10 @@ private sealed class TestMetadataProviderFactory : IMetadataProviderFactory { public int InitializeAsyncCallCount { get; private set; } + public List InitializationOrder { get; } = new(); + + public CancellationToken CancellationToken { get; private set; } + /// /// When set, InitializeAsync() returns a faulted task carrying it, standing in for a /// metadata inference failure such as an unreachable database or an entity that is @@ -177,9 +258,13 @@ private sealed class TestMetadataProviderFactory : IMetadataProviderFactory /// public Exception? InitializeAsyncException { get; init; } - public Task InitializeAsync() + public Task InitializeAsync() => InitializeAsync(CancellationToken.None); + + public Task InitializeAsync(CancellationToken cancellationToken) { InitializeAsyncCallCount++; + InitializationOrder.Add("metadata"); + CancellationToken = cancellationToken; return InitializeAsyncException is null ? Task.CompletedTask : Task.FromException(InitializeAsyncException); @@ -200,6 +285,29 @@ public List GetAllMetadataExceptions() => new(); } + private sealed class TestMcpToolRegistryRefreshService : IMcpToolRegistryRefreshService + { + private readonly List _initializationOrder; + + public TestMcpToolRegistryRefreshService(List initializationOrder) + { + _initializationOrder = initializationOrder; + } + + public int EnsureInitializedCallCount { get; private set; } + + public CancellationToken CancellationToken { get; private set; } + + public void EnsureInitialized() => EnsureInitialized(CancellationToken.None); + + public void EnsureInitialized(CancellationToken cancellationToken) + { + EnsureInitializedCallCount++; + CancellationToken = cancellationToken; + _initializationOrder.Add("registry"); + } + } + private sealed class TestHost : IHost { public TestHost(System.IServiceProvider services) @@ -255,11 +363,15 @@ private sealed class TestMcpStdioServer : IMcpStdioServer public CancellationToken CancellationToken { get; private set; } + public Exception? RunAsyncException { get; init; } + public Task RunAsync(CancellationToken cancellationToken) { RunAsyncCallCount++; CancellationToken = cancellationToken; - return Task.CompletedTask; + return RunAsyncException is null + ? Task.CompletedTask + : Task.FromException(RunAsyncException); } } } diff --git a/src/Service.Tests/UnitTests/McpStdioServerInitializeTests.cs b/src/Service.Tests/UnitTests/McpStdioServerInitializeTests.cs index 53cff20188..9c599d74eb 100644 --- a/src/Service.Tests/UnitTests/McpStdioServerInitializeTests.cs +++ b/src/Service.Tests/UnitTests/McpStdioServerInitializeTests.cs @@ -85,7 +85,32 @@ public void HandleInitialize_ClientRequests2025_11_25_WithoutDescription_EmitsNe Assert.AreEqual(1, CountOutputLines(stdoutCapture)); } - private static McpStdioServer CreateServer(string? description, out StringWriter stdoutCapture) + [TestMethod] + public void HandleInitialize_WithoutNotifier_AdvertisesListChangedFalse() + { + McpStdioServer server = CreateServer( + description: null, + out StringWriter stdoutCapture, + registerNotifier: false); + + JsonElement responseRoot = InvokeHandleInitialize( + server, + stdoutCapture, + """ + {"jsonrpc":"2.0","id":3,"method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"client","version":"1.0.0"}}} + """); + + AssertInitializeEnvelopeAndCapabilities( + responseRoot, + expectedId: 3, + expectedProtocolVersion: "2025-11-25", + expectedListChanged: false); + } + + private static McpStdioServer CreateServer( + string? description, + out StringWriter stdoutCapture, + bool registerNotifier = true) { stdoutCapture = new StringWriter(); McpStdoutWriter stdoutWriter = new(stdoutCapture); @@ -102,11 +127,17 @@ private static McpStdioServer CreateServer(string? description, out StringWriter RuntimeConfigProvider runtimeConfigProvider = new StubRuntimeConfigProvider(runtimeConfig); IConfiguration configuration = new ConfigurationBuilder().Build(); - ServiceProvider serviceProvider = new ServiceCollection() - .AddSingleton(configuration) - .AddSingleton(stdoutWriter) - .AddSingleton(runtimeConfigProvider) - .BuildServiceProvider(); + ServiceCollection services = new(); + services.AddSingleton(configuration); + services.AddSingleton(stdoutWriter); + services.AddSingleton(runtimeConfigProvider); + if (registerNotifier) + { + services.AddSingleton( + new McpStdioToolListChangedNotifier(stdoutWriter)); + } + + ServiceProvider serviceProvider = services.BuildServiceProvider(); return new McpStdioServer(new McpToolRegistry(), serviceProvider); } @@ -120,14 +151,20 @@ private static JsonElement InvokeHandleInitialize(McpStdioServer server, StringW JsonElement requestRoot = request.RootElement; JsonElement? id = requestRoot.TryGetProperty("id", out JsonElement idElement) ? idElement : null; - handleInitialize.Invoke(server, new object?[] { id, requestRoot }); + handleInitialize.Invoke( + server, + new object?[] { id, requestRoot }); string output = ExtractSingleOutputLine(stdoutCapture); using JsonDocument response = JsonDocument.Parse(output); return response.RootElement.Clone(); } - private static void AssertInitializeEnvelopeAndCapabilities(JsonElement responseRoot, object expectedId, string expectedProtocolVersion) + private static void AssertInitializeEnvelopeAndCapabilities( + JsonElement responseRoot, + object expectedId, + string expectedProtocolVersion, + bool expectedListChanged = true) { Assert.AreEqual("2.0", responseRoot.GetProperty("jsonrpc").GetString()); if (expectedId is int expectedNumericId) @@ -143,7 +180,9 @@ private static void AssertInitializeEnvelopeAndCapabilities(JsonElement response Assert.AreEqual(expectedProtocolVersion, result.GetProperty("protocolVersion").GetString()); JsonElement capabilities = result.GetProperty("capabilities"); - Assert.IsTrue(capabilities.GetProperty("tools").GetProperty("listChanged").GetBoolean()); + Assert.AreEqual( + expectedListChanged, + capabilities.GetProperty("tools").GetProperty("listChanged").GetBoolean()); Assert.AreEqual(JsonValueKind.Object, capabilities.GetProperty("logging").ValueKind); JsonElement serverInfo = result.GetProperty("serverInfo"); diff --git a/src/Service.Tests/UnitTests/McpStdioServerProtocolTests.cs b/src/Service.Tests/UnitTests/McpStdioServerProtocolTests.cs index e20b9defed..d4f8fc0754 100644 --- a/src/Service.Tests/UnitTests/McpStdioServerProtocolTests.cs +++ b/src/Service.Tests/UnitTests/McpStdioServerProtocolTests.cs @@ -119,8 +119,9 @@ public async Task RunAsync_InitializeConfigurationFailureReturnsInternalError() public async Task RunAsync_ListToolsReturnsOnlyEnabledToolMetadata() { McpToolRegistry registry = new(); - registry.RegisterTool(new RecordingTool("enabled", isEnabled: true)); - registry.RegisterTool(new RecordingTool("disabled", isEnabled: false)); + registry.ReplaceAll( + new[] { new RecordingTool("enabled", isEnabled: true), new RecordingTool("disabled", isEnabled: false) }, + CreateRuntimeConfig()); string input = Request(id: 4, method: "tools/list") + Environment.NewLine; (McpStdioServer server, StringWriter output, _) = CreateServer(input, registry: registry); @@ -134,16 +135,24 @@ public async Task RunAsync_ListToolsReturnsOnlyEnabledToolMetadata() } [TestMethod] - public async Task RunAsync_ListToolsConfigurationFailureReturnsInternalError() + public async Task RunAsync_ListToolsConfigurationFailureReturnsPublishedSnapshot() { + // Discovery serves the last published snapshot without re-reading runtime configuration. + McpToolRegistry registry = new(); + registry.ReplaceAll(new[] { new RecordingTool("retained_tool") }, CreateRuntimeConfig()); string input = Request(id: 3, method: "tools/list") + Environment.NewLine; (McpStdioServer server, StringWriter output, _) = CreateServer( input, + registry: registry, runtimeConfigProvider: new ThrowingRuntimeConfigProvider()); await server.RunAsync(CancellationToken.None); - AssertError(ParseResponses(output).Single(), 3L, McpStdioJsonRpcErrorCodes.INTERNAL_ERROR, "Internal error"); + JsonElement response = ParseResponses(output).Single(); + Assert.AreEqual(3, response.GetProperty("id").GetInt32()); + JsonElement tools = response.GetProperty("result").GetProperty("tools"); + Assert.AreEqual(1, tools.GetArrayLength()); + Assert.AreEqual("retained_tool", tools[0].GetProperty("name").GetString()); } [DataTestMethod] @@ -168,7 +177,7 @@ public async Task RunAsync_CallToolSupportsNameAndArguments() { RecordingTool tool = new("test_tool"); McpToolRegistry registry = new(); - registry.RegisterTool(tool); + registry.ReplaceAll(new[] { tool }, CreateRuntimeConfig()); string input = Request( id: 12, method: "tools/call", @@ -188,7 +197,7 @@ public async Task RunAsync_CallToolSupportsLegacyToolNameAndMissingArguments() { RecordingTool tool = new("legacy_tool"); McpToolRegistry registry = new(); - registry.RegisterTool(tool); + registry.ReplaceAll(new[] { tool }, CreateRuntimeConfig()); string input = Request( id: 13, method: "tools/call", @@ -208,8 +217,7 @@ public async Task RunAsync_CallToolPrefersStandardNameOverLegacyToolName() RecordingTool standardTool = new("standard"); RecordingTool legacyTool = new("legacy"); McpToolRegistry registry = new(); - registry.RegisterTool(standardTool); - registry.RegisterTool(legacyTool); + registry.ReplaceAll(new[] { standardTool, legacyTool }, CreateRuntimeConfig()); string input = Request( id: 14, method: "tools/call", @@ -228,7 +236,7 @@ public async Task RunAsync_CallToolWithConfiguredRoleProvidesScopedIdentityAndCl { RecordingTool tool = new("role_tool"); McpToolRegistry registry = new(); - registry.RegisterTool(tool); + registry.ReplaceAll(new[] { tool }, CreateRuntimeConfig()); HttpContextAccessor accessor = new(); IConfiguration configuration = new ConfigurationBuilder() .AddInMemoryCollection(new Dictionary { ["MCP:Role"] = "writer" }) @@ -257,7 +265,7 @@ public async Task RunAsync_ThrowingToolReturnsInternalErrorAndClearsRoleContext( { RecordingTool tool = new("throwing_tool", exception: new InvalidOperationException("failure")); McpToolRegistry registry = new(); - registry.RegisterTool(tool); + registry.ReplaceAll(new[] { tool }, CreateRuntimeConfig()); HttpContextAccessor accessor = new(); IConfiguration configuration = new ConfigurationBuilder() .AddInMemoryCollection(new Dictionary { ["MCP:Role"] = "reader" }) diff --git a/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs b/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs index 224d534158..a6996aab07 100644 --- a/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs +++ b/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs @@ -12,6 +12,7 @@ using Azure.DataApiBuilder.Mcp.Model; using Microsoft.Extensions.DependencyInjection; using Microsoft.VisualStudio.TestTools.UnitTesting; +using Moq; namespace Azure.DataApiBuilder.Service.Tests.UnitTests { @@ -59,6 +60,39 @@ public async Task RunAsync_BlankLineThenShutdown_IgnoresBlankLineAndHandlesShutd "Expected shutdown response result.ok to be true."); } + [TestMethod] + public async Task RunAsync_CompleteInitializationHandshake_MarksToolListNotifierReady() + { + Mock notifier = new(); + string input = + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":\"2025-11-25\",\"capabilities\":{},\"clientInfo\":{\"name\":\"test\",\"version\":\"1.0\"}}}" + Environment.NewLine + + "{\"jsonrpc\":\"2.0\",\"method\":\"notifications/initialized\"}" + Environment.NewLine + + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"shutdown\"}" + Environment.NewLine; + (McpStdioServer server, _) = CreateServerWithCapturedOutput( + new StringReader(input), + notifier.Object); + + await server.RunAsync(CancellationToken.None); + + notifier.Verify(value => value.MarkInitialized(), Times.Once); + } + + [TestMethod] + public async Task RunAsync_InitializedNotificationBeforeInitialize_DoesNotMarkNotifierReady() + { + Mock notifier = new(); + string input = + "{\"jsonrpc\":\"2.0\",\"method\":\"notifications/initialized\"}" + Environment.NewLine + + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"shutdown\"}" + Environment.NewLine; + (McpStdioServer server, _) = CreateServerWithCapturedOutput( + new StringReader(input), + notifier.Object); + + await server.RunAsync(CancellationToken.None); + + notifier.Verify(value => value.MarkInitialized(), Times.Never); + } + [TestMethod] public async Task RunAsync_OutOfRangeNumericId_PreservesIdAndContinuesProcessing() { @@ -91,7 +125,9 @@ public async Task RunAsync_OutOfRangeNumericId_PreservesIdAndContinuesProcessing Assert.IsTrue(shutdownResponse.RootElement.GetProperty("result").GetProperty("ok").GetBoolean()); } - private static (McpStdioServer server, StringWriter stdoutCapture) CreateServerWithCapturedOutput(TextReader inputReader) + private static (McpStdioServer server, StringWriter stdoutCapture) CreateServerWithCapturedOutput( + TextReader inputReader, + IMcpStdioToolListChangedNotifier? notifier = null) { StringWriter stdoutCapture = new(); McpStdoutWriter stdoutWriter = new(stdoutCapture); @@ -99,6 +135,11 @@ private static (McpStdioServer server, StringWriter stdoutCapture) CreateServerW ServiceCollection services = new(); services.AddSingleton(stdoutWriter); services.AddSingleton(); + if (notifier is not null) + { + services.AddSingleton(notifier); + } + IServiceProvider serviceProvider = services.BuildServiceProvider(); McpStdioServer server = new( diff --git a/src/Service.Tests/UnitTests/McpStdioToolListChangedNotifierTests.cs b/src/Service.Tests/UnitTests/McpStdioToolListChangedNotifierTests.cs new file mode 100644 index 0000000000..fdd10c486f --- /dev/null +++ b/src/Service.Tests/UnitTests/McpStdioToolListChangedNotifierTests.cs @@ -0,0 +1,272 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#nullable enable + +using System; +using System.IO; +using System.Linq; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Azure.DataApiBuilder.Mcp.Core; +using Microsoft.Extensions.Logging; +using Microsoft.VisualStudio.TestTools.UnitTesting; +using Moq; + +namespace Azure.DataApiBuilder.Service.Tests.UnitTests +{ + [TestClass] + public class McpStdioToolListChangedNotifierTests + { + [TestMethod] + public void NotifyToolsListChanged_BeforeInitialized_DoesNotWrite() + { + StringWriter output = new(); + using McpStdoutWriter stdoutWriter = new(output); + McpStdioToolListChangedNotifier notifier = new(stdoutWriter); + + notifier.NotifyToolsListChanged(); + + Assert.AreEqual(string.Empty, output.ToString()); + } + + [TestMethod] + public void NotifyToolsListChanged_AfterInitialized_WritesProtocolFrame() + { + SignalingStringWriter output = new(); + using McpStdoutWriter stdoutWriter = new(output); + McpStdioToolListChangedNotifier notifier = new(stdoutWriter); + notifier.MarkInitialized(); + + notifier.NotifyToolsListChanged(); + Assert.IsTrue( + output.LineWritten.Wait(TimeSpan.FromSeconds(5)), + "The queued tool-list notification was not written."); + + string[] lines = output.ToString().Split( + Environment.NewLine, + StringSplitOptions.RemoveEmptyEntries); + Assert.AreEqual(1, lines.Length); + + using JsonDocument document = JsonDocument.Parse(lines[0]); + JsonElement root = document.RootElement; + Assert.AreEqual("2.0", root.GetProperty("jsonrpc").GetString()); + Assert.AreEqual( + "notifications/tools/list_changed", + root.GetProperty("method").GetString()); + Assert.AreEqual(JsonValueKind.Object, root.GetProperty("params").ValueKind); + Assert.AreEqual(0, root.GetProperty("params").EnumerateObject().Count()); + Assert.IsFalse(root.TryGetProperty("id", out _), + "JSON-RPC notifications must not include a request id."); + } + + [TestMethod] + public void MarkInitialized_IsIdempotent() + { + SignalingStringWriter output = new(); + using McpStdoutWriter stdoutWriter = new(output); + McpStdioToolListChangedNotifier notifier = new(stdoutWriter); + + notifier.MarkInitialized(); + notifier.MarkInitialized(); + notifier.NotifyToolsListChanged(); + Assert.IsTrue( + output.LineWritten.Wait(TimeSpan.FromSeconds(5)), + "The queued tool-list notification was not written."); + + string[] lines = output.ToString().Split( + Environment.NewLine, + StringSplitOptions.RemoveEmptyEntries); + Assert.AreEqual(1, lines.Length); + } + + [TestMethod] + public async Task NotifyToolsListChanged_WhenStdoutBlocks_ReturnsWithoutWaitingForWrite() + { + BlockingStringWriter output = new(); + using McpStdoutWriter stdoutWriter = new(output); + McpStdioToolListChangedNotifier notifier = new(stdoutWriter); + notifier.MarkInitialized(); + + Task notificationCall = Task.Run(notifier.NotifyToolsListChanged); + try + { + Assert.IsTrue( + output.WriteEntered.Wait(TimeSpan.FromSeconds(5)), + "The notification worker did not begin the stdout write."); + Assert.IsTrue( + await Task.WhenAny(notificationCall, Task.Delay(TimeSpan.FromSeconds(1))) == notificationCall, + "NotifyToolsListChanged must enqueue transport I/O instead of blocking the reload pipeline."); + } + finally + { + output.ReleaseWrite.Set(); + await notificationCall.WaitAsync(TimeSpan.FromSeconds(5)); + } + + Assert.IsTrue( + output.LineWritten.Wait(TimeSpan.FromSeconds(5)), + "The notification was not written after stdout resumed."); + } + + [TestMethod] + public void NotifyToolsListChanged_WhenPrimarySchedulingFails_UsesFallbackWorker() + { + SignalingStringWriter output = new(); + using McpStdoutWriter stdoutWriter = new(output); + int schedulingAttempts = 0; + McpStdioToolListChangedNotifier notifier = new( + stdoutWriter, + logger: null, + tryScheduleWorker: _ => + { + Interlocked.Increment(ref schedulingAttempts); + return false; + }); + notifier.MarkInitialized(); + + notifier.NotifyToolsListChanged(); + + Assert.IsTrue( + output.LineWritten.Wait(TimeSpan.FromSeconds(5)), + "The dedicated fallback worker did not deliver the pending notification."); + Assert.AreEqual(1, Volatile.Read(ref schedulingAttempts)); + Assert.AreEqual(1, output.LineCount); + } + + [TestMethod] + public void NotifyToolsListChanged_WhileWriteIsBlocked_CoalescesPendingChanges() + { + BlockingStringWriter output = new(); + using McpStdoutWriter stdoutWriter = new(output); + McpStdioToolListChangedNotifier notifier = new(stdoutWriter); + notifier.MarkInitialized(); + + notifier.NotifyToolsListChanged(); + try + { + Assert.IsTrue( + output.WriteEntered.Wait(TimeSpan.FromSeconds(5)), + "The first notification did not reach the blocking writer."); + + notifier.NotifyToolsListChanged(); + notifier.NotifyToolsListChanged(); + notifier.NotifyToolsListChanged(); + output.ReleaseWrite.Set(); + + Assert.IsTrue( + SpinWait.SpinUntil(() => output.LineCount == 2, TimeSpan.FromSeconds(5)), + "Expected one in-flight notification and one coalesced pending notification."); + Assert.AreEqual(2, output.LineCount); + } + finally + { + output.ReleaseWrite.Set(); + } + } + + [TestMethod] + public void NotifyToolsListChanged_WhenQueuedWriteFails_LogsError() + { + ThrowingStringWriter output = new(); + using McpStdoutWriter stdoutWriter = new(output); + Mock> logger = new(); + McpStdioToolListChangedNotifier notifier = new(stdoutWriter, logger.Object); + notifier.MarkInitialized(); + + notifier.NotifyToolsListChanged(); + + Assert.IsTrue( + output.WriteAttempted.Wait(TimeSpan.FromSeconds(5)), + "The notification worker did not attempt the stdout write."); + Assert.IsTrue( + SpinWait.SpinUntil(() => logger.Invocations.Count > 0, TimeSpan.FromSeconds(5)), + "The asynchronous notification failure was not logged."); + logger.Verify( + value => value.Log( + LogLevel.Error, + It.IsAny(), + It.Is((state, _) => + state.ToString()!.Contains( + "Failed to write an MCP tool-list change notification.", + StringComparison.Ordinal)), + It.IsAny(), + (Func)It.IsAny()), + Times.Once); + } + + private class SignalingStringWriter : StringWriter + { + private int _lineCount; + + public ManualResetEventSlim LineWritten { get; } = new(); + + public int LineCount => Volatile.Read(ref _lineCount); + + public override void WriteLine(string? value) + { + base.WriteLine(value); + Interlocked.Increment(ref _lineCount); + LineWritten.Set(); + } + + protected override void Dispose(bool disposing) + { + base.Dispose(disposing); + if (disposing) + { + LineWritten.Dispose(); + } + } + } + + private sealed class BlockingStringWriter : SignalingStringWriter + { + public ManualResetEventSlim WriteEntered { get; } = new(); + + public ManualResetEventSlim ReleaseWrite { get; } = new(); + + public override void WriteLine(string? value) + { + WriteEntered.Set(); + if (!ReleaseWrite.Wait(TimeSpan.FromSeconds(10))) + { + throw new TimeoutException("Timed out waiting to release the test stdout writer."); + } + + base.WriteLine(value); + } + + protected override void Dispose(bool disposing) + { + base.Dispose(disposing); + if (disposing) + { + WriteEntered.Dispose(); + ReleaseWrite.Dispose(); + } + } + } + + private sealed class ThrowingStringWriter : StringWriter + { + public ManualResetEventSlim WriteAttempted { get; } = new(); + + public override void WriteLine(string? value) + { + WriteAttempted.Set(); + throw new IOException("Expected stdout failure."); + } + + protected override void Dispose(bool disposing) + { + base.Dispose(disposing); + if (disposing) + { + WriteAttempted.Dispose(); + } + } + } + } +} diff --git a/src/Service.Tests/UnitTests/McpStdoutWriterTests.cs b/src/Service.Tests/UnitTests/McpStdoutWriterTests.cs index 420b6842a6..92e4746218 100644 --- a/src/Service.Tests/UnitTests/McpStdoutWriterTests.cs +++ b/src/Service.Tests/UnitTests/McpStdoutWriterTests.cs @@ -206,6 +206,29 @@ public void Dispose_DuringConcurrentWrites_DoesNotThrow() $"Producer task did not complete successfully. Status: {producer.Status}, Exception: {producer.Exception?.Message}"); } + [TestMethod] + public async Task Dispose_WhileWriteIsBlocked_ReturnsWithoutWaitingForStdout() + { + BlockingTextWriter inner = new(); + McpStdoutWriter writer = new(inner); + Task blockedWrite = Task.Run(() => writer.WriteLine("blocked")); + + Assert.IsTrue( + inner.WriteEntered.Wait(TimeSpan.FromSeconds(5)), + "The test writer did not enter its blocking write."); + + Task dispose = Task.Run(writer.Dispose); + Assert.AreSame( + dispose, + await Task.WhenAny(dispose, Task.Delay(TimeSpan.FromSeconds(1))), + "Disposal must not wait indefinitely for a blocked stdout write."); + + writer.WriteLine("after-dispose"); + inner.ReleaseWrite.Set(); + await blockedWrite.WaitAsync(TimeSpan.FromSeconds(5)); + Assert.AreEqual(1, inner.WriteCount); + } + /// /// The default constructor must NOT open the real stdout stream. /// This is critical: DI registers the writer eagerly during host build, @@ -225,5 +248,26 @@ public void Constructor_DoesNotOpenStdout() // Assert — no exception is the success criterion. } + + private sealed class BlockingTextWriter : StringWriter + { + internal ManualResetEventSlim WriteEntered { get; } = new(false); + + internal ManualResetEventSlim ReleaseWrite { get; } = new(false); + + internal int WriteCount { get; private set; } + + public override void WriteLine(string? value) + { + WriteEntered.Set(); + if (!ReleaseWrite.Wait(TimeSpan.FromSeconds(10))) + { + throw new TimeoutException("The blocking test writer was not released."); + } + + WriteCount++; + base.WriteLine(value); + } + } } } diff --git a/src/Service.Tests/UnitTests/SqlMetadataProviderHelperTests.cs b/src/Service.Tests/UnitTests/SqlMetadataProviderHelperTests.cs index bb74006d57..e8ae0b1825 100644 --- a/src/Service.Tests/UnitTests/SqlMetadataProviderHelperTests.cs +++ b/src/Service.Tests/UnitTests/SqlMetadataProviderHelperTests.cs @@ -8,6 +8,7 @@ using System.Reflection; using System.Runtime.CompilerServices; using System.Text.Json.Nodes; +using System.Threading; using Azure.DataApiBuilder.Config.DatabasePrimitives; using Azure.DataApiBuilder.Config.ObjectModel; using Azure.DataApiBuilder.Core.Configurations; @@ -241,7 +242,8 @@ public async System.Threading.Tasks.Task GenerateAutoentitiesIntoEntities_NullCo { MsSqlMetadataProvider provider = CreateProvider(); - System.Threading.Tasks.Task task = (System.Threading.Tasks.Task)GetMsSqlMethod("GenerateAutoentitiesIntoEntities") + System.Threading.Tasks.Task task = (System.Threading.Tasks.Task)GetMsSqlMethod( + "GenerateAutoentitiesIntoEntities", typeof(IReadOnlyDictionary)) .Invoke(provider, new object?[] { null })!; await task; @@ -257,7 +259,8 @@ public async System.Threading.Tasks.Task GenerateAutoentitiesIntoEntities_NullRe ["all"] = new Autoentity(null, null, null) }; - System.Threading.Tasks.Task task = (System.Threading.Tasks.Task)GetMsSqlMethod("GenerateAutoentitiesIntoEntities") + System.Threading.Tasks.Task task = (System.Threading.Tasks.Task)GetMsSqlMethod( + "GenerateAutoentitiesIntoEntities", typeof(IReadOnlyDictionary)) .Invoke(provider, new object?[] { autoentities })!; DataApiBuilderException exception = await Assert.ThrowsExceptionAsync(() => task); @@ -278,7 +281,8 @@ public async System.Threading.Tasks.Task GenerateAutoentitiesIntoEntities_Incomp ["all"] = new Autoentity(null, null, null) }; - System.Threading.Tasks.Task task = (System.Threading.Tasks.Task)GetMsSqlMethod("GenerateAutoentitiesIntoEntities") + System.Threading.Tasks.Task task = (System.Threading.Tasks.Task)GetMsSqlMethod( + "GenerateAutoentitiesIntoEntities", typeof(IReadOnlyDictionary)) .Invoke(provider, new object?[] { autoentities })!; await task; @@ -360,7 +364,8 @@ public void BaseVirtualMetadataOperations_UseDefaultBehavior() "Book", "Author", "dbo.book_authors", new Dictionary() }); - MethodInfo generateAutoentities = GetBaseMethod("GenerateAutoentitiesIntoEntities"); + MethodInfo generateAutoentities = GetBaseMethod( + "GenerateAutoentitiesIntoEntities", typeof(IReadOnlyDictionary)); TargetInvocationException exception = Assert.ThrowsException( () => generateAutoentities.Invoke(provider, new object?[] { null })); Assert.IsInstanceOfType(exception.InnerException); @@ -401,7 +406,7 @@ public async System.Threading.Tasks.Task FillSchemaForStoredProcedureAsync_Trans System.Threading.Tasks.Task task = (System.Threading.Tasks.Task)method.Invoke(provider, new object[] { - procedure, "Book", "dbo", "get_books", new StoredProcedureDefinition() + procedure, "Book", "dbo", "get_books", new StoredProcedureDefinition(), CancellationToken.None })!; DataApiBuilderException exception = @@ -489,11 +494,17 @@ private static Dictionary> GetMap(MsSqlMetada (Dictionary>)typeof(MsSqlMetadataProvider).BaseType! .GetField($"<{propertyName}>k__BackingField", BindingFlags.Instance | BindingFlags.NonPublic)!.GetValue(provider)!; - private static MethodInfo GetBaseMethod(string methodName) => - typeof(MsSqlMetadataProvider).BaseType!.GetMethod(methodName, BindingFlags.Instance | BindingFlags.NonPublic)!; + private static MethodInfo GetBaseMethod(string methodName, params Type[] parameterTypes) => + parameterTypes.Length == 0 + ? typeof(MsSqlMetadataProvider).BaseType!.GetMethod(methodName, BindingFlags.Instance | BindingFlags.NonPublic)! + : typeof(MsSqlMetadataProvider).BaseType!.GetMethod( + methodName, BindingFlags.Instance | BindingFlags.NonPublic, parameterTypes)!; - private static MethodInfo GetMsSqlMethod(string methodName) => - typeof(MsSqlMetadataProvider).GetMethod(methodName, BindingFlags.Instance | BindingFlags.NonPublic)!; + private static MethodInfo GetMsSqlMethod(string methodName, params Type[] parameterTypes) => + parameterTypes.Length == 0 + ? typeof(MsSqlMetadataProvider).GetMethod(methodName, BindingFlags.Instance | BindingFlags.NonPublic)! + : typeof(MsSqlMetadataProvider).GetMethod( + methodName, BindingFlags.Instance | BindingFlags.NonPublic, parameterTypes)!; private static void ConfigureAutoentityQuery(MsSqlMetadataProvider provider, JsonArray result) { @@ -503,6 +514,7 @@ private static void ConfigureAutoentityQuery(MsSqlMetadataProvider provider, Jso It.IsAny>(), It.IsAny?, System.Threading.Tasks.Task>>(), It.IsAny(), + It.IsAny(), It.IsAny(), It.IsAny>())) .ReturnsAsync(result); diff --git a/src/Service.Tests/UnitTests/SqlMetadataProviderUnitTests.cs b/src/Service.Tests/UnitTests/SqlMetadataProviderUnitTests.cs index dd6ad7d27e..8dae4104b7 100644 --- a/src/Service.Tests/UnitTests/SqlMetadataProviderUnitTests.cs +++ b/src/Service.Tests/UnitTests/SqlMetadataProviderUnitTests.cs @@ -8,6 +8,7 @@ using System.IO.Abstractions; using System.Net; using System.Text.Json.Nodes; +using System.Threading; using System.Threading.Tasks; using Azure.DataApiBuilder.Config.DatabasePrimitives; using Azure.DataApiBuilder.Config.ObjectModel; @@ -487,6 +488,7 @@ public async Task ValidateExceptionForInvalidResultFieldNames(string invalidFiel It.IsAny>(), It.IsAny, Task>>(), It.IsAny(), + It.IsAny(), It.IsAny(), It.IsAny>())) .ReturnsAsync(invalidFieldJsonArray); @@ -514,7 +516,9 @@ public async Task ValidateExceptionForInvalidResultFieldNames(string invalidFiel { Assert.AreEqual(HttpStatusCode.ServiceUnavailable, ex.StatusCode); Assert.AreEqual(DataApiBuilderException.SubStatusCodes.ErrorInInitialization, ex.SubStatusCode); - Assert.IsTrue(ex.Message.Contains("returns a column without a name")); + Assert.IsTrue( + ex.Message.Contains("returns a column without a name"), + $"Unexpected validation exception: {ex.Message}"); } TestHelper.UnsetAllDABEnvironmentVariables(); diff --git a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs index bdd7be4fbf..ea1c029759 100644 --- a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs +++ b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs @@ -1224,6 +1224,150 @@ await executor.SetManagedIdentityAccessTokenIfAnyAsync( #endregion + [DataTestMethod] + [DataRow(true, DisplayName = "RequestAborted cancels an explicitly cancellable query")] + [DataRow(false, DisplayName = "The explicit token cancels a query with RequestAborted")] + public async Task ExecuteQueryAsync_WithExplicitAndRequestTokens_ObservesEitherCancellation( + bool cancelRequest) + { + RuntimeConfig mockConfig = new( + Schema: string.Empty, + DataSource: new(DatabaseType.MSSQL, string.Empty, new()), + Runtime: new( + Rest: new(), + GraphQL: new(), + Mcp: new(), + Host: new(null, null)), + Entities: new(new Dictionary())); + MockFileSystem fileSystem = new(); + fileSystem.AddFile( + FileSystemRuntimeConfigLoader.DEFAULT_CONFIG_FILE_NAME, + new MockFileData(mockConfig.ToJson())); + FileSystemRuntimeConfigLoader loader = new(fileSystem); + RuntimeConfigProvider provider = new(loader) { IsLateConfigured = true }; + Mock> logger = new(); + DefaultHttpContext context = new(); + Mock httpContextAccessor = new(); + httpContextAccessor.Setup(accessor => accessor.HttpContext).Returns(context); + DbExceptionParser dbExceptionParser = new MsSqlDbExceptionParser(provider); + Mock queryExecutor = new( + provider, + dbExceptionParser, + logger.Object, + httpContextAccessor.Object, + null, + null) + { + CallBase = true + }; + queryExecutor + .Setup(executor => executor.CreateConnection(It.IsAny())) + .Returns(new SqlConnection()); + queryExecutor + .Setup(executor => executor.SetManagedIdentityAccessTokenIfAnyAsync( + It.IsAny(), + It.IsAny(), + It.IsAny())) + .Returns(Task.CompletedTask); + + CancellationToken observedToken = default; + TaskCompletionSource executionEntered = new( + TaskCreationOptions.RunContinuationsAsynchronously); + queryExecutor + .Setup(executor => executor.ExecuteQueryAgainstDbAsync( + It.IsAny(), + It.IsAny(), + It.IsAny>(), + It.IsAny, Task>>(), + It.IsAny(), + It.IsAny(), + It.IsAny>(), + It.IsAny())) + .Returns( + (SqlConnection connection, + string sql, + IDictionary parameters, + Func, Task> handler, + HttpContext requestContext, + string dataSourceName, + List arguments, + CancellationToken token) => + { + observedToken = token; + executionEntered.TrySetResult(); + return WaitUntilCanceledAsync(token); + }); + + using CancellationTokenSource explicitCancellation = new(); + using CancellationTokenSource requestCancellation = new(); + context.RequestAborted = requestCancellation.Token; + Task queryTask = queryExecutor.Object.ExecuteQueryAsync( + sqltext: string.Empty, + parameters: new Dictionary(), + dataReaderHandler: null, + dataSourceName: provider.GetConfig().DefaultDataSourceName, + cancellationToken: explicitCancellation.Token, + httpContext: context, + args: null); + + try + { + await executionEntered.Task.WaitAsync(TimeSpan.FromSeconds(5)); + if (cancelRequest) + { + requestCancellation.Cancel(); + } + else + { + explicitCancellation.Cancel(); + } + + try + { + await queryTask.WaitAsync(TimeSpan.FromSeconds(5)); + Assert.Fail("Expected linked query cancellation."); + } + catch (OperationCanceledException) + { + // Expected from either linked source. + } + + Assert.IsTrue(observedToken.IsCancellationRequested); + Assert.AreEqual(cancelRequest, requestCancellation.IsCancellationRequested); + Assert.AreEqual(!cancelRequest, explicitCancellation.IsCancellationRequested); + queryExecutor.Verify(executor => executor.ExecuteQueryAgainstDbAsync( + It.IsAny(), + It.IsAny(), + It.IsAny>(), + It.IsAny, Task>>(), + It.IsAny(), + It.IsAny(), + It.IsAny>()), + Times.Never); + } + finally + { + explicitCancellation.Cancel(); + requestCancellation.Cancel(); + try + { + await queryTask.WaitAsync(TimeSpan.FromSeconds(5)); + } + catch (OperationCanceledException) + { + // Expected during cleanup. + } + + await loader.StopAsync(CancellationToken.None).WaitAsync(TimeSpan.FromSeconds(5)); + } + + static async Task WaitUntilCanceledAsync(CancellationToken token) + { + await Task.Delay(Timeout.InfiniteTimeSpan, token); + return null; + } + } + /// /// Validates that when the CancellationToken from httpContext.RequestAborted times out /// during a long-running query execution (simulating ExecuteReaderAsync being interrupted diff --git a/src/Service/Program.cs b/src/Service/Program.cs index 76af52ba97..615cb3d95c 100644 --- a/src/Service/Program.cs +++ b/src/Service/Program.cs @@ -174,6 +174,11 @@ public static IHostBuilder CreateHostBuilder(string[] args, bool runMcpStdio, st { services.AddSingleton(_mcpStdoutWriter); services.AddSingleton(_mcpNotificationWriter); + services.AddSingleton(); + services.AddSingleton(serviceProvider => + serviceProvider.GetRequiredService()); + services.AddSingleton(serviceProvider => + serviceProvider.GetRequiredService()); } }) .ConfigureLogging(logging => diff --git a/src/Service/Startup.cs b/src/Service/Startup.cs index b41550bf2e..6bbd51fdbf 100644 --- a/src/Service/Startup.cs +++ b/src/Service/Startup.cs @@ -138,6 +138,7 @@ public void ConfigureServices(IServiceCollection services) services.AddSingleton(fileSystem); services.AddSingleton(sp => configLoader); services.AddSingleton(sp => configProvider); + services.AddSingleton(); bool runtimeConfigAvailable = configProvider.TryGetConfig(out RuntimeConfig? runtimeConfig); @@ -526,7 +527,11 @@ public void ConfigureServices(IServiceCollection services) // Subscribe the GraphQL schema refresh method to the specific hot-reload event _hotReloadEventHandler.Subscribe( DabConfigEvents.GRAPHQL_SCHEMA_REFRESH_ON_CONFIG_CHANGED, - (_, _) => RefreshGraphQLSchema(services)); + (_, args) => + { + args.CancellationToken.ThrowIfCancellationRequested(); + RefreshGraphQLSchema(services); + }); // Cache config IFusionCacheBuilder fusionCacheBuilder = services.AddFusionCache() @@ -605,6 +610,12 @@ public void ConfigureServices(IServiceCollection services) ConfigureResponseCompression(services, runtimeConfig); services.AddControllers(); + + // Hosted services stop in reverse registration order. Register the loader drain last + // so reload work exits before any other hosted service begins shutting down and before + // the root provider disposes reload subscribers or their dependencies. + services.AddSingleton(serviceProvider => + serviceProvider.GetRequiredService()); } /// @@ -1008,7 +1019,11 @@ public void Configure(IApplicationBuilder app, IWebHostEnvironment env, RuntimeC IRequestExecutorManager requestExecutorManager = app.ApplicationServices.GetRequiredService(); _hotReloadEventHandler.Subscribe( "GRAPHQL_SCHEMA_EVICTION_ON_CONFIG_CHANGED", - (_, _) => EvictGraphQLSchema(requestExecutorManager)); + (_, args) => + { + args.CancellationToken.ThrowIfCancellationRequested(); + EvictGraphQLSchema(requestExecutorManager); + }); app.UseEndpoints(endpoints => { @@ -1432,18 +1447,12 @@ private async Task PerformOnConfigChangeAsync(IApplicationBuilder app) { try { - RuntimeConfigProvider runtimeConfigProvider = app.ApplicationServices.GetService()!; - RuntimeConfig runtimeConfig = runtimeConfigProvider.GetConfig(); - + RuntimeConfig runtimeConfig = + await RuntimeInitializationHelper.InitializeRuntimeDependenciesAsync( + app.ApplicationServices); RuntimeConfigValidator runtimeConfigValidator = app.ApplicationServices.GetService()!; - // Now that the configuration has been set, perform validation of the runtime config - // itself. - - runtimeConfigValidator.ValidateConfigProperties(); - IMetadataProviderFactory sqlMetadataProviderFactory = app.ApplicationServices.GetRequiredService(); - await sqlMetadataProviderFactory.InitializeAsync(); // Manually trigger DI service instantiation of GraphQLSchemaCreator and RestService // to attempt to reduce chances that the first received client request diff --git a/src/Service/Telemetry/LogLevelInitializer.cs b/src/Service/Telemetry/LogLevelInitializer.cs index dc28242357..4d1f3bafe2 100644 --- a/src/Service/Telemetry/LogLevelInitializer.cs +++ b/src/Service/Telemetry/LogLevelInitializer.cs @@ -42,6 +42,7 @@ public void SetLogLevel() private void OnConfigChanged(object? sender, HotReloadEventArgs args) { + args.CancellationToken.ThrowIfCancellationRequested(); SetLogLevel(); } } diff --git a/src/Service/Utilities/McpStdioHelper.cs b/src/Service/Utilities/McpStdioHelper.cs index 65307336e3..c93e3a65bb 100644 --- a/src/Service/Utilities/McpStdioHelper.cs +++ b/src/Service/Utilities/McpStdioHelper.cs @@ -6,12 +6,13 @@ using System.Diagnostics.CodeAnalysis; using System.IO; using System.Text; -using Azure.DataApiBuilder.Core.Services.MetadataProviders; +using System.Threading; +using Azure.DataApiBuilder.Config; using Azure.DataApiBuilder.Mcp.Core; -using Azure.DataApiBuilder.Mcp.Model; using Microsoft.Extensions.Configuration; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Hosting; +using Microsoft.Extensions.Options; namespace Azure.DataApiBuilder.Service.Utilities { @@ -87,28 +88,23 @@ public static void ConfigureMcpStdio(IConfigurationBuilder builder, string? mcpR /// Runs the MCP stdio host. /// /// The host to run. - /// True when the stdio loop ran to completion; false when startup failed and was + /// True when the stdio loop ran to completion; false when startup or the loop failed and was /// reported, which Program.Main surfaces as a non-zero exit code. public static bool RunMcpStdioHost(IHost host) { try { - // Stdio mode never calls host.Run(), so Startup.Configure -- and with it - // PerformOnConfigChangeAsync, the only caller of IMetadataProviderFactory - // .InitializeAsync() -- never executes. Without this, entities are known to - // the tool registry (their names come from config) while no entity ever gets - // a database object, and every tool call fails with - // "Database object for entity '' has not been inferred." - IMetadataProviderFactory metadataProviderFactory = - host.Services.GetRequiredService(); - metadataProviderFactory.InitializeAsync().GetAwaiter().GetResult(); - - McpToolRegistry registry = - host.Services.GetRequiredService(); - IEnumerable tools = - host.Services.GetServices(); - - McpToolRegistry.InitializeAndRegisterTools(tools, registry, host.Services); + // This process entry point is deliberately synchronous and runs without an + // ASP.NET, UI, or other custom SynchronizationContext. Bridging the two async + // operations with GetAwaiter().GetResult() therefore cannot deadlock on a + // captured context and preserves direct exception propagation. + // Stdio deliberately does not start the web host, so Startup.Configure does not + // initialize runtime dependencies. Run the same serialized validation, metadata, + // and registry sequence used by HTTP startup before opening the stdio loop. + RuntimeInitializationHelper + .InitializeRuntimeDependenciesAsync(host.Services) + .GetAwaiter() + .GetResult(); IHostApplicationLifetime lifetime = host.Services.GetRequiredService(); @@ -148,6 +144,28 @@ public static bool RunMcpStdioHost(IHost host) } finally { + FileSystemRuntimeConfigLoader? configLoader = + host.Services.GetService(); + if (configLoader is not null) + { + TimeSpan shutdownTimeout = host.Services + .GetService>()? + .Value.ShutdownTimeout ?? new HostOptions().ShutdownTimeout; + using CancellationTokenSource shutdownCancellation = new(shutdownTimeout); + try + { + configLoader + .StopAsync(shutdownCancellation.Token) + .GetAwaiter() + .GetResult(); + } + catch (OperationCanceledException) + when (shutdownCancellation.IsCancellationRequested) + { + // Match Generic Host shutdown semantics: cancellation bounds the drain. + } + } + host.Dispose(); } } diff --git a/src/Service/Utilities/RuntimeConfigLoaderShutdownService.cs b/src/Service/Utilities/RuntimeConfigLoaderShutdownService.cs new file mode 100644 index 0000000000..1c533df119 --- /dev/null +++ b/src/Service/Utilities/RuntimeConfigLoaderShutdownService.cs @@ -0,0 +1,30 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Threading; +using System.Threading.Tasks; +using Azure.DataApiBuilder.Config; +using Microsoft.Extensions.Hosting; + +namespace Azure.DataApiBuilder.Service.Utilities +{ + /// + /// Stops and drains serialized runtime configuration work during the hosted-service shutdown + /// phase, before the root service provider disposes any hot-reload subscriber dependencies. + /// + internal sealed class RuntimeConfigLoaderShutdownService( + FileSystemRuntimeConfigLoader configLoader) : IHostedService + { + public Task StartAsync(CancellationToken cancellationToken) + { + return Task.CompletedTask; + } + + public Task StopAsync(CancellationToken cancellationToken) + { + // The loader first requests cancellation through its own token, then this host token + // bounds the drain according to HostOptions.ShutdownTimeout. + return configLoader.StopAsync(cancellationToken); + } + } +} diff --git a/src/Service/Utilities/RuntimeInitializationHelper.cs b/src/Service/Utilities/RuntimeInitializationHelper.cs new file mode 100644 index 0000000000..b46246cb23 --- /dev/null +++ b/src/Service/Utilities/RuntimeInitializationHelper.cs @@ -0,0 +1,62 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Threading.Tasks; +using Azure.DataApiBuilder.Config; +using Azure.DataApiBuilder.Config.ObjectModel; +using Azure.DataApiBuilder.Core.Configurations; +using Azure.DataApiBuilder.Core.Services.MetadataProviders; +using Azure.DataApiBuilder.Mcp.Core; +using Microsoft.Extensions.DependencyInjection; + +namespace Azure.DataApiBuilder.Service.Utilities +{ + /// + /// Coordinates initial configuration-dependent service construction for HTTP and stdio. + /// + internal static class RuntimeInitializationHelper + { + /// + /// Captures and validates the active configuration, initializes its database metadata, + /// and publishes the initial MCP registry while excluding file-triggered hot reloads. + /// + /// The application service provider. + /// The configuration generation initialized by this operation. + public static async Task InitializeRuntimeDependenciesAsync( + IServiceProvider serviceProvider) + { + ArgumentNullException.ThrowIfNull(serviceProvider); + + FileSystemRuntimeConfigLoader configLoader = + serviceProvider.GetRequiredService(); + RuntimeConfig? initializedConfig = null; + + await configLoader.ExecuteWithHotReloadSerializationAsync(async cancellationToken => + { + cancellationToken.ThrowIfCancellationRequested(); + RuntimeConfigProvider runtimeConfigProvider = + serviceProvider.GetRequiredService(); + initializedConfig = runtimeConfigProvider.GetConfig(); + + RuntimeConfigValidator runtimeConfigValidator = + serviceProvider.GetRequiredService(); + runtimeConfigValidator.ValidateConfigProperties(); + + IMetadataProviderFactory metadataProviderFactory = + serviceProvider.GetRequiredService(); + await metadataProviderFactory + .InitializeAsync(cancellationToken) + .ConfigureAwait(false); + + // MCP services are absent when MCP was disabled at startup. + cancellationToken.ThrowIfCancellationRequested(); + IMcpToolRegistryRefreshService? mcpToolRegistryRefreshService = + serviceProvider.GetService(); + mcpToolRegistryRefreshService?.EnsureInitialized(cancellationToken); + }).ConfigureAwait(false); + + return initializedConfig!; + } + } +}