diff --git a/docs/design/McpToolRegistryHotReload.md b/docs/design/McpToolRegistryHotReload.md new file mode 100644 index 0000000000..dd86c8b4cf --- /dev/null +++ b/docs/design/McpToolRegistryHotReload.md @@ -0,0 +1,1141 @@ +# Design Document: MCP Tool Registry Hot-Reload + +## Status + +Implemented by this change; retained as the design and review record. + +This document describes the implemented design for refreshing Data API builder's MCP tool registry +when runtime configuration is hot-reloaded. It records not only the final behavior but also the +compatibility boundaries, lifecycle guarantees, rejected alternatives, and known limitations that +are important during review. + +Source links point to the implementation delivered by this change. + +## Summary + +Before this change, DAB built its MCP tool registry once at startup. A hot-reloaded `RuntimeConfig` +could change which custom tools existed and could change their names, descriptions, and input +schemas, but the registry continued serving the startup tool instances and startup metadata. + +The implemented design keeps `McpToolRegistry` as a singleton and changes its contents to an atomically published immutable snapshot. A singleton refresh service builds a complete candidate snapshot after the existing metadata, engine, and authorization hot-reload handlers have run. If candidate construction succeeds, the service atomically swaps the snapshot. If it fails, the previous snapshot remains active. + +The design also sends `notifications/tools/list_changed` to an initialized stdio client when the advertised tool list or metadata changes. HTTP requests always read the current snapshot, but HTTP push notifications are deferred because the installed MCP SDK requires experimental session tracking APIs for broadcast notifications. + +## Reviewer Guide: Deliberate Design Calls + +The following decisions are intentional. They are summarized here because they are the places most +likely to look surprising when reviewing the implementation in isolation. + +| Design call | Why this design was chosen | Accepted consequence | +|---|---|---| +| Replace the complete registry instead of mutating it or DI | One atomic reference swap gives lookup and discovery one generation and preserves the previous generation on failure. | Every applicable reload rebuilds all generated custom tools. | +| Keep both a loader serialization gate and a registry writer lock | The loader gate protects cross-component metadata/config ordering; the registry lock protects publication from direct or out-of-band callers. They have different ownership and neither is redundant. | Registry rebuilds are serialized even outside the normal file-loader path. | +| Defer initial publication from hosted-service `StartAsync()` | Generic hosted services start before `Startup.Configure` completes database metadata initialization. Publishing there would advertise config-only or stale schemas. | The hosted service subscribes early, while the shared startup helper performs strict initial publication after metadata is ready. | +| Use `RuntimeConfig` reference identity as the generation token | The loader creates and publishes a new configuration object for each parsed generation. Reference identity is cheaper and less error-prone than structural comparison or maintaining a second version counter. | Callers that replace the current configuration object create a new generation even when values are equal. | +| Publish a fresh generation even when discovery metadata is equivalent | Generated tool instances must align with the current configuration and metadata generation. | Registry version and instances change, but no `list_changed` notification is sent for semantically equivalent discovery. | +| Retain the previous registry when a reload candidate fails | A complete previously validated snapshot is safer than a partial or invalid new snapshot. | Until a later successful reload, runtime configuration can be newer than MCP discovery; execution-time revalidation prevents obsolete authorization or execution. | +| Permit configuration-schema fallback when DB enrichment is unavailable | Existing custom-tool behavior already has a usable configuration-derived schema, and metadata unavailability should not silently remove an otherwise callable tool. | Discovery can be less precise; the fallback reason is logged. All other construction failures reject the whole candidate. | +| Keep ordered hot-reload callbacks synchronous | Existing DAB ordering is defined by synchronous event completion. Returning early or introducing an untracked queue would let later components publish before earlier ones finish. | A watcher callback can remain occupied for the reload duration; concurrent callbacks wait on the loader gate and shutdown cancels queued waiters. | +| Add cancellation through Core rather than only canceling gate waits | Shutdown cannot safely drain a reload if metadata connection opening, schema discovery, query execution, or token acquisition ignores cancellation. | The change necessarily touches shared query and metadata paths; default interface implementations preserve existing implementers. | +| Replace two tokenless `SqlMetadataProvider` protected virtual slots with token-bearing slots | `FillSchemaForStoredProcedureAsync()` and `GetColumnsAsync()` directly own database metadata I/O that must observe loader shutdown. They are implementation hooks for DAB's closed set of built-in providers, not part of a documented custom-provider plug-in contract. | A manually authored subclass that overrode either old signature must recompile and update the override. This is an explicit, limited source/binary compatibility exception. | +| Link an explicit query token with `HttpContext.RequestAborted` | Either the owning operation or the disconnected HTTP client must be able to cancel the same database work. | A linked token source is allocated only when both distinct tokens are cancellable; tokenless legacy calls retain their established virtual dispatch. | +| Bound shutdown by `HostOptions.ShutdownTimeout` | An unbounded drain can hang process shutdown forever when an extension callback does not cooperate. | A successful drain guarantees dependency safety; after timeout, the host may dispose dependencies while non-cooperative extension code is still running. .NET cannot forcibly terminate that code. | +| Make synchronous `Dispose()` nonblocking | `Dispose()` can run after the host's bounded drain has already timed out and must not reintroduce an infinite wait. | Coordinated hosts and direct consumers that need a drain must call `StopAsync()` before disposing dependent services. | +| Dispose the OS watcher and blocked stdout resources without joining them | `FileSystemWatcher.Dispose()` and an abandoned stdout pipe can block independently of reload correctness. | Event admission is stopped synchronously; exceptional resource cleanup is left to a background worker or process teardown rather than delaying host shutdown. | +| Coalesce stdio invalidations and do not replay pre-initialization changes | `list_changed` is an invalidation, not a change log; one frame tells the client to fetch the latest complete snapshot. Before initialization the client has no established cache to invalidate. | Intermediate generations are not individually reported. | +| Advertise stdio push but not HTTP push | Stdio has one owned connection. HTTP broadcast requires experimental MCP SDK session interception and tracking. | HTTP clients see the latest snapshot on their next explicit `tools/list` request but receive no push invalidation in this change. | +| Remove CLR-public MCP implementation plumbing but preserve supported Core interfaces | Incremental registry methods permit a second, non-atomic construction model and the MCP assembly is not a supported reference package. Core interfaces are supported extension points. | Manual consumers of unsupported MCP runtime members must migrate; Core implementers remain source- and binary-compatible through default interface methods. | +| Allow an in-flight call to retain its resolved old tool instance | Invalidating or canceling arbitrary requests at the instant of publication would require per-request generation tracking. | The call can finish, but generated tools revalidate current configuration, metadata, and authorization before execution. | + +The detailed sections below define the exact guarantees and alternatives behind these calls. + +## Motivation + +Custom MCP tools are generated from stored-procedure entities with `mcp.custom-tool` enabled. Before +this change, those tools were constructed from the startup configuration and registered as DI +singletons. The registry was then populated once by a hosted service. + +Consequently, a configuration hot-reload can leave MCP discovery stale in several ways: + +- A newly enabled custom tool does not appear. +- A removed or disabled custom tool remains registered. +- Renaming an entity does not update the tool name. +- Changing an entity description does not update the tool description. +- Changing stored-procedure parameters does not update the advertised input schema. +- A custom tool can retain metadata derived from an old database metadata generation. + +Built-in tool visibility already evaluated the current configuration during each `tools/list` +request, but combined it with a fixed startup registry. This avoided some stale built-in visibility, +but did not solve stale custom tools or provide a single consistent registry generation. + +## Previous Implementation + +### Registry construction + +Before this change, [McpServiceCollectionExtensions.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpServiceCollectionExtensions.cs): + +1. Registers `McpToolRegistry` as a singleton. +2. Registers `McpToolRegistryInitializer` as a hosted service. +3. Discovers built-in `IMcpTool` implementations and registers them as singletons. +4. Builds custom tools from the startup `RuntimeConfig` and registers each custom tool as a singleton. + +The former `McpToolRegistryInitializer` resolved every `IMcpTool` and registered it once when the host started. + +[McpStdioHelper.cs](../../src/Service/Utilities/McpStdioHelper.cs) separately initializes the registry because stdio mode deliberately builds, but does not start, the ASP.NET Core web host. + +### Registry state + +[McpToolRegistry.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs) stored tools in a mutable, case-insensitive `Dictionary`. It supported individual registration, lookup by name, and filtering enabled tools using a supplied `RuntimeConfig`. + +The dictionary was safe under startup-only mutation, but could not be modified concurrently with MCP requests. + +### Custom tool metadata + +[DynamicCustomTool.cs](../../src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs) captures an `Entity` at construction time. Its tool name, description, and configuration-based parameter schema therefore belong to that configuration generation. `InitializeMetadata(IServiceProvider)` may cache a schema enriched from database metadata. + +Execution is safer than discovery: `ExecuteAsync()` retrieves the current `RuntimeConfig`, verifies that the entity still exists, verifies that it is still a stored procedure with custom-tool enabled, and uses current database metadata and authorization state. A stale tool can therefore fail safely, but it can still be advertised with stale metadata. + +### MCP request handlers + +[McpServerConfiguration.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs) implements HTTP `tools/list` and `tools/call` handlers using the registry singleton. + +[McpStdioServer.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs) implements the equivalent stdio JSON-RPC handlers. + +Both transports combined a fixed registry with the latest runtime configuration during `tools/list`. + +### Existing hot-reload pipeline + +[RuntimeConfigLoader.cs](../../src/Config/RuntimeConfigLoader.cs) raises named hot-reload events in an intentional order: + +1. Query-manager factory. +2. Metadata-provider factory. +3. Query-engine factory. +4. Mutation-engine factory. +5. Documentation. +6. Authorization resolver. +7. GraphQL schema operations. +8. Log-level initialization. + +The MCP registry needs refreshed database metadata and must not publish newly callable tools before their query, mutation, and authorization dependencies are ready. It must therefore participate in this ordered pipeline rather than subscribing directly to the earlier runtime change token. + +## Goals + +1. Refresh custom tool membership after a successful runtime configuration hot-reload. +2. Refresh custom tool names, descriptions, and input schemas. +3. Preserve current built-in tool enablement behavior. +4. Ensure `tools/list` observes one internally consistent registry generation. +5. Ensure `tools/call` never observes a partially rebuilt registry. +6. Keep registry reads lock-free or effectively lock-free. +7. Preserve the previous registry when a hot-reload rebuild fails. +8. Use the same construction path for HTTP startup, stdio startup, and hot-reload. +9. Notify an initialized stdio client when advertised tools change. +10. Keep the implementation compatible with the existing ordered hot-reload architecture. +11. Avoid dynamic mutation of the DI container. +12. Provide deterministic behavior and sufficient logging for diagnosis. +13. Preserve independently DI-registered `IMcpTool` extensions across registry generations. + +## Non-Goals + +This work does not: + +- Dynamically enable MCP when it was disabled at application startup. +- Dynamically disable or unmap an MCP endpoint. +- Dynamically change `runtime.mcp.path`. +- Rebuild ASP.NET Core endpoint routing or middleware. +- Change initialize instructions for an already-established MCP session. +- Broadcast HTTP tool-list notifications through experimental MCP SDK session APIs. +- Make the complete DAB hot-reload pipeline transactional. +- Solve cross-component generation isolation for every DAB hot-reload consumer. +- Change MCP authorization semantics. +- Change the wire shape of existing tools. + +Changes to startup-bound MCP settings continue to require a process restart. Improving global hot-reload rollback and transactionality is tracked separately. + +## Design Decisions + +### 1. `McpToolRegistry` remains a singleton + +HTTP and stdio request handlers already retain a reference to the registry. Keeping one singleton avoids rebuilding MCP servers, handlers, transports, or service providers. + +The singleton no longer exposes a dictionary that is incrementally mutated during normal operation. Instead, it holds one current immutable snapshot reference. + +### 2. Registry generations are immutable snapshots + +The registry snapshot conceptually contains: + +```csharp +internal sealed record McpToolRegistrySnapshot( + long Version, + ImmutableDictionary Tools, + int AdvertisedToolCount, + string DiscoveryJson, + string DiscoveryCanonicalJson); +``` + +`Tools` contains: + +- Every built-in tool, including built-ins currently disabled by DML tool configuration. +- Every independently DI-registered `IMcpTool` implementation. +- Every configuration-generated custom tool enabled in the configuration used to build the snapshot. + +`DiscoveryJson` contains precomputed metadata for tools whose `IsEnabled(config)` result was true +for that same configuration generation. Tools are sorted deterministically by name, while nested +schema-property insertion order is preserved for clients that render parameters in wire order. +`DiscoveryCanonicalJson` is a separate recursively property-sorted representation used only for +semantic change comparison. `AdvertisedToolCount` avoids retaining a duplicate object graph solely +for diagnostics. + +Keeping lookup state and advertised metadata in the same snapshot prevents a request from combining tools from one generation with enablement or metadata from another generation. + +Protocol `Tool` objects are mutable SDK models, so the registry defensively clones metadata during +candidate construction and again when returning public discovery results. Candidate publication +pre-serializes the order-preserving discovery representation, which the accessor deserializes to +produce caller-owned clones without reserializing every tool per request. The separate canonical +representation is never served. Neither a tool retaining its source metadata object nor a caller +mutating a returned object can modify a published snapshot. + +Tool names must be nonempty and must not contain leading or trailing whitespace. Rejecting rather +than trimming guarantees that every exact name returned by `tools/list` resolves through +`TryGetTool()`. + +### 3. Publication is atomic + +A candidate snapshot is built completely before the live registry is changed. The registry publishes it with a single `Interlocked.Exchange` or equivalent atomic reference swap. + +Readers capture the current snapshot once per operation: + +- `tools/list` deserializes caller-owned metadata from one snapshot's order-preserving discovery JSON. +- `tools/call` resolves a tool from `Tools` in one snapshot. + +Readers do not acquire the rebuild lock. They observe either the complete previous snapshot or the complete replacement snapshot. + +### 4. DI-owned and configuration-generated tool lifetimes differ + +Built-in tools remain DI-owned application singletons because they are stateless and their execution paths already read current request/configuration state. Independently registered custom `IMcpTool` implementations are also DI-owned and remain part of every candidate; `IMcpTool` remains an extension point regardless of the implementation's `ToolType` value. + +Only `DynamicCustomTool` objects generated from entity configuration are removed from automatic DI registration. They are configuration-generation objects and are recreated for every registry candidate. A name collision between an independently registered tool and a generated tool rejects the candidate rather than silently discarding either implementation. + +This avoids treating the immutable DI service collection as a dynamic registry. + +### 5. Refresh orchestration is separate from state storage + +A singleton `McpToolRegistryRefreshService` coordinates initialization and hot-reload. It also implements `IHostedService` for normal HTTP-host startup. + +Its responsibilities are: + +1. Capture the current `RuntimeConfig` generation. +2. Obtain all DI-owned `IMcpTool` implementations, including independently registered custom tools. +3. Create fresh configuration-generated custom tools from the captured configuration. +4. Enrich custom tool schemas from refreshed database metadata. +5. Ask the registry to validate and build a complete candidate snapshot. +6. Verify that the captured configuration is still current. +7. Atomically publish the candidate. +8. Notify configured transports if advertised metadata changed. +9. Log success or failure. + +`McpToolRegistry` owns registry invariants and publication. The refresh service owns lifecycle and dependencies. The bulk construction and publication surface is internal to the MCP runtime assembly rather than an advertised embedding API. + +### 6. Custom tool creation is strict + +`CustomMcpToolFactory` previously caught broad exceptions and skipped individual entities. Retaining +that behavior would allow a partial candidate to be published. + +For registry initialization and refresh: + +- Unexpected custom tool construction failures reject the complete candidate. +- The exception identifies the source entity. +- Empty names and case-insensitive collisions reject the candidate. +- Collisions are checked across built-in and custom tools. +- No candidate tool is silently omitted because construction failed. + +Database metadata unavailability is a deliberate exception to this strict behavior. `DynamicCustomTool` already supports a configuration-derived schema fallback. The candidate may use that fallback, but the reason must be logged so reduced schema accuracy is visible. + +### 7. Metadata initialization uses explicit dependencies + +The metadata initialization path receives explicit dependencies: + +```csharp +void InitializeMetadata( + RuntimeConfig config, + IMetadataProviderFactory metadataProviderFactory); +``` + +`McpMetadataHelper` has an overload that accepts `IMetadataProviderFactory` directly. Execution call sites retain the service-provider overload where resolving request services is appropriate. + +This ensures that: + +- Configuration-derived metadata belongs to the captured generation. +- Database metadata comes from the factory already refreshed earlier in the ordered pipeline. +- Candidate construction does not resolve arbitrary application services. + +### 8. MCP receives a dedicated ordered hot-reload event + +Add a named event such as: + +```text +MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED +``` + +The event is added to `DabConfigEvents` and `HotReloadEventHandler`, then raised by `RuntimeConfigLoader` after `AUTHZ_RESOLVER_ON_CONFIG_CHANGED` and before GraphQL schema events. + +The order is intentional: + +1. Query and metadata dependencies are refreshed first. +2. Query and mutation engines are refreshed. +3. Authorization state is refreshed. +4. MCP builds and publishes tools that depend on those services. +5. GraphQL performs its independent schema lifecycle. + +The refresh callback catches and logs hot-reload failures so an MCP candidate failure does not prevent later hot-reload handlers from running. Startup initialization remains strict because there is no previous valid registry to preserve. + +### 9. Initialization is shared and idempotent + +The refresh service exposes one idempotent initialization path. + +Direct `EnsureInitialized()` calls remain no-ops after the current `RuntimeConfig` reference has +successfully been applied. The ordered hot-reload event is intentionally different: it always +rebuilds the current configuration because the configuration becomes visible before its metadata +event runs. This distinction prevents an out-of-band initialization during that interval from +permanently publishing the new configuration with the previous metadata generation. + +#### HTTP mode + +The host resolves the singleton through `IHostedService` early enough to subscribe to ordered +hot-reload events, but `StartAsync()` does not publish the initial snapshot. ASP.NET Core starts +hosted services before `Startup.Configure` finishes initializing database metadata. + +`Startup.PerformOnConfigChangeAsync()` invokes a shared runtime-initialization helper. Configuration +capture and validation, `IMetadataProviderFactory` initialization, and the refresh service's +idempotent registry publication execute as one asynchronous operation under the +`FileSystemRuntimeConfigLoader` serialization gate. A file callback therefore cannot replace the +active configuration between those steps. + +The DI registration must ensure that resolving the concrete refresh service and resolving `IHostedService` return the same object, for example by registering the concrete singleton and mapping `IHostedService` to it. + +#### Stdio mode + +Stdio intentionally does not start the ASP.NET Core host, so `Startup.Configure` does not initialize +database metadata. `McpStdioHelper` invokes the same serialized initial dependency operation used by +HTTP startup before starting the stdio loop. + +This replaces the former duplicate per-tool registration path and gives both transports identical validation and metadata behavior. + +### 10. A stale candidate is never published + +Distinct file edits can produce overlapping hot-reload callbacks even though duplicate notifications for one file content are suppressed by `ConfigFileWatcher`. + +`FileSystemRuntimeConfigLoader` serializes initial dependency construction and the complete reload +operation per loader instance. Its async-capable gate is held across initial configuration capture +and validation, metadata initialization, and MCP registry publication. For reload, the gate is +acquired before loading the new configuration and remains held until every synchronous +`SignalConfigChanged()` handler returns. Consequently, one complete generation finishes before +another path can replace the active configuration or begin updating dependencies. + +#### Serialization layers and lock ordering + +The refresh service retains its own writer gate and stale-generation guard as defense in depth: + +1. Acquire the refresh writer gate. +2. Capture `RuntimeConfig config = runtimeConfigProvider.GetConfig()`. +3. Build the candidate against `config`. +4. Before publication, verify that `runtimeConfigProvider.GetConfig()` is still the same configuration object. +5. If it changed, discard the candidate without notifying clients. + +A callback for the newer configuration will build the latest snapshot. The service tracks only successfully applied configuration references, so a later event can retry after an earlier failure. + +The loader gate prevents mixed dependency generations within the normal serialized startup/reload +path. An out-of-band `EnsureInitialized()` can still run after a new configuration is installed but +before that configuration's metadata event. Such a call may publish an intermediate mixed snapshot. +The later ordered MCP event therefore bypasses config-reference idempotency and rebuilds after +metadata and authorization refresh, replacing the intermediate snapshot. The stale guard also +prevents an older, slower registry rebuild initiated outside the file-loader pipeline from +overwriting a newer registry generation. Neither mechanism provides transactional rollback after a +handler failure; that remains separate work. + +The two locks protect different invariants. The loader gate spans configuration publication and all +ordered component callbacks. The refresh lock spans only registry candidate creation and +publication, including direct `EnsureInitialized()` calls that do not own the loader gate. Normal +startup and file reload acquire them in loader-then-refresh order. The refresh service never calls +back into an operation that acquires the loader gate while holding its own lock, and transport +notification occurs after the refresh lock is released. This fixed ownership prevents lock-order +cycles while retaining defense against out-of-band callers. + +Reference identity is deliberately checked after candidate construction, immediately before +publication. Checking only before construction would not detect a configuration replacement that +occurs while tools and metadata are being materialized. + +#### Shutdown contract + +Loader shutdown does not wait to acquire this gate. It atomically stops admission, cancels callbacks +waiting to enter the gate, and requests cooperative cancellation of the active generation. +Cancellation is propagated through every DAB-owned ordered handler and metadata database operation, +including connection opening, schema discovery, query execution, access-token acquisition, and the +otherwise synchronous `FillSchema` command. Cancellation is checked before later MCP, GraphQL, +authorization, and logging publications, so a canceled owned generation does not continue through +the remaining pipeline. + +The loader tracks each operation that owns the serialization gate and exposes an idempotent +`StopAsync(CancellationToken)` drain. An `IHostedService` invokes that drain during the host stopping +phase, before the root service provider disposes reload subscribers or their dependencies. The +host-supplied token bounds the drain according to `HostOptions.ShutdownTimeout`. Stdio composition, +which constructs but does not start the ASP.NET Core host, explicitly applies the same configured +timeout before disposing its host. Synchronous `FileSystemRuntimeConfigLoader.Dispose()` only stops +admission and requests cancellation; it does not reintroduce an unbounded wait after a host timeout. + +.NET cannot forcibly terminate arbitrary synchronous subscriber code. A successful drain guarantees +that no reload operation remains and dependency disposal is safe. If the host timeout expires, the +host stops waiting; DAB-owned subscribers are designed to observe cancellation, while extension +subscribers are contractually required to observe `HotReloadEventArgs.CancellationToken` and avoid +indefinite blocking. Isolating non-cooperative extension callbacks from root-provider lifetime would +require a separately owned dependency container and is outside this feature's scope. + +The resulting guarantees are: + +| Shutdown path | Guarantee | +|---|---| +| Active DAB-owned handler observes cancellation | `StopAsync()` waits for the gate owner and cancellation callbacks to exit, then dependencies may be safely disposed. | +| Queued reload waiting for the loader gate | The loader token cancels the wait; the callback never becomes an active generation. | +| Extension handler ignores cancellation but exits before the host timeout | The drain still completes and dependency disposal remains ordered after the handler. | +| Extension handler remains blocked when the host timeout expires | `StopAsync()` observes the host token and returns cancellation; the host may continue disposal. No stronger lifetime guarantee is possible for arbitrary in-process synchronous code. | +| Direct synchronous `Dispose()` | New work is rejected and cancellation starts, but no drain is promised. A direct owner requiring dependency ordering must call `StopAsync()` first. | + +Cancellation callbacks are also extension points and can block. Shutdown uses `CancelAsync()` and +tracks callback completion rather than running callbacks inline on the host stopping thread. The +tracked callback task participates in a successful drain, while the caller's shutdown token still +bounds the wait. + +`SemaphoreSlim` and `CancellationTokenSource` are loader-owned disposable resources. Cleanup cannot +run merely when the gate owner exits because a canceled waiter may still be unwinding from `Wait()`. +The loader therefore counts every operation admitted before shutdown, including gate waiters. The +last operation releases the gate before marking itself drained. A single shutdown-completion task +then waits for both admitted operations and `CancelAsync()` callbacks before disposing the semaphore +and token source. `StopAsync()` joins that task; synchronous `Dispose()` starts the same task and +returns without blocking. If the host times out behind non-cooperative work, cleanup remains pending +and occurs when that work eventually exits rather than racing it. + +#### Watcher callback and resource-lifetime decision + +`ConfigFileWatcher` invokes the reload entry point synchronously rather than creating a detached +task. This preserves the existing ordered event contract and ensures every admitted reload is +tracked as either a gate waiter or gate owner. File-system implementations may invoke another +callback concurrently; the loader gate serializes those callbacks and shutdown cancels queued +waiters. Introducing another queue would require separate task ownership, ordering, coalescing, and +drain rules without improving registry atomicity. + +Watcher callback admission is disabled synchronously under a separate watcher-lifecycle lock. +Potentially blocking operating-system watcher resource disposal is scheduled on a background worker; +this resource cleanup is independent of the reload-operation drain and does not retain access to the +reload subscribers. + +### 11. Existing tool-call safety is preserved + +After a successful swap: + +- Removed custom tools no longer resolve for new calls. +- Renamed custom tools resolve only under the new name. +- Disabled built-in tools remain in the lookup map, preserving current behavior in which execution returns a structured tool-disabled result. + +A request that resolved a tool immediately before a swap may finish with that tool instance. `DynamicCustomTool.ExecuteAsync()` still validates the current configuration, entity type, custom-tool enablement, database metadata, and authorization before execution. This makes retirement of old custom tool objects safe without explicit cancellation or disposal. + +### 12. Stdio sends tool-list change notifications + +Production stdio composition registers a tool-list notifier and advertises +`tools.listChanged = true`. Alternative composition without that notifier advertises `false` so the +initialize response never promises a notification path that is unavailable. + +After a successful noninitial refresh, send: + +```json +{ + "jsonrpc": "2.0", + "method": "notifications/tools/list_changed", + "params": {} +} +``` + +Notification rules: + +- Do not notify for initial registry construction. +- Do not notify before the server successfully completes the `initialize` response and the client + subsequently sends `notifications/initialized`. +- Do not queue a missed pre-initialization notification; the client has not yet established its cache and will request the initial list. +- Notify only when the advertised tool list or metadata changed. +- Send after atomic publication. +- Invoke transport notifiers after releasing the registry writer lock. +- Check initialization state before enqueueing delivery, then perform potentially blocking stdout + I/O on a worker so the reload pipeline can continue to later handlers. +- Route the frame through the shared `McpStdoutWriter` so it cannot interleave with responses or logging notifications. +- Notification write failure is logged and does not roll back the registry. + +A small stdio notifier service owns initialization state and queued frame writing. Multiple changes +while one stdout write is pending may be coalesced because any delivered invalidation causes the +client to request the latest complete snapshot. `McpStdioServer` tracks successful initialize-response +completion for the connection and marks the notifier initialized only when a subsequent +`notifications/initialized` arrives; an out-of-order notification is ignored. The refresh service +depends on zero or more tool-list notifiers; HTTP mode has no notifier registered in this iteration. +If the thread pool rejects the initial worker request, the notifier retains the pending invalidation +and starts one dedicated background fallback worker. This rare fallback keeps reload callbacks +nonblocking and avoids losing the only invalidation when no later configuration change occurs. +Because an abandoned client can leave stdout blocked indefinitely, `McpStdoutWriter.Dispose()` +first rejects future writes and releases the underlying writer only when its serialization lock is +immediately available. Host disposal therefore does not wait behind a blocked notification; in that +exceptional case the process-owned stdout handle is reclaimed when the process exits. + +### 13. HTTP reads are immediately current, but HTTP push is deferred + +The HTTP MCP SDK handlers execute against the registry singleton per request. Once the snapshot is swapped, the next `tools/list` and `tools/call` request sees it without rebuilding the MCP server. + +HTTP `listChanged` capability is explicitly set to false and must not be advertised as supported +until HTTP notification delivery is implemented. + +The installed MCP SDK can send a notification through an individual `McpServer` session, but broadcasting requires tracking all active sessions through an experimental `RunSessionHandler`. Depending on that experimental API is not necessary for registry correctness and is deferred to focused follow-up work. + +### 14. Notify only for a semantic discovery change + +Every applicable configuration hot-reload rebuilds and publishes a generation so custom tool instances align with the current configuration. However, an unrelated configuration change should not claim that the tool list changed. + +The registry compares the previous and candidate advertised metadata in deterministic name order. +Before comparison it canonicalizes serialized JSON recursively by sorting object properties while +sorting primitive-string arrays only for the order-insensitive JSON Schema set keywords `required`, +`type`, and `enum`. All other arrays preserve order because values such as an array-valued `default` +or `examples` entry can be order-sensitive advertised metadata. The comparison therefore ignores +semantically irrelevant object insertion and schema-set order while still covering the complete +tool metadata, including name, description, input schema, and any future advertised fields. + +Canonical JSON is comparison-only. The separately stored discovery representation preserves nested +object insertion order, including stored-procedure parameter order, so canonicalization does not +introduce a wire-order behavior change for clients that render schema properties in received order. + +The swap still occurs when advertised metadata is equal, but `notifications/tools/list_changed` is emitted only when discovery metadata differs. + +## Registry API Behavior + +The production path changes from incremental registration to bulk replacement. + +Conceptual operations are: + +```csharp +IReadOnlyList GetAdvertisedTools(); + +bool TryGetTool(string toolName, out IMcpTool? tool); + +internal McpToolRegistryUpdateResult ReplaceAll( + IEnumerable tools, + RuntimeConfig config); +``` + +`ReplaceAll`: + +1. Materializes the input once. +2. Retrieves and validates metadata once per tool. +3. Builds the case-insensitive lookup map. +4. Evaluates `IsEnabled(config)` for advertised metadata. +5. Sorts advertised metadata deterministically. +6. Compares advertised metadata with the current snapshot. +7. Atomically publishes the candidate. +8. Returns the new version and whether discovery metadata changed. + +The registry no longer exposes incremental registration or caller-configured filtering. Those +helpers and the former startup initializer were implementation plumbing in the MCP runtime assembly, +not documented APIs from a supported reference package. Removing them leaves one construction model: +the refresh service builds and atomically publishes a complete configuration-aware generation. +Production discovery uses `GetAdvertisedTools()`, whose visibility and metadata were captured with +that same snapshot. + +### Intentional cleanup of CLR-public implementation members + +This change deliberately removes or narrows members that happened to be declared `public`, including +the `RegisterTool()` overloads, `GetEnabledTools()`, `InitializeAndRegisterTools()`, +`DynamicCustomTool.InitializeMetadata(IServiceProvider)`, and the former +`McpToolRegistryInitializer`. `CustomMcpToolFactory.CreateCustomTools()` now exposes its concrete +generated-tool result type, and bulk candidate construction/publication is `internal`. These are +intentional source/API-surface changes, not accidental compatibility omissions. + +Those members were used only by DAB's own startup and test implementation. They were not documented +as extension APIs, and `Azure.DataApiBuilder.Mcp.dll` is runtime payload of the DAB tool rather than a +supported reference package. The supported `Microsoft.DataApiBuilder.Core` embedding package does +not expose the MCP implementation assembly. A consumer that manually referenced the runtime DLL and +called these methods was depending on unsupported internals solely because their CLR visibility was +too broad. Retaining obsolete shims would preserve two conflicting construction models and undermine +the atomic-snapshot invariant, so the implementation-only surface is removed instead. + +Cancellation-aware overloads added to the supported `Microsoft.DataApiBuilder.Core` interfaces are +different: they use default interface implementations that check pre-cancellation and delegate to +the established parameterless members. Existing third-party implementations therefore remain +source- and binary-compatible. DAB's built-in implementations override the overloads to propagate +cancellation through database I/O. + +### Cancellation API compatibility and legacy behavior + +The cancellation additions intentionally extend rather than replace established contracts: + +- `IMetadataProviderFactory`, `ISqlMetadataProvider`, and `IQueryExecutor` retain all legacy + members. Their new token overloads are default interface methods that reject an already canceled + call and otherwise delegate to the legacy member. +- A legacy third-party implementation therefore compiles and runs unchanged. Once its legacy call + begins, DAB cannot impose cooperative cancellation on work that does not accept a token; this is + the compatibility trade-off. DAB-owned implementations override the new members and carry the + token to connection opening, commands, readers, retries, and access-token acquisition. +- `QueryExecutor.ExecuteQueryAsync()` combines an explicit operation token with + `HttpContext.RequestAborted` when both are cancellable, so either owner can stop retries and + database I/O. Calls through the established tokenless overload still dispatch through the + established virtual database-execution member; this preserves existing subclasses and Moq + setups while that member continues to consume `RequestAborted`. +- `HotReloadEventArgs` adds a read-only token and a three-argument constructor while retaining the + original two-argument constructor. Existing compiled and source callers remain valid, while the + file loader uses the new overload with its owned shutdown token. +- Except for the two explicitly documented `SqlMetadataProvider` metadata-I/O slots below, existing + public and protected Config/Core members retain their original signatures and virtual slots. + `TryLoadConfig()`, `SignalConfigChanged()`, `PopulateTriggerMetadataForTable()`, + `GenerateAutoentitiesIntoEntities()`, and `QueryAutoentitiesAsync()` delegate to or are invoked + by explicit token-aware overloads. Existing compiled callers and derived providers therefore do + not need to recompile or add overrides for those members. +- `IMcpToolRegistryRefreshService` retains parameterless `EnsureInitialized()` and provides the + token overload through a default implementation. The shared startup path uses the token overload; + existing test or embedding implementations that only provide the original member continue to + work. +- The otherwise synchronous provider `FillSchema()` call cannot be made truly asynchronous. DAB + registers cancellation to invoke `DbCommand.Cancel()` and converts a provider exception observed + after cancellation into `OperationCanceledException`. This is best-effort and remains dependent + on the database provider honoring command cancellation. + +#### Intentional replacement of two protected metadata-provider virtual slots + +`SqlMetadataProvider` is a public class in the supported +`Microsoft.DataApiBuilder.Core` package. This design therefore explicitly acknowledges that replacing +these protected virtual members is observable to a consumer that subclasses the implementation: + +```csharp +// Previous slots +FillSchemaForStoredProcedureAsync( + Entity, string, string, string, StoredProcedureDefinition) + +GetColumnsAsync(string, string) + +// Replacement slots +FillSchemaForStoredProcedureAsync( + Entity, string, string, string, StoredProcedureDefinition, CancellationToken) + +GetColumnsAsync(string, string, CancellationToken) +``` + +This is an intentional and narrowly scoped source/binary compatibility break, not an accidental +omission. It is accepted for the following reasons: + +1. Neither member appears on `ISqlMetadataProvider`, in the Core package usage documentation, or in + a supported custom-database-provider registration contract. `MetadataProviderFactory` constructs + a closed set of DAB-owned providers from `DatabaseType`; it has no provider plug-in registration + point. The only overrides in this repository are the DAB-owned SQL Server and MySQL providers. +2. These methods are exactly where potentially long-running schema calls acquire credentials, open + connections, and query provider metadata. Bounded shutdown requires the loader token to reach + those operations; prechecking a token only at the caller is insufficient once I/O has begun. +3. Keeping a tokenless overload but invoking only the new overload would preserve metadata shape + while silently bypassing an existing override. Invoking the old override would preserve dispatch + but create an uncancelable hole in the shutdown path. Either shim would imply a compatibility + guarantee that the runtime could not honor together with end-to-end cancellation. +4. Normal package consumers, embedders that use `ISqlMetadataProvider`, and implementations of the + public Core interfaces are unaffected. The impact is limited to consumers that directly derived + from this implementation class and overrode one of these undocumented provider-internal hooks. + +Migration for such a subclass is mechanical: recompile against the new Core package, add the +`CancellationToken` parameter to the override, and pass it through credential acquisition, +`OpenAsync()`, and schema/query APIs. A future supported custom-provider SPI should define lifecycle +and cancellation as part of its initial contract rather than relying on protected implementation +hooks. + +`FileSystemRuntimeConfigLoader.StopAsync(CancellationToken)` is public because the Service assembly +and direct hosts must coordinate Config-owned work before disposing their own dependency graph. It +is a lifecycle operation, not a replacement for `Dispose()`; repeated calls join the same shutdown +and may use different caller-side timeout tokens. + +## Detailed Flows + +### Initial HTTP startup + +```mermaid +sequenceDiagram + participant Host + participant Startup + participant Refresh as McpToolRegistryRefreshService + participant Factory as CustomMcpToolFactory + participant Metadata as IMetadataProviderFactory + participant Registry as McpToolRegistry + + Host->>Refresh: StartAsync (subscribe only) + Startup->>Metadata: InitializeAsync(loader cancellation token) + Metadata-->>Startup: DB metadata ready + Startup->>Refresh: EnsureInitialized(loader cancellation token) + Refresh->>Refresh: Capture current RuntimeConfig + Refresh->>Factory: Create custom tools(config) + Factory-->>Refresh: Fresh custom tools + Refresh->>Metadata: Resolve refreshed DB metadata + Refresh->>Refresh: Initialize custom schemas + Refresh->>Registry: ReplaceAll(built-ins + custom, config) + Registry-->>Refresh: Published version + Note over Refresh,Registry: No list-changed notification on initial construction +``` + +An invalid or duplicate tool causes startup to fail, preserving current strict startup behavior. + +### Initial stdio startup + +```mermaid +sequenceDiagram + participant Helper as McpStdioHelper + participant Refresh as McpToolRegistryRefreshService + participant Metadata as IMetadataProviderFactory + participant Registry as McpToolRegistry + participant Server as McpStdioServer + + Helper->>Metadata: InitializeAsync(loader cancellation token) + Metadata-->>Helper: DB metadata ready + Helper->>Refresh: EnsureInitialized(loader cancellation token) + Refresh->>Registry: Build and publish initial snapshot + Helper->>Server: RunAsync() + Server->>Server: Handle initialize + Server->>Server: Handle notifications/initialized + Server->>Server: Mark tool-list notifier initialized +``` + +### Successful hot-reload + +```mermaid +sequenceDiagram + participant Loader as RuntimeConfigLoader + participant Metadata as MetadataProviderFactory + participant Auth as AuthorizationResolver + participant Refresh as McpToolRegistryRefreshService + participant Registry as McpToolRegistry + participant Notifier as Stdio notifier + + Loader->>Metadata: METADATA_PROVIDER_FACTORY_ON_CONFIG_CHANGED + Metadata-->>Loader: Metadata refreshed + Loader->>Auth: AUTHZ_RESOLVER_ON_CONFIG_CHANGED + Auth-->>Loader: Authorization refreshed + Loader->>Refresh: MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED + Refresh->>Refresh: Capture config and build candidate + Refresh->>Refresh: Verify captured config is still current + Refresh->>Registry: Atomic ReplaceAll + Registry-->>Refresh: Published with discovery changes + Refresh->>Notifier: NotifyToolsListChanged() +``` + +### Failed hot-reload candidate + +```mermaid +sequenceDiagram + participant Loader as RuntimeConfigLoader + participant Refresh as McpToolRegistryRefreshService + participant Registry as McpToolRegistry + + Loader->>Refresh: MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED + Refresh->>Refresh: Build candidate + Refresh--xRefresh: Invalid name, collision, or construction failure + Refresh->>Refresh: Log error with entity/tool context + Note over Registry: Previous immutable snapshot remains active + Refresh-->>Loader: Return without throwing +``` + +## Concurrency Model + +### Reads + +Registry reads capture one snapshot reference and do not lock. Immutable lookup and metadata collections are safe for concurrent requests. + +### Writes + +Initialization and refresh operations use one writer gate. Candidate construction happens while registry rebuilds are serialized, but the live snapshot remains available to readers. + +### In-flight requests + +An in-flight `tools/list` serializes the snapshot it captured. It returns either the complete old list or complete new list. + +An in-flight `tools/call` retains the resolved tool instance. A later swap does not invalidate that object. Dynamic custom tools revalidate current configuration and authorization before database execution. + +### Multiple configuration changes + +The per-loader gate ensures initial metadata and registry construction cannot overlap a file reload, +and that one file generation completes all ordered handlers before the next notification begins +loading. The stale-generation guard additionally prevents publication when the active `RuntimeConfig` +changes during candidate construction through another code path. The latest callback eventually +publishes the latest generation. + +Serialization is scoped to each `FileSystemRuntimeConfigLoader`; independent loaders do not block one +another. Transactional rollback is still tracked separately. + +## Failure Semantics + +| Scenario | Registry result | Client notification | Hot-reload pipeline | +|---|---|---|---| +| Initial construction succeeds | Initial snapshot published | None | Startup continues | +| Initial construction fails | No usable snapshot | None | Startup fails | +| Hot-reload construction succeeds and metadata changes | New snapshot published | Stdio notification after initialization | Continues | +| Hot-reload construction succeeds with equivalent metadata | New snapshot published | None | Continues | +| Custom tool construction fails | Previous snapshot retained | None | Error logged; continues | +| Tool name is invalid or duplicated | Previous snapshot retained | None | Error logged; continues | +| DB metadata is unavailable but config fallback works | New snapshot published with fallback schema | Notify if metadata changed | Warning logged; continues | +| Candidate becomes stale before publication | Candidate discarded | None | Newer callback is expected to refresh | +| Stdio notification write fails | New snapshot remains published | Delivery failed | Error logged; continues | + +### Temporary limitation before transactional hot-reload + +The current DAB hot-reload pipeline commits the new `RuntimeConfig` before component refresh callbacks complete. If MCP candidate construction fails, DAB can temporarily have a newer runtime configuration and metadata generation with the previous MCP registry snapshot. + +This design chooses the safest local behavior: + +- Never publish a partial or invalid registry. +- Keep the previous discovery snapshot. +- Let stale custom tools fail safely through current execution-time validation. +- Retry on a later configuration event. +- Log the degraded condition clearly. + +Once transactional hot-reload exists, MCP candidate construction should become a transaction participant and reject the complete configuration candidate before any component publishes it. + +## Configuration Behavior + +### Changes applied live + +The following changes are reflected after a successful registry refresh: + +- Adding a stored-procedure entity with `mcp.custom-tool: true`. +- Removing such an entity. +- Enabling or disabling `mcp.custom-tool` on a stored-procedure entity. +- Renaming a custom-tool entity and therefore its normalized tool name. +- Changing an entity description. +- Changing configuration-declared stored-procedure parameter metadata. +- Changing DB-discovered stored-procedure parameter metadata. +- Changing global built-in DML tool enablement flags. +- Introducing or resolving a custom/custom or custom/built-in name collision. + +### Changes that remain startup-bound + +- `runtime.mcp.enabled`. +- `runtime.mcp.path`. +- HTTP route registration. +- Existing session initialization instructions. + +The actual route and service registration remain those selected at startup. If these options are modified in a hot-reloaded file, a restart is required for them to take effect consistently. + +## Dependency Injection Changes + +The intended registrations are: + +- `McpToolRegistry`: singleton. +- Built-in `IMcpTool` implementations: singleton. +- Independently registered custom `IMcpTool` implementations: DI-owned and retained across generations. +- `McpToolRegistryRefreshService`: singleton. +- `IHostedService`: resolves the same refresh-service singleton. +- Configuration-generated `DynamicCustomTool` instances: not registered in DI. +- Stdio tool-list notifier: singleton, registered only in stdio mode. + +The refresh service preserves every implementation in its DI-provided `IEnumerable`. Reflection-based discovery of DAB's own implementations continues to exclude `DynamicCustomTool`; those instances come only from the per-generation factory during normal startup and reload. + +## Implementation Map + +The implementation is split across the following touchpoints: + +### Config project + +- [DabConfigEvents.cs](../../src/Config/DabConfigEvents.cs): add the MCP registry event name. +- [HotReloadEventHandler.cs](../../src/Config/HotReloadEventHandler.cs): register the event slot. +- [RuntimeConfigLoader.cs](../../src/Config/RuntimeConfigLoader.cs): raise the event at the agreed position. +- [FileSystemRuntimeConfigLoader.cs](../../src/Config/FileSystemRuntimeConfigLoader.cs): serialize initial dependency construction and complete file-reload pipelines with one async-capable per-loader gate. + +### MCP project + +- [McpToolRegistry.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs): immutable snapshots, bulk replacement, and atomic reads/publication. +- Remove the former `McpToolRegistryInitializer`; the refresh service owns initialization. +- [McpServiceCollectionExtensions.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpServiceCollectionExtensions.cs): register the shared refresh service and stop registering custom tools in DI. +- [CustomMcpToolFactory.cs](../../src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs): strict candidate creation. +- [DynamicCustomTool.cs](../../src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs): explicit metadata dependencies and removal of the stale-metadata assumption. +- [McpMetadataHelper.cs](../../src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs): optional explicit metadata-factory overload. +- [McpServerConfiguration.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs): list directly from one registry snapshot. +- [McpStdioServer.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs): list from the snapshot and mark notification readiness. +- New stdio tool-list notifier type or types. + +### Service project + +- [RuntimeInitializationHelper.cs](../../src/Service/Utilities/RuntimeInitializationHelper.cs): coordinate serialized configuration validation, metadata initialization, and initial registry publication for both transports. +- [McpStdioHelper.cs](../../src/Service/Utilities/McpStdioHelper.cs): invoke shared idempotent initialization. +- [Program.cs](../../src/Service/Program.cs): register the stdio notifier with the shared stdout writer. + +### Tests + +- Extend registry, factory, dynamic-tool, stdio, and hot-reload tests. +- Add focused refresh-service and concurrency tests. + +## Testing Strategy + +### Registry unit tests + +1. Bulk replacement publishes unique tools. +2. Names are case-insensitive. +3. Empty and whitespace names reject the candidate. +4. Built-in/custom and custom/custom collisions reject the candidate. +5. A rejected replacement leaves the exact previous snapshot active. +6. Added custom tools become discoverable and callable. +7. Removed custom tools disappear from lookup and discovery. +8. Renamed custom tools remove the old name and add the new name atomically. +9. Built-in visibility is computed from the candidate configuration. +10. Advertised tools have deterministic ordering. +11. Equivalent metadata does not report a discovery change. +12. Name, description, input-schema, addition, removal, and visibility changes report a discovery change. +13. Concurrent readers observe no exceptions or partial snapshots during repeated swaps. +14. Served input-schema properties preserve their source insertion order even though comparison is canonicalized. +15. Reordering string values in JSON Schema `required`, `type`, or `enum` arrays does not report a discovery change. +16. Reordering an order-sensitive primitive array such as an array-valued `default` still reports a discovery change. + +### Refresh-service unit tests + +1. Initial construction uses all DI-owned tools and newly created configuration-generated tools. +2. Every refresh creates new custom tool instances. +3. Built-in instances are reused. +4. Metadata initialization uses the captured configuration and refreshed metadata factory. +5. A fallback schema is published and logged when DB metadata is unavailable. +6. Construction failure preserves the previous registry. +7. Startup failure propagates. +8. Hot-reload failure is caught and logged. +9. A stale candidate is discarded. +10. Repeated direct initialization for an already successfully applied configuration does not publish duplicate generations unnecessarily. +11. An ordered MCP event rebuilds after refreshed metadata even when an out-of-band initialization already applied the same configuration reference. +12. Notifications occur only after a successful noninitial semantic discovery change. +13. Independently DI-registered custom implementations remain published after refresh. + +### Handler and transport tests + +1. HTTP `tools/list` reads only registry snapshot metadata. +2. HTTP `tools/call` resolves from the current snapshot. +3. Stdio `tools/list` reads only registry snapshot metadata. +4. Stdio ignores an out-of-order `notifications/initialized` and does not notify before a complete + successful initialization handshake. +5. Stdio does not notify for initial construction. +6. Stdio emits the exact `notifications/tools/list_changed` frame after an applicable refresh. +7. Stdio serializes notifications through `McpStdoutWriter` without interleaving. +8. Stdio notification failure does not revert the registry. +9. HTTP does not advertise `listChanged` in this iteration. +10. A blocked stdio writer does not block the ordered hot-reload pipeline. +11. The real HTTP handler omits tools disabled in the published registry snapshot. +12. The physical HTTP failure-path test observes completion of the rejected candidate before + checking retained discovery and writing a recovery configuration. +13. Stdio advertises `listChanged = false` when no stdio notifier is composed. +14. A rejected primary worker schedule retains and delivers the pending notification through the fallback worker. +15. Multiple changes while a write is blocked coalesce to one additional pending notification. +16. Disposing the shared stdout writer while a notification write is blocked returns immediately, + rejects later writes, and does not corrupt the in-flight frame when the pipe resumes. + +### Hot-reload integration tests + +1. The MCP registry event runs after metadata and authorization refresh. +2. Enabling a custom tool makes it appear after a real configuration reload. +3. Disabling or removing a custom tool removes it. +4. Description changes are reflected. +5. Stored-procedure input-schema changes are reflected. +6. Built-in DML visibility changes are reflected. +7. A duplicate name leaves the previous registry active. +8. Correcting a failed configuration allows the next refresh to succeed. +9. Rapid successive configurations cannot publish an older registry after a newer one. +10. A reload paused before metadata refresh cannot overlap initial metadata and registry construction; + the final advertised schema comes from the reload generation's database metadata. +11. A physical stdio config-file write traverses the complete ordered pipeline, emits exactly one + notification for one net-new file content, and returns updated discovery. +12. Shutdown cancels an active cancellation-aware reload handler, prevents later ordered handlers + from running, cancels callbacks queued on the loader gate, and waits for the active handler to + exit before the drain completes. +13. Hosted shutdown drains active reload work before earlier hosted services stop and before the + root provider disposes reload subscribers or their dependencies. +14. Hosted shutdown observes its supplied cancellation token when a non-cooperative subscriber + prevents the drain from completing. +15. Potentially blocking operating-system watcher disposal does not delay the reload-operation + drain after watcher callbacks have been synchronously disabled and detached. + +Database-backed schema tests should reuse existing MCP stored-procedure fixtures where database metadata is required. Pure membership, collision, notification, and atomicity behavior should remain unit-testable without a live database. + +## Logging and Diagnostics + +Use structured logs from the refresh service for: + +- Initial registry version and built-in/custom/advertised counts. +- Successful hot-reload version and counts. +- Whether advertised discovery metadata changed. +- Candidate discard because a newer configuration became active. +- Config-schema fallback, including entity name and reason. +- Candidate failure, including entity/tool context and exception. +- Notification delivery failure. +- Primary notification-worker scheduling failure and fallback-worker creation failure. + +Do not log connection strings, stored-procedure argument values, or other secrets. + +## Security Considerations + +The registry refresh does not change authentication or authorization policy. + +- Built-in and custom tools continue to authorize at execution time. +- `DynamicCustomTool` continues to validate current entity existence, type, enablement, metadata, and role permissions. +- Publishing a tool in `tools/list` does not bypass execution authorization. +- Retaining an old snapshot after failed refresh does not authorize obsolete execution because current configuration and authorization checks still run. + +## Alternatives Considered + +### Mutate the existing dictionary in place + +Rejected because concurrent reads could observe partial state or race with dictionary mutation. It also makes rollback difficult. + +### Rebuild custom tools on every `tools/list` + +Rejected because it moves validation and metadata work into request handling, repeats work, makes failures request-time failures, and does not naturally solve `tools/call` lookup consistency. + +### Create custom tools lazily on `tools/call` + +Rejected because clients still need accurate discovery metadata and collisions should be rejected before invocation. + +### Subscribe directly to `RuntimeConfigProvider.GetChangeToken()` + +Rejected because that signal occurs before the ordered metadata, engine, and authorization refresh events. The registry could publish schemas derived from stale dependencies. + +### Dynamically add and remove DI registrations + +Rejected because the built service provider is not a dynamic registry. Rebuilding it would create duplicate singleton graphs and lifecycle problems. + +### Rebuild the MCP server or HTTP endpoints + +Rejected because handlers already dereference the registry per request. Rebuilding transports and endpoint routing is unnecessary and would complicate active sessions. + +### Implement HTTP broadcast notifications now + +Deferred because registry correctness does not require it and the current SDK exposes active-session interception through an experimental API. + +### Queue file-watcher reloads onto detached tasks + +Rejected because the existing ordered pipeline defines completion synchronously. A detached queue +would need a second ordering and coalescing model, explicit task ownership, exception observation, +and another shutdown drain. Keeping callbacks synchronous and serializing them at the loader gate +makes every admitted operation observable to shutdown. + +### Wait without a timeout until every reload subscriber exits + +Rejected because extension callbacks are arbitrary synchronous code and can block forever. An +unbounded wait would make `HostOptions.ShutdownTimeout` ineffective. The implementation requests +cooperative cancellation and drains normally, but honors the host's timeout when code does not +cooperate. + +### Dispose dependencies immediately without draining reload work + +Rejected because DAB-owned reload handlers use singleton metadata, query, authorization, and +logging dependencies. Disposing those while a cooperative generation is unwinding creates avoidable +use-after-dispose races. The shutdown hosted service is registered last so it stops first and drains +the loader before earlier hosted services and the root provider stop. + +### Isolate reload subscribers in a separately owned service provider + +Deferred. A child provider could remain alive after the root host timeout and would provide a +stronger lifetime boundary for non-cooperative extensions, but it would duplicate or proxy a large +singleton graph and change existing hot-reload ownership. That is disproportionate to this registry +feature and does not make arbitrary code forcibly cancelable. + +### Run cancellation callbacks inline on the stopping thread + +Rejected because token registrations are extension points and may themselves block. `CancelAsync()` +allows cancellation to be requested without trapping the host stopping thread inside a callback, +while callback completion still participates in a successful bounded drain. + +### Preserve obsolete public registry methods as compatibility shims + +Rejected because `RegisterTool()`, caller-configured filtering, and bulk initialization expose a +second incremental construction path that can violate the atomic-generation invariant. These were +implementation members of the unsupported MCP runtime assembly. Supported Core interface additions +instead use default methods to preserve compatibility. + +### Retain the two tokenless metadata-provider virtual slots as shims + +Rejected for this change. Calling a legacy override would allow arbitrary schema I/O to outlive the +bounded reload drain because the override has no cancellation input. Keeping the old slots but no +longer dispatching through them would avoid some loader failures while still changing subclass +behavior silently. The design instead makes the break explicit and requires any direct subclass to +adopt the token-bearing contract. + +### Notify for every published generation + +Rejected because configuration and metadata generations can change without changing MCP discovery. +Notifications are client cache invalidations, so emitting them for equivalent advertised metadata +causes unnecessary `tools/list` traffic. The implementation still publishes fresh tool instances +but compares canonical discovery metadata before notifying. + +### Canonicalize the JSON served to clients + +Rejected because recursively sorting schema properties would alter stored-procedure parameter wire +order. The registry stores separate order-preserving serving JSON and canonical comparison JSON so +semantic comparison does not change client-visible ordering. + +### Block disposal until an abandoned stdio write completes + +Rejected because a client can stop reading while leaving the pipe open, causing the write lock to +remain held indefinitely. Disposal marks the writer closed to new work and releases resources only +when the lock is immediately available; otherwise process teardown reclaims the process-owned +stdout handle. + +## Risks and Mitigations + +| Risk | Mitigation | +|---|---| +| Candidate built from stale configuration | Capture config, serialize registry rebuilds, and verify reference identity before publication. | +| Candidate uses stale DB metadata | Place MCP event after metadata refresh and inject metadata factory explicitly. | +| Invalid custom tool removes otherwise valid tools | Reject the candidate and retain the complete previous snapshot. | +| Old tool runs after swap | Execution revalidates current configuration and authorization. | +| Unrelated reload sends unnecessary notification | Compare deterministic advertised metadata before notifying. | +| Stdio notification corrupts JSON-RPC output | Use the shared locked `McpStdoutWriter`. | +| HTTP clients cache old list | Every explicit list request is current; HTTP push is tracked as follow-up. | +| Config commits while MCP refresh fails | Preserve valid registry, log degraded state, and rely on future transactional hot-reload work for global rollback. | +| Mutable SDK metadata is changed by a tool or caller | Deep-clone metadata into the candidate and return caller-owned clones from discovery. | +| Raw JSON order creates false discovery changes | Compare separately canonicalized JSON while preserving source property order in served JSON. | +| Startup publishes before DB metadata exists | Hosted service subscribes only; the shared startup helper publishes under the loader gate after metadata initialization. | +| Shutdown cancels only gate waiters but not active database work | Propagate the loader token through DAB-owned handlers, metadata providers, query executors, connections, commands, and token acquisition. | +| Synchronous `FillSchema()` ignores cancellation | Register `DbCommand.Cancel()` as best-effort provider cancellation and translate cancellation-triggered provider failures. | +| Cancellation callback blocks the host stopping thread | Request cancellation with `CancelAsync()` and include callback completion in the bounded drain. | +| Loader synchronization fields leak or are disposed while waiters unwind | Track gate waiters as admitted operations, release before signaling drain, then dispose the semaphore and token source from the shared shutdown-completion task. | +| Explicit query cancellation hides `RequestAborted` | Link distinct cancellable tokens and use the result for access-token acquisition, retry waits, connection opening, and command execution. | +| Non-cooperative extension prevents shutdown drain | Honor the host timeout and document that extension callbacks must observe the event token; stronger isolation is follow-up architecture. | +| OS watcher disposal blocks behind a callback | Disable and detach admission synchronously, then dispose the watcher independently on a background worker. | +| Stdio client stops reading stdout | Keep notification I/O off the reload path and make writer disposal reject new writes without waiting for a blocked write lock. | +| Notification queue grows while stdout is blocked | Store one pending invalidation bit and coalesce all intermediate changes. | +| New cancellation overloads break third-party Core implementations | Use default interface implementations that precheck cancellation and delegate to legacy members. | +| Replacing protected metadata-provider slots breaks a direct subclass | Accept the limited break because there is no documented custom-provider SPI, identify both slots explicitly, and provide a mechanical token-propagation migration path. | + +## Acceptance Criteria + +The implementation is complete when: + +1. HTTP and stdio startup use one shared registry initialization path. +2. Custom tools are no longer DI singletons tied to startup configuration. +3. A successful hot-reload atomically updates custom tool membership and metadata. +4. `tools/list` and `tools/call` never observe a partially rebuilt registry. +5. `tools/list` no longer combines a registry generation with an independently read configuration generation. +6. Duplicate or invalid candidate tools cannot partially update the registry. +7. A hot-reload rebuild failure leaves the previous registry usable and does not stop later hot-reload handlers. +8. A startup registry failure still fails startup. +9. Metadata initialization uses the exact captured configuration and refreshed metadata provider. +10. A stale rebuild cannot overwrite a newer registry generation. +11. An initialized stdio client receives `notifications/tools/list_changed` only after a successful semantic discovery change. +12. HTTP requests immediately observe the new snapshot without experimental session APIs. +13. Existing execution-time configuration and authorization validation remains intact. +14. Unit, concurrency, transport, and hot-reload tests cover the behaviors listed in this document. +15. Independently DI-registered `IMcpTool` extensions survive startup and every refresh. +16. Initial construction and every DAB-owned reload generation observe loader shutdown cancellation. +17. A successful loader drain completes before dependent hosted services and the root provider are disposed. +18. A non-cooperative extension cannot make the host ignore `HostOptions.ShutdownTimeout`. +19. Existing Core interface implementers and legacy `HotReloadEventArgs` construction remain source- and binary-compatible; the only accepted Core virtual-slot breaks are the two metadata-provider hooks documented above. +20. Blocked watcher or stdio resource cleanup cannot indefinitely delay coordinated host shutdown. + +## Follow-Up Work + +The following work remains intentionally separate: + +- Transactional application-wide hot-reload preparation, commit, and rollback. +- Dynamic MCP endpoint enablement and path changes. +- HTTP session tracking and `notifications/tools/list_changed` broadcast. +- Dynamic initialize instructions for future HTTP sessions, if required. +- Separate dependency ownership or process isolation for non-cooperative extension callbacks, if a + stronger post-timeout lifetime guarantee becomes a product requirement. +- Truly asynchronous schema-table discovery if database providers add an alternative to synchronous + `DbDataAdapter.FillSchema()`; command cancellation remains best-effort until then. +- A supported custom metadata-provider SPI, if required, with explicit registration, ownership, + lifecycle, compatibility, and cancellation contracts instead of protected implementation hooks. 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 1a4afaf8b6..4cd49bc48d 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..0e05e65d30 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; /// - /// Registers a tool in the registry + /// Replaces the complete registry with a snapshot built for . + /// The candidate is validated and materialized before it is atomically published. /// - /// Thrown when tool name is invalid or duplicate - public void RegisterTool(IMcpTool tool) + internal McpToolRegistryUpdateResult ReplaceAll(IEnumerable tools, RuntimeConfig config) { - Tool metadata = tool.GetToolMetadata(); - string toolName = metadata.Name?.Trim() ?? string.Empty; + 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); - // Reject empty or whitespace-only tool names + 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)); + } + + /// + /// Atomically publishes a previously built and validated candidate. + /// + internal McpToolRegistryUpdateResult PublishCandidate(McpToolRegistryCandidate candidate) + { + ArgumentNullException.ThrowIfNull(candidate); + + 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,204 @@ 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(); } - /// - /// Initializes and registers all MCP tools, enriching custom tools with DB metadata schemas. - /// Shared by both HTTP hosted-service and stdio startup paths. - /// - public static void InitializeAndRegisterTools( - IEnumerable tools, - McpToolRegistry registry, - IServiceProvider serviceProvider) + private static string CreateDiscoveryJson(ImmutableArray metadata) { - foreach (IMcpTool tool in tools) + 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)) + { + WriteCanonicalJson(writer, serializedMetadata.RootElement); + } + + 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, + string? propertyName = null, + bool isWithinJsonSchema = false) + { + 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, + property.Name, + isWithinJsonSchema || + property.NameEquals("inputSchema") || + property.NameEquals("outputSchema")); + } + + writer.WriteEndObject(); + break; + + case JsonValueKind.Array: + writer.WriteStartArray(); + if (!TryWriteOrderInsensitiveJsonSchemaStringArray( + writer, + element, + propertyName, + isWithinJsonSchema)) + { + foreach (JsonElement item in element.EnumerateArray()) + { + WriteCanonicalJson( + writer, + item, + propertyName: null, + isWithinJsonSchema); + } + } + + 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."); + } + } + + private static bool TryWriteOrderInsensitiveJsonSchemaStringArray( + Utf8JsonWriter writer, + JsonElement element, + string? propertyName, + bool isWithinJsonSchema) + { + // JSON Schema defines these arrays as sets, so their order does not change validation + // semantics. Do not sort every primitive array: values under keywords such as + // "default" and "examples" can be ordered JSON array instances whose order is part of + // the advertised metadata. + if (!isWithinJsonSchema || + propertyName is not ("required" or "type" or "enum")) + { + return false; + } + + List values = new(); + foreach (JsonElement item in element.EnumerateArray()) { - if (tool is DynamicCustomTool customTool) + if (item.ValueKind != JsonValueKind.String) { - customTool.InitializeMetadata(serviceProvider); + return false; } - registry.RegisterTool(tool); + values.Add(item.GetString()!); } + + values.Sort(StringComparer.Ordinal); + foreach (string value in values) + { + writer.WriteStringValue(value); + } + + return true; + } + + 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 205dc3d646..056646e0c8 100644 --- a/src/Core/Authorization/AuthorizationResolver.cs +++ b/src/Core/Authorization/AuthorizationResolver.cs @@ -76,6 +76,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 d449c396c0..c5ff45f49e 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..2895d6b89d 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 . See + /// docs/design/McpToolRegistryHotReload.md for the compatibility decision. + /// 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 . See + /// docs/design/McpToolRegistryHotReload.md for the compatibility decision. + /// 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 e66968de13..e580d46f5c 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 b3debd48f1..e9643fa278 100644 --- a/src/Service.Tests/Mcp/DynamicCustomToolTests.cs +++ b/src/Service.Tests/Mcp/DynamicCustomToolTests.cs @@ -566,19 +566,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 @@ -852,9 +858,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; } @@ -868,9 +879,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()); } @@ -887,7 +903,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), @@ -935,7 +952,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 1602bdc6bc..64959d58d5 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.Threading; @@ -24,17 +26,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); @@ -77,7 +109,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] @@ -132,7 +164,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"); } #region Helpers 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..a7195499d8 --- /dev/null +++ b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs @@ -0,0 +1,826 @@ +// 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); + } + + [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..dcd3a0de0d 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryTests.cs @@ -2,9 +2,9 @@ // 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.Threading; using System.Threading.Tasks; @@ -25,343 +25,443 @@ 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)); - - // Try to register a custom tool with a conflicting name - IMcpTool customTool = new MockMcpTool("read_records", ToolType.Custom); - - // Assert - Should throw - DataApiBuilderException exception = Assert.ThrowsException( - () => registry.RegisterTool(customTool) - ); + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, inputSchemaJson: SCHEMA_BA) }, + config); - Assert.IsTrue(exception.Message.Contains("read_records")); - Assert.IsTrue(exception.Message.Contains("built-in tool")); + Assert.IsTrue(result.DiscoveryChanged); } /// - /// 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. + /// 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. /// [TestMethod] - public void RegisterTool_WithLeadingTrailingWhitespace_DetectsDuplicate() + public void GetAdvertisedTools_PreservesInputSchemaPropertyOrder() { - // Arrange + const string SCHEMA = + "{\"type\":\"object\",\"properties\":{" + + "\"second\":{\"type\":\"string\"}," + + "\"first\":{\"type\":\"integer\"}}," + + "\"required\":[\"second\",\"first\"]}"; McpToolRegistry registry = new(); - IMcpTool tool1 = new MockMcpTool("my_tool", ToolType.BuiltIn); - IMcpTool tool2 = new MockMcpTool(" my_tool ", ToolType.Custom); - - // Act - registry.RegisterTool(tool1); - - // Assert - trimmed name should collide - Assert.ThrowsException( - () => registry.RegisterTool(tool2) - ); + 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); } /// - /// Parameterized test verifying GetEnabledTools returns only enabled tools. + /// A real input-schema change remains client-visible after canonicalization. /// - [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 ReplaceAll_WithChangedInputSchema_ReportsDiscoveryChange() { - // Arrange McpToolRegistry registry = new(); - for (int i = 0; i < enabledCount; i++) - { - registry.RegisterTool(new MockMcpTool($"enabled_{i}", ToolType.BuiltIn, isEnabledFunc: _ => true)); - } + 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); + } - for (int i = 0; i < disabledCount; i++) + /// + /// 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))}"); } /// @@ -461,12 +561,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 +587,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 +606,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 daf0c9e3b1..7394a9f53a 100644 --- a/src/Service.Tests/UnitTests/McpStdioHelperTests.cs +++ b/src/Service.Tests/UnitTests/McpStdioHelperTests.cs @@ -3,13 +3,20 @@ #nullable enable +using System.Collections.Generic; +using System.IO.Abstractions.TestingHelpers; using System.Threading; using System.Threading.Tasks; +using Azure.DataApiBuilder.Config; +using Azure.DataApiBuilder.Core.Configurations; +using Azure.DataApiBuilder.Core.Services.MetadataProviders; using Azure.DataApiBuilder.Mcp.Core; using Azure.DataApiBuilder.Service.Utilities; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Hosting; +using Microsoft.Extensions.Logging.Abstractions; using Microsoft.VisualStudio.TestTools.UnitTesting; +using Moq; namespace Azure.DataApiBuilder.Service.Tests.UnitTests { @@ -22,8 +29,30 @@ public void RunMcpStdioHost_DoesNotStartWebHost() ServiceCollection services = new(); TestApplicationLifetime lifetime = new(); TestMcpStdioServer stdioServer = new(); - - services.AddSingleton(); + List initializationOrder = new(); + 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); + Mock metadataProviderFactory = new(); + metadataProviderFactory + .Setup(factory => factory.InitializeAsync(It.IsAny())) + .Callback(() => initializationOrder.Add("metadata")) + .Returns(Task.CompletedTask); + TestMcpToolRegistryRefreshService refreshService = new(initializationOrder); + + services.AddSingleton(configLoader); + services.AddSingleton(runtimeConfigProvider); + services.AddSingleton(runtimeConfigValidator); + services.AddSingleton(metadataProviderFactory.Object); + services.AddSingleton(refreshService); services.AddSingleton(lifetime); services.AddSingleton(stdioServer); @@ -39,12 +68,39 @@ 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."); + metadataProviderFactory.Verify( + factory => factory.InitializeAsync(It.IsAny()), + Times.Once); + CollectionAssert.AreEqual( + new[] { "metadata", "registry" }, + initializationOrder, + "MCP stdio mode should initialize metadata before publishing the registry."); 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."); } + private sealed class TestMcpToolRegistryRefreshService : IMcpToolRegistryRefreshService + { + private readonly List _initializationOrder; + + public TestMcpToolRegistryRefreshService(List initializationOrder) + { + _initializationOrder = initializationOrder; + } + + public int EnsureInitializedCallCount { get; private set; } + + public void EnsureInitialized() + { + EnsureInitializedCallCount++; + _initializationOrder.Add("registry"); + } + } + private sealed class TestHost : IHost { public TestHost(System.IServiceProvider services) 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/McpStdioServerRunAsyncTests.cs b/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs index 3477e775a5..222b1f0ae4 100644 --- a/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs +++ b/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs @@ -11,6 +11,7 @@ using Azure.DataApiBuilder.Mcp.Core; using Microsoft.Extensions.DependencyInjection; using Microsoft.VisualStudio.TestTools.UnitTesting; +using Moq; namespace Azure.DataApiBuilder.Service.Tests.UnitTests { @@ -58,7 +59,42 @@ public async Task RunAsync_BlankLineThenShutdown_IgnoresBlankLineAndHandlesShutd "Expected shutdown response result.ok to be true."); } - private static (McpStdioServer server, StringWriter stdoutCapture) CreateServerWithCapturedOutput(TextReader inputReader) + [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); + } + + private static (McpStdioServer server, StringWriter stdoutCapture) CreateServerWithCapturedOutput( + TextReader inputReader, + IMcpStdioToolListChangedNotifier? notifier = null) { StringWriter stdoutCapture = new(); McpStdoutWriter stdoutWriter = new(stdoutCapture); @@ -66,6 +102,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/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 b640c79cd9..6be65e431a 100644 --- a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs +++ b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs @@ -1015,6 +1015,150 @@ private static Mock CreateHttpContextAccessorWithAuthentic #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 4ee403b98e..cb66e4c36c 100644 --- a/src/Service/Utilities/McpStdioHelper.cs +++ b/src/Service/Utilities/McpStdioHelper.cs @@ -4,9 +4,12 @@ using System; using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; +using System.Threading; +using Azure.DataApiBuilder.Config; using Microsoft.Extensions.Configuration; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Hosting; +using Microsoft.Extensions.Options; namespace Azure.DataApiBuilder.Service.Utilities { @@ -78,12 +81,17 @@ public static bool RunMcpStdioHost(IHost host) { try { - Mcp.Core.McpToolRegistry registry = - host.Services.GetRequiredService(); - IEnumerable tools = - host.Services.GetServices(); - - Mcp.Core.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(); @@ -96,6 +104,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!; + } + } +}