From b2e0d650b272e17817a997dd9e668abc31f75ed8 Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Thu, 30 Jul 2026 16:19:39 -0700 Subject: [PATCH 01/21] docs(mcp): design tool registry hot reload --- docs/design/McpToolRegistryHotReload.md | 702 ++++++++++++++++++++++++ 1 file changed, 702 insertions(+) create mode 100644 docs/design/McpToolRegistryHotReload.md diff --git a/docs/design/McpToolRegistryHotReload.md b/docs/design/McpToolRegistryHotReload.md new file mode 100644 index 0000000000..9d088ae06b --- /dev/null +++ b/docs/design/McpToolRegistryHotReload.md @@ -0,0 +1,702 @@ +# Design Document: MCP Tool Registry Hot-Reload + +## Status + +Proposed. + +This document describes the agreed design for refreshing Data API builder's MCP tool registry when runtime configuration is hot-reloaded. It is intended to guide implementation and review. + +Source links in this document point to the current implementation that will be changed; they are not examples of the proposed implementation. + +## Summary + +DAB currently builds its MCP tool registry once at startup. A hot-reloaded `RuntimeConfig` can change which custom tools exist and can change their names, descriptions, and input schemas, but the registry continues serving the startup tool instances and startup metadata. + +The proposed 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. + +## Motivation + +Custom MCP tools are generated from stored-procedure entities with `mcp.custom-tool` enabled. Today, those tools are constructed from the startup configuration and registered as DI singletons. The registry is 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 evaluates the current configuration during each `tools/list` request, but it is combined with a fixed startup registry. This avoids some stale built-in visibility, but does not solve stale custom tools or provide a single consistent registry generation. + +## Current Implementation + +### Registry construction + +[McpServiceCollectionExtensions.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpServiceCollectionExtensions.cs) currently: + +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. + +[McpToolRegistryInitializer.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs) resolves every `IMcpTool` and registers it once when the host starts. + +[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) stores tools in a mutable, case-insensitive `Dictionary`. It supports individual registration, lookup by name, and filtering enabled tools using a supplied `RuntimeConfig`. + +The dictionary is safe under current startup-only mutation, but it cannot 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 currently combine 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. + +## Non-Goals + +This work will 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 will no longer expose a dictionary that is incrementally mutated during normal operation. Instead, it will hold one current immutable snapshot reference. + +### 2. Registry generations are immutable snapshots + +The registry snapshot will conceptually contain: + +```csharp +internal sealed record McpToolRegistrySnapshot( + long Version, + ImmutableDictionary Tools, + ImmutableArray AdvertisedTools); +``` + +`Tools` contains: + +- Every built-in tool, including built-ins currently disabled by DML tool configuration. +- Every custom tool enabled in the configuration used to build the snapshot. + +`AdvertisedTools` contains precomputed metadata for tools whose `IsEnabled(config)` result was true for that same configuration generation. It is sorted deterministically by tool name. + +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 stored in the snapshot are treated as immutable after construction. A list handler may create a new list container, but it must not mutate snapshot metadata. + +### 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` reads `AdvertisedTools` from one snapshot. +- `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. Built-in and custom 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. + +Custom tools are removed from DI registration. They are configuration-generation objects and are recreated for every registry candidate. + +This avoids treating the immutable DI service collection as a dynamic registry. + +### 5. Refresh orchestration is separate from state storage + +A singleton `McpToolRegistryRefreshService` will coordinate initialization and hot-reload. It will also implement `IHostedService` for normal HTTP-host startup. + +Its responsibilities are: + +1. Capture the current `RuntimeConfig` generation. +2. Obtain the DI-owned built-in tools. +3. Create fresh 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. + +### 6. Custom tool creation is strict + +`CustomMcpToolFactory` currently catches broad exceptions and skips individual entities. That 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 + +`DynamicCustomTool.InitializeMetadata(IServiceProvider)` is a service-locator pattern and allows metadata initialization to retrieve a `RuntimeConfig` different from the generation being built. + +The initialization path will instead receive explicit dependencies, conceptually: + +```csharp +void InitializeMetadata( + RuntimeConfig config, + IMetadataProviderFactory metadataProviderFactory); +``` + +`McpMetadataHelper` may gain an overload that accepts `IMetadataProviderFactory` directly. Existing execution call sites may 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. + +#### HTTP mode + +The host starts the same singleton through `IHostedService.StartAsync()`. + +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 hosted services do not run. `McpStdioHelper` explicitly resolves the same refresh-service singleton and invokes its idempotent initialization method before starting the stdio loop. + +This replaces the current 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`. + +The refresh service serializes its own rebuilds and uses a stale-generation guard: + +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. + +This guard prevents an older, slower registry rebuild from overwriting a newer registry generation. It does not claim to make the full DAB hot-reload pipeline transactional; that remains separate work. + +### 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 + +The stdio initialize response already advertises `tools.listChanged = true`, so the server must implement the corresponding notification. + +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 client 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. +- 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 will own initialization state and frame writing. `McpStdioServer` marks it initialized when handling `notifications/initialized`. The refresh service depends on zero or more tool-list notifiers; HTTP mode has no notifier registered in this iteration. + +### 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 must not be advertised 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. The comparison covers the complete serialized tool metadata, including name, description, input schema, and any future advertised fields. + +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); + +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 current public `RegisterTool` method is not used by the new production path. Because it is public, implementation should preserve it unless API review explicitly approves removal. If retained, it must use copy-on-write under the writer gate and must never mutate a published dictionary in place. + +## Detailed Flows + +### Initial HTTP startup + +```mermaid +sequenceDiagram + participant Host + participant Refresh as McpToolRegistryRefreshService + participant Factory as CustomMcpToolFactory + participant Metadata as IMetadataProviderFactory + participant Registry as McpToolRegistry + + Host->>Refresh: StartAsync() + 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 Registry as McpToolRegistry + participant Server as McpStdioServer + + Helper->>Refresh: EnsureInitialized() + 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 stale-generation guard prevents publication when the active `RuntimeConfig` changed during candidate construction. The latest callback eventually publishes the latest generation. + +The design intentionally does not lock the entire DAB hot-reload pipeline. Global serialization and transactional rollback are 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. +- `McpToolRegistryRefreshService`: singleton. +- `IHostedService`: resolves the same refresh-service singleton. +- Custom tools: not registered in DI. +- Stdio tool-list notifier: singleton, registered only in stdio mode. + +The refresh service receives `IEnumerable` containing built-ins only. Reflection-based built-in discovery remains unchanged except for continuing to exclude `DynamicCustomTool`. + +## Anticipated Source Changes + +The exact file split may change during implementation, but the expected touchpoints are: + +### 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. + +### MCP project + +- [McpToolRegistry.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs): immutable snapshots, bulk replacement, and atomic reads/publication. +- [McpToolRegistryInitializer.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs): replace with or evolve into the refresh service. +- [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 + +- [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. + +### Refresh-service unit tests + +1. Initial construction uses DI-owned built-ins and newly created custom 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 callbacks for an already successfully applied configuration do not publish duplicate generations unnecessarily. +11. Notifications occur only after a successful noninitial semantic discovery change. + +### 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 does not notify before `notifications/initialized`. +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. + +### 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. + +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. + +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. + +## 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. | +| Public `RegisterTool` API conflicts with snapshot design | Preserve via safe copy-on-write unless API review approves removal. | + +## 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. + +## 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. From d572ff85ff441bfe94b1a89fd73bd07d2563d0bf Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Thu, 30 Jul 2026 16:37:15 -0700 Subject: [PATCH 02/21] refactor(mcp): add atomic tool registry snapshots --- .../Core/CustomMcpToolFactory.cs | 6 +- .../Core/DynamicCustomTool.cs | 48 +++-- .../Core/McpToolRegistry.cs | 201 +++++++++++++++--- .../Utils/McpMetadataHelper.cs | 42 +++- src/Service.Tests/Mcp/McpToolRegistryTests.cs | 195 ++++++++++++++++- 5 files changed, 440 insertions(+), 52 deletions(-) diff --git a/src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs b/src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs index f688eeb80a..67303f6dbf 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs @@ -50,8 +50,12 @@ public static IEnumerable CreateCustomTools(RuntimeConfig config, ILog { logger?.LogError( ex, - "Failed to create custom tool for entity '{EntityName}'. Skipping.", + "Failed to create custom MCP tool for entity '{EntityName}'.", entityName); + + 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..0a9917067b 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 { @@ -82,7 +78,29 @@ public DynamicCustomTool(string entityName, Entity entity) public void InitializeMetadata(IServiceProvider serviceProvider) { ArgumentNullException.ThrowIfNull(serviceProvider); - _cachedInputSchema = BuildInputSchemaFromDbMetadata(serviceProvider); + + RuntimeConfigProvider? configProvider = serviceProvider.GetService(); + IMetadataProviderFactory? metadataProviderFactory = serviceProvider.GetService(); + if (configProvider is null || metadataProviderFactory is null) + { + _cachedInputSchema = null; + return; + } + + InitializeMetadata(configProvider.GetConfig(), metadataProviderFactory); + } + + /// + /// Initializes the input schema using an explicit configuration and metadata-provider + /// generation. Falls back to config-based metadata when database metadata is unavailable. + /// + public void InitializeMetadata( + RuntimeConfig config, + IMetadataProviderFactory metadataProviderFactory) + { + ArgumentNullException.ThrowIfNull(config); + ArgumentNullException.ThrowIfNull(metadataProviderFactory); + _cachedInputSchema = BuildInputSchemaFromDbMetadata(config, metadataProviderFactory); } /// @@ -322,20 +340,14 @@ 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) { - 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 _, diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs index f8275c61d4..f7ebce796f 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs @@ -1,7 +1,10 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +using System.Collections.Immutable; using System.Net; +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,58 +18,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 + /// Registers a tool in the registry using copy-on-write publication. + /// This compatibility API is retained for callers that incrementally construct a registry; + /// production initialization and hot-reload use . /// /// Thrown when tool name is invalid or duplicate public void RegisterTool(IMcpTool tool) { + ArgumentNullException.ThrowIfNull(tool); + Tool metadata = tool.GetToolMetadata(); - string toolName = metadata.Name?.Trim() ?? string.Empty; + string toolName = ValidateToolName(metadata); - // Reject empty or whitespace-only tool names - if (string.IsNullOrWhiteSpace(toolName)) + lock (_writerLock) { - throw new DataApiBuilderException( - message: "MCP tool name cannot be null, empty, or whitespace.", - statusCode: HttpStatusCode.ServiceUnavailable, - subStatusCode: DataApiBuilderException.SubStatusCodes.ErrorInInitialization); + McpToolRegistrySnapshot current = _snapshot; + + if (current.Tools.TryGetValue(toolName, out IMcpTool? existingTool)) + { + if (ReferenceEquals(existingTool, tool)) + { + return; + } + + throw CreateDuplicateToolException(toolName, existingTool, tool); + } + + ImmutableDictionary tools = current.Tools.Add(toolName, tool); + ImmutableArray advertisedTools = SortMetadata(current.AdvertisedTools.Add(metadata)); + string fingerprint = CreateDiscoveryFingerprint(advertisedTools); + + Interlocked.Exchange( + ref _snapshot, + new McpToolRegistrySnapshot( + Version: current.Version + 1, + Tools: tools, + AdvertisedTools: advertisedTools, + DiscoveryFingerprint: fingerprint)); } + } - // Check for duplicate tool names (case-insensitive) - if (_tools.TryGetValue(toolName, out IMcpTool? existingTool)) + /// + /// Replaces the complete registry with a snapshot built for . + /// The candidate is validated and materialized before it is atomically published. + /// + public McpToolRegistryUpdateResult ReplaceAll(IEnumerable tools, RuntimeConfig config) + { + ArgumentNullException.ThrowIfNull(tools); + ArgumentNullException.ThrowIfNull(config); + + lock (_writerLock) { - // 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)) + ImmutableDictionary.Builder toolBuilder = + ImmutableDictionary.CreateBuilder(StringComparer.OrdinalIgnoreCase); + List advertisedMetadata = new(); + + foreach (IMcpTool tool in tools) { - return; + ArgumentNullException.ThrowIfNull(tool); + + Tool metadata = tool.GetToolMetadata(); + string toolName = ValidateToolName(metadata); + + if (toolBuilder.TryGetValue(toolName, out IMcpTool? existingTool)) + { + if (ReferenceEquals(existingTool, tool)) + { + continue; + } + + throw CreateDuplicateToolException(toolName, existingTool, tool); + } + + toolBuilder.Add(toolName, tool); + if (tool.IsEnabled(config)) + { + advertisedMetadata.Add(metadata); + } } - string existingToolType = existingTool.ToolType == ToolType.BuiltIn ? "built-in" : "custom"; - string newToolType = tool.ToolType == ToolType.BuiltIn ? "built-in" : "custom"; + ImmutableArray advertisedTools = SortMetadata(advertisedMetadata); + string fingerprint = CreateDiscoveryFingerprint(advertisedTools); + McpToolRegistrySnapshot current = _snapshot; + McpToolRegistrySnapshot replacement = new( + Version: current.Version + 1, + Tools: toolBuilder.ToImmutable(), + AdvertisedTools: advertisedTools, + DiscoveryFingerprint: fingerprint); - 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.", - statusCode: HttpStatusCode.ServiceUnavailable, - subStatusCode: DataApiBuilderException.SubStatusCodes.ErrorInInitialization); + Interlocked.Exchange(ref _snapshot, replacement); + + return new McpToolRegistryUpdateResult( + Version: replacement.Version, + DiscoveryChanged: !string.Equals( + current.DiscoveryFingerprint, + replacement.DiscoveryFingerprint, + StringComparison.Ordinal), + RegisteredToolCount: replacement.Tools.Count, + AdvertisedToolCount: replacement.AdvertisedTools.Length); } + } - _tools[toolName] = tool; + /// + /// Gets the metadata snapshot advertised by tools/list. + /// + public IReadOnlyList GetAdvertisedTools() + { + return Volatile.Read(ref _snapshot).AdvertisedTools; } /// /// Gets metadata for all registered tools that are enabled in the given runtime configuration. + /// Retained for compatibility; MCP handlers should use so + /// lookup and discovery come from the same registry generation. /// public IEnumerable GetEnabledTools(RuntimeConfig config) { - return _tools.Values + ArgumentNullException.ThrowIfNull(config); + McpToolRegistrySnapshot snapshot = Volatile.Read(ref _snapshot); + return snapshot.Tools.Values .Where(t => t.IsEnabled(config)) .Select(t => t.GetToolMetadata()); } @@ -76,7 +155,7 @@ public IEnumerable GetEnabledTools(RuntimeConfig config) /// public bool TryGetTool(string toolName, out IMcpTool? tool) { - return _tools.TryGetValue(toolName, out tool); + return Volatile.Read(ref _snapshot).Tools.TryGetValue(toolName, out tool); } /// @@ -98,5 +177,71 @@ public static void InitializeAndRegisterTools( registry.RegisterTool(tool); } } + + private static string ValidateToolName(Tool metadata) + { + string toolName = metadata.Name?.Trim() ?? string.Empty; + if (string.IsNullOrWhiteSpace(toolName)) + { + throw new DataApiBuilderException( + message: "MCP tool name cannot be null, empty, or whitespace.", + statusCode: HttpStatusCode.ServiceUnavailable, + subStatusCode: DataApiBuilderException.SubStatusCodes.ErrorInInitialization); + } + + return toolName; + } + + private static DataApiBuilderException CreateDuplicateToolException( + string toolName, + IMcpTool existingTool, + IMcpTool newTool) + { + 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); + } + + private static ImmutableArray SortMetadata(IEnumerable metadata) + { + return metadata + .OrderBy(tool => tool.Name, StringComparer.OrdinalIgnoreCase) + .ThenBy(tool => tool.Name, StringComparer.Ordinal) + .ToImmutableArray(); + } + + private static string CreateDiscoveryFingerprint(ImmutableArray metadata) + { + return JsonSerializer.Serialize(metadata.ToArray(), _discoveryJsonOptions); + } + + private sealed record McpToolRegistrySnapshot( + long Version, + ImmutableDictionary Tools, + ImmutableArray AdvertisedTools, + string DiscoveryFingerprint) + { + public static McpToolRegistrySnapshot Empty { get; } = new( + Version: 0, + Tools: ImmutableDictionary.Create(StringComparer.OrdinalIgnoreCase), + AdvertisedTools: ImmutableArray.Empty, + DiscoveryFingerprint: "[]"); + } } + + /// + /// Describes the result of atomically replacing an MCP registry snapshot. + /// + public readonly record struct McpToolRegistryUpdateResult( + long Version, + bool DiscoveryChanged, + int RegisteredToolCount, + int AdvertisedToolCount); } diff --git a/src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs b/src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs index 2d79649bbb..2ca7e116e9 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, @@ -65,14 +67,48 @@ public static bool TryResolveMetadata( } // 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; + // Resolve datasource name for the entity. try { diff --git a/src/Service.Tests/Mcp/McpToolRegistryTests.cs b/src/Service.Tests/Mcp/McpToolRegistryTests.cs index d8c5dc0b59..f45e6e218a 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryTests.cs @@ -2,6 +2,7 @@ // Licensed under the MIT License. using System; +using System.Collections.Concurrent; using System.Collections.Generic; using System.Linq; using System.Net; @@ -364,6 +365,190 @@ public void GetEnabledTools_MixedBuiltInAndCustomTools() Assert.IsFalse(enabledTools.Any(t => t.Name == "delete_record")); } + /// + /// Replacing the registry publishes a complete, deterministically ordered snapshot and + /// removes tools that belonged only to the previous generation. + /// + [TestMethod] + public void ReplaceAll_PublishesCompleteOrderedSnapshot() + { + McpToolRegistry registry = new(); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll( + new[] { new MockMcpTool("old_tool", ToolType.Custom) }, + config); + + 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); + } + + /// + /// A candidate containing a duplicate name is rejected before publication, leaving the + /// complete previous snapshot active. + /// + [TestMethod] + public void ReplaceAll_WithDuplicateName_PreservesPreviousSnapshot() + { + McpToolRegistry registry = new(); + RuntimeConfig config = CreateRuntimeConfig(); + IMcpTool previousTool = new MockMcpTool("previous_tool", ToolType.BuiltIn); + registry.ReplaceAll(new[] { previousTool }, config); + + 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()); + } + + /// + /// Replacing tool instances with semantically identical discovery metadata advances the + /// registry generation without reporting a client-visible discovery change. + /// + [TestMethod] + public void ReplaceAll_WithEquivalentMetadata_DoesNotReportDiscoveryChange() + { + McpToolRegistry registry = new(); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, description: "Same description") }, + config); + + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, description: "Same description") }, + config); + + Assert.AreEqual(2, result.Version); + Assert.IsFalse(result.DiscoveryChanged); + } + + /// + /// 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); + + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, description: "New description") }, + config); + + Assert.IsTrue(result.DiscoveryChanged); + Assert.AreEqual("New description", registry.GetAdvertisedTools().Single().Description); + } + + /// + /// 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 ReplaceAll_CapturesVisibilityFromCandidateConfig() + { + McpToolRegistry registry = new(); + IMcpTool configAwareTool = new MockMcpTool( + "create_record", + ToolType.BuiltIn, + isEnabledFunc: config => config.McpDmlTools?.CreateRecord == true); + + RuntimeConfig disabledConfig = CreateRuntimeConfig(new DmlToolsConfig(createRecord: false)); + registry.ReplaceAll(new[] { configAwareTool }, disabledConfig); + + Assert.AreEqual(0, registry.GetAdvertisedTools().Count); + Assert.IsTrue(registry.TryGetTool("create_record", out _)); + + RuntimeConfig enabledConfig = CreateRuntimeConfig(new DmlToolsConfig(createRecord: true)); + registry.ReplaceAll(new[] { configAwareTool }, enabledConfig); + + Assert.AreEqual(1, registry.GetAdvertisedTools().Count); + } + + /// + /// Concurrent readers see only a complete old or complete new advertised snapshot while + /// registry generations are repeatedly replaced. + /// + [TestMethod] + public void ReplaceAll_WithConcurrentReaders_NeverExposesPartialSnapshot() + { + McpToolRegistry registry = new(); + 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); + + 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))}"); + } + /// /// Validates IsEnabled for each real built-in tool matches the DmlToolsConfig flag value. /// @@ -461,12 +646,18 @@ private class MockMcpTool : IMcpTool { private readonly string _toolName; private readonly Func? _isEnabledFunc; + private readonly string _description; - public MockMcpTool(string toolName, ToolType toolType, Func? isEnabledFunc = null) + public MockMcpTool( + string toolName, + ToolType toolType, + Func? isEnabledFunc = null, + string? description = null) { _toolName = toolName; ToolType = toolType; _isEnabledFunc = isEnabledFunc; + _description = description ?? $"Mock {toolType} tool"; } public ToolType ToolType { get; } @@ -483,7 +674,7 @@ public Tool GetToolMetadata() return new Tool { Name = _toolName, - Description = $"Mock {ToolType} tool", + Description = _description, InputSchema = doc.RootElement.Clone() }; } From 8a3ce27162f16d8c53c8849c1159522b9644451f Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Thu, 30 Jul 2026 16:58:21 -0700 Subject: [PATCH 03/21] feat(mcp): refresh tool registry on config reload --- .../Core/DynamicCustomTool.cs | 5 +- .../Core/McpServerConfiguration.cs | 8 +- .../Core/McpServiceCollectionExtensions.cs | 24 +- .../Core/McpStdioServer.cs | 9 +- .../Core/McpToolRegistry.cs | 104 ++--- .../Core/McpToolRegistryInitializer.cs | 32 +- .../Core/McpToolRegistryRefreshService.cs | 189 +++++++++ src/Config/DabConfigEvents.cs | 1 + src/Config/HotReloadEventHandler.cs | 1 + src/Config/RuntimeConfigLoader.cs | 4 + .../Mcp/McpToolRegistryRefreshServiceTests.cs | 360 ++++++++++++++++++ src/Service.Tests/Mcp/McpToolRegistryTests.cs | 2 +- .../UnitTests/McpStdioHelperTests.cs | 15 +- src/Service/Utilities/McpStdioHelper.cs | 9 +- 14 files changed, 671 insertions(+), 92 deletions(-) create mode 100644 src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs create mode 100644 src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs diff --git a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs index 0a9917067b..d991e00277 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs @@ -87,20 +87,21 @@ public void InitializeMetadata(IServiceProvider serviceProvider) return; } - InitializeMetadata(configProvider.GetConfig(), metadataProviderFactory); + _ = InitializeMetadata(configProvider.GetConfig(), metadataProviderFactory); } /// /// Initializes the input schema using an explicit configuration and metadata-provider /// generation. Falls back to config-based metadata when database metadata is unavailable. /// - public void InitializeMetadata( + public bool InitializeMetadata( RuntimeConfig config, IMetadataProviderFactory metadataProviderFactory) { ArgumentNullException.ThrowIfNull(config); ArgumentNullException.ThrowIfNull(metadataProviderFactory); _cachedInputSchema = BuildInputSchemaFromDbMetadata(config, metadataProviderFactory); + return _cachedInputSchema.HasValue; } /// diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs index 20040588fe..8b5c943863 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) => diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpServiceCollectionExtensions.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpServiceCollectionExtensions.cs index c88cae148d..5bc5bcd278 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,16 @@ 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 and register built-in MCP tools. Custom tools are configuration- + // generation objects and are created by McpToolRegistryRefreshService. 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 +69,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..08e5d7471a 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs @@ -287,16 +287,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/McpToolRegistry.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs index f7ebce796f..7f2bc50117 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs @@ -72,48 +72,70 @@ public void RegisterTool(IMcpTool tool) /// The candidate is validated and materialized before it is atomically published. /// public McpToolRegistryUpdateResult ReplaceAll(IEnumerable tools, RuntimeConfig config) + { + return PublishCandidate(CreateCandidate(tools, config)); + } + + /// + /// Builds and validates a complete replacement without publishing it. + /// + internal static McpToolRegistryCandidate CreateCandidate( + IEnumerable tools, + RuntimeConfig config) { ArgumentNullException.ThrowIfNull(tools); ArgumentNullException.ThrowIfNull(config); - lock (_writerLock) - { - ImmutableDictionary.Builder toolBuilder = - ImmutableDictionary.CreateBuilder(StringComparer.OrdinalIgnoreCase); - List advertisedMetadata = new(); + ImmutableDictionary.Builder toolBuilder = + ImmutableDictionary.CreateBuilder(StringComparer.OrdinalIgnoreCase); + List advertisedMetadata = new(); - foreach (IMcpTool tool in tools) - { - ArgumentNullException.ThrowIfNull(tool); + foreach (IMcpTool tool in tools) + { + ArgumentNullException.ThrowIfNull(tool); - Tool metadata = tool.GetToolMetadata(); - string toolName = ValidateToolName(metadata); + Tool metadata = tool.GetToolMetadata(); + string toolName = ValidateToolName(metadata); - if (toolBuilder.TryGetValue(toolName, out IMcpTool? existingTool)) + if (toolBuilder.TryGetValue(toolName, out IMcpTool? existingTool)) + { + if (ReferenceEquals(existingTool, tool)) { - if (ReferenceEquals(existingTool, tool)) - { - continue; - } - - throw CreateDuplicateToolException(toolName, existingTool, tool); + continue; } - toolBuilder.Add(toolName, tool); - if (tool.IsEnabled(config)) - { - advertisedMetadata.Add(metadata); - } + throw CreateDuplicateToolException(toolName, existingTool, tool); } - ImmutableArray advertisedTools = SortMetadata(advertisedMetadata); - string fingerprint = CreateDiscoveryFingerprint(advertisedTools); + toolBuilder.Add(toolName, tool); + if (tool.IsEnabled(config)) + { + advertisedMetadata.Add(metadata); + } + } + + ImmutableArray advertisedTools = SortMetadata(advertisedMetadata); + return new McpToolRegistryCandidate( + Tools: toolBuilder.ToImmutable(), + AdvertisedTools: advertisedTools, + DiscoveryFingerprint: CreateDiscoveryFingerprint(advertisedTools)); + } + + /// + /// 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: toolBuilder.ToImmutable(), - AdvertisedTools: advertisedTools, - DiscoveryFingerprint: fingerprint); + Tools: candidate.Tools, + AdvertisedTools: candidate.AdvertisedTools, + DiscoveryFingerprint: candidate.DiscoveryFingerprint); Interlocked.Exchange(ref _snapshot, replacement); @@ -158,26 +180,6 @@ public bool TryGetTool(string toolName, out IMcpTool? tool) return Volatile.Read(ref _snapshot).Tools.TryGetValue(toolName, out tool); } - /// - /// 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) - { - foreach (IMcpTool tool in tools) - { - if (tool is DynamicCustomTool customTool) - { - customTool.InitializeMetadata(serviceProvider); - } - - registry.RegisterTool(tool); - } - } - private static string ValidateToolName(Tool metadata) { string toolName = metadata.Name?.Trim() ?? string.Empty; @@ -244,4 +246,12 @@ public readonly record struct McpToolRegistryUpdateResult( bool DiscoveryChanged, int RegisteredToolCount, int AdvertisedToolCount); + + /// + /// A fully materialized and validated registry generation awaiting publication. + /// + internal sealed record McpToolRegistryCandidate( + ImmutableDictionary Tools, + ImmutableArray AdvertisedTools, + string DiscoveryFingerprint); } diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs index a7c323a967..b37e824929 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs @@ -8,14 +8,18 @@ namespace Azure.DataApiBuilder.Mcp.Core { /// - /// Hosted service to initialize the MCP tool registry + /// Compatibility hosted service for callers that previously constructed the registry initializer + /// directly. DAB startup uses . /// + [Obsolete($"Use {nameof(McpToolRegistryRefreshService)} instead.")] public class McpToolRegistryInitializer : IHostedService { private readonly IServiceProvider _serviceProvider; private readonly McpToolRegistry _toolRegistry; - public McpToolRegistryInitializer(IServiceProvider serviceProvider, McpToolRegistry toolRegistry) + public McpToolRegistryInitializer( + IServiceProvider serviceProvider, + McpToolRegistry toolRegistry) { _serviceProvider = serviceProvider; _toolRegistry = toolRegistry; @@ -23,8 +27,28 @@ public McpToolRegistryInitializer(IServiceProvider serviceProvider, McpToolRegis public Task StartAsync(CancellationToken cancellationToken) { - IEnumerable tools = _serviceProvider.GetServices(); - McpToolRegistry.InitializeAndRegisterTools(tools, _toolRegistry, _serviceProvider); + cancellationToken.ThrowIfCancellationRequested(); + + IMcpToolRegistryRefreshService? refreshService = + _serviceProvider.GetService(); + if (refreshService is not null) + { + refreshService.EnsureInitialized(); + return Task.CompletedTask; + } + + // Preserve the legacy behavior for manually assembled service providers that have not + // registered the refresh service. + foreach (IMcpTool tool in _serviceProvider.GetServices()) + { + if (tool is DynamicCustomTool customTool) + { + customTool.InitializeMetadata(_serviceProvider); + } + + _toolRegistry.RegisterTool(tool); + } + 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..5ddff8aee6 --- /dev/null +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs @@ -0,0 +1,189 @@ +// 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(); + } + + /// + /// 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 _builtInTools; + private readonly McpToolRegistry _toolRegistry; + private readonly IMetadataProviderFactory _metadataProviderFactory; + private readonly IReadOnlyList _notifiers; + private readonly ILogger _logger; + private readonly object _refreshLock = new(); + private RuntimeConfig? _lastAppliedConfig; + + public McpToolRegistryRefreshService( + RuntimeConfigProvider runtimeConfigProvider, + IEnumerable tools, + McpToolRegistry toolRegistry, + IMetadataProviderFactory metadataProviderFactory, + IEnumerable notifiers, + ILogger logger, + HotReloadEventHandler? hotReloadEventHandler = null) + { + _runtimeConfigProvider = runtimeConfigProvider; + _builtInTools = tools + .Where(tool => tool.ToolType == ToolType.BuiltIn) + .ToArray(); + _toolRegistry = toolRegistry; + _metadataProviderFactory = metadataProviderFactory; + _notifiers = notifiers.ToArray(); + _logger = logger; + + hotReloadEventHandler?.Subscribe( + MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED, + OnConfigChanged); + } + + /// + public void EnsureInitialized() + { + RefreshRegistry(); + } + + /// + public Task StartAsync(CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + EnsureInitialized(); + return Task.CompletedTask; + } + + /// + public Task StopAsync(CancellationToken cancellationToken) + { + return Task.CompletedTask; + } + + private void OnConfigChanged(object? sender, HotReloadEventArgs args) + { + try + { + RefreshRegistry(); + } + catch (Exception ex) + { + _logger.LogError( + ex, + "Failed to refresh the MCP tool registry after a runtime configuration change. " + + "The previous registry snapshot remains active."); + } + } + + private void RefreshRegistry() + { + lock (_refreshLock) + { + RuntimeConfig config = _runtimeConfigProvider.GetConfig(); + if (ReferenceEquals(config, _lastAppliedConfig)) + { + return; + } + + List customTools = CustomMcpToolFactory + .CreateCustomTools(config, _logger) + .ToList(); + + foreach (DynamicCustomTool customTool in customTools.Cast()) + { + bool initializedFromDatabase = customTool.InitializeMetadata( + config, + _metadataProviderFactory); + if (!initializedFromDatabase) + { + _logger.LogWarning( + "Database metadata was unavailable for custom MCP tool '{ToolName}' " + + "on entity '{EntityName}'. Using configuration-derived input schema.", + customTool.GetToolMetadata().Name, + customTool.EntityName); + } + } + + McpToolRegistryCandidate candidate = McpToolRegistry.CreateCandidate( + _builtInTools.Concat(customTools), + config); + + 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; + } + + bool isInitialGeneration = _lastAppliedConfig is null; + McpToolRegistryUpdateResult result = _toolRegistry.PublishCandidate(candidate); + _lastAppliedConfig = config; + + _logger.LogInformation( + "Published MCP tool registry version {Version} with {RegisteredToolCount} " + + "registered tools and {AdvertisedToolCount} advertised tools.", + result.Version, + result.RegisteredToolCount, + result.AdvertisedToolCount); + + if (!isInitialGeneration && result.DiscoveryChanged) + { + NotifyToolsListChanged(); + } + } + } + + 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 + { + /// + /// Notifies a connected, initialized client that it should refresh tools/list. + /// + void NotifyToolsListChanged(); + } +} 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/HotReloadEventHandler.cs b/src/Config/HotReloadEventHandler.cs index 666c3c227b..28464650ad 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 }, diff --git a/src/Config/RuntimeConfigLoader.cs b/src/Config/RuntimeConfigLoader.cs index b40d0b084f..2e05e208b4 100644 --- a/src/Config/RuntimeConfigLoader.cs +++ b/src/Config/RuntimeConfigLoader.cs @@ -105,6 +105,10 @@ protected void SignalConfigChanged(string message = "") // this function is called. OnConfigChangedEvent(new HotReloadEventArgs(AUTHZ_RESOLVER_ON_CONFIG_CHANGED, message)); + // Custom MCP tool schemas depend on refreshed database metadata. Publish the new + // registry only after query, mutation, and authorization dependencies are ready. + OnConfigChangedEvent(new HotReloadEventArgs(MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED, message)); + // 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)); diff --git a/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs new file mode 100644 index 0000000000..bf56702712 --- /dev/null +++ b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs @@ -0,0 +1,360 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +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.Logging.Abstractions; +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 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_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 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 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(new Dictionary()); + Mock metadataProviderFactory = new(); + metadataProviderFactory + .Setup(factory => factory.GetMetadataProvider(It.IsAny())) + .Returns(sqlMetadataProvider.Object); + + Mock notifier = new(); + McpToolRegistry registry = new(); + HotReloadEventHandler hotReloadEventHandler = new(); + McpToolRegistryRefreshService service = new( + configProvider.Object, + builtInTools, + registry, + metadataProviderFactory.Object, + new[] { notifier.Object }, + NullLogger.Instance, + hotReloadEventHandler); + + return new TestContext( + service, + registry, + notifier, + hotReloadEventHandler); + } + + 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 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 sealed record TestContext( + McpToolRegistryRefreshService Service, + McpToolRegistry Registry, + Mock Notifier, + HotReloadEventHandler HotReloadEventHandler); + + 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 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 f45e6e218a..68d25a4146 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryTests.cs @@ -144,7 +144,7 @@ public void RegisterTool_WithDifferentCasing_ThrowsException() /// /// 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. + /// This preserves the compatibility registration API's idempotent behavior. /// [TestMethod] public void RegisterTool_SameInstanceTwice_IsIdempotent() diff --git a/src/Service.Tests/UnitTests/McpStdioHelperTests.cs b/src/Service.Tests/UnitTests/McpStdioHelperTests.cs index daf0c9e3b1..9d7fe855ac 100644 --- a/src/Service.Tests/UnitTests/McpStdioHelperTests.cs +++ b/src/Service.Tests/UnitTests/McpStdioHelperTests.cs @@ -22,8 +22,9 @@ public void RunMcpStdioHost_DoesNotStartWebHost() ServiceCollection services = new(); TestApplicationLifetime lifetime = new(); TestMcpStdioServer stdioServer = new(); + TestMcpToolRegistryRefreshService refreshService = new(); - services.AddSingleton(); + services.AddSingleton(refreshService); services.AddSingleton(lifetime); services.AddSingleton(stdioServer); @@ -39,12 +40,24 @@ public void RunMcpStdioHost_DoesNotStartWebHost() "MCP stdio mode should not stop a host that was never started."); Assert.AreEqual(1, stdioServer.RunAsyncCallCount, "MCP stdio mode should still run the stdio JSON-RPC loop."); + Assert.AreEqual(1, refreshService.EnsureInitializedCallCount, + "MCP stdio mode should initialize the shared tool registry before running the loop."); Assert.AreEqual(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 + { + public int EnsureInitializedCallCount { get; private set; } + + public void EnsureInitialized() + { + EnsureInitializedCallCount++; + } + } + private sealed class TestHost : IHost { public TestHost(System.IServiceProvider services) diff --git a/src/Service/Utilities/McpStdioHelper.cs b/src/Service/Utilities/McpStdioHelper.cs index 4ee403b98e..799f344e80 100644 --- a/src/Service/Utilities/McpStdioHelper.cs +++ b/src/Service/Utilities/McpStdioHelper.cs @@ -78,12 +78,9 @@ 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); + Mcp.Core.IMcpToolRegistryRefreshService refreshService = + host.Services.GetRequiredService(); + refreshService.EnsureInitialized(); IHostApplicationLifetime lifetime = host.Services.GetRequiredService(); From d6f49affe75d8d55b958097f699f944f2779b2a2 Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Thu, 30 Jul 2026 17:14:20 -0700 Subject: [PATCH 04/21] feat(mcp): notify stdio clients when tools change --- .../Core/McpStdioServer.cs | 3 + .../Core/McpStdioToolListChangedNotifier.cs | 59 +++++++++++++++ .../UnitTests/McpStdioServerRunAsyncTests.cs | 26 ++++++- .../McpStdioToolListChangedNotifierTests.cs | 74 +++++++++++++++++++ src/Service/Program.cs | 5 ++ 5 files changed, 166 insertions(+), 1 deletion(-) create mode 100644 src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs create mode 100644 src/Service.Tests/UnitTests/McpStdioToolListChangedNotifierTests.cs diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs index 08e5d7471a..fcf34ca2ba 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(); @@ -131,6 +133,7 @@ public async Task RunAsync(CancellationToken cancellationToken) break; case "notifications/initialized": + _toolListChangedNotifier?.MarkInitialized(); break; case "tools/list": diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs new file mode 100644 index 0000000000..827a82855c --- /dev/null +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs @@ -0,0 +1,59 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Text.Json; +using Azure.DataApiBuilder.Mcp.Model; +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 readonly McpStdoutWriter _stdoutWriter; + private int _isInitialized; + + public McpStdioToolListChangedNotifier(McpStdoutWriter stdoutWriter) + { + _stdoutWriter = stdoutWriter; + } + + /// + public void MarkInitialized() + { + Interlocked.Exchange(ref _isInitialized, 1); + } + + /// + public void NotifyToolsListChanged() + { + if (Volatile.Read(ref _isInitialized) == 0) + { + return; + } + + var notification = new + { + jsonrpc = McpStdioJsonRpcErrorCodes.JSON_RPC_VERSION, + method = NotificationMethods.ToolListChangedNotification, + @params = new { } + }; + + _stdoutWriter.WriteLine(JsonSerializer.Serialize(notification)); + } + } +} diff --git a/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs b/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs index 3477e775a5..61328f4c0f 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,25 @@ 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_InitializedNotification_MarksToolListNotifierReady() + { + 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.Once); + } + + private static (McpStdioServer server, StringWriter stdoutCapture) CreateServerWithCapturedOutput( + TextReader inputReader, + IMcpStdioToolListChangedNotifier? notifier = null) { StringWriter stdoutCapture = new(); McpStdoutWriter stdoutWriter = new(stdoutCapture); @@ -66,6 +85,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..f4266074e0 --- /dev/null +++ b/src/Service.Tests/UnitTests/McpStdioToolListChangedNotifierTests.cs @@ -0,0 +1,74 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#nullable enable + +using System; +using System.IO; +using System.Linq; +using System.Text.Json; +using Azure.DataApiBuilder.Mcp.Core; +using Microsoft.VisualStudio.TestTools.UnitTesting; + +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() + { + StringWriter output = new(); + using McpStdoutWriter stdoutWriter = new(output); + McpStdioToolListChangedNotifier notifier = new(stdoutWriter); + notifier.MarkInitialized(); + + notifier.NotifyToolsListChanged(); + + 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() + { + StringWriter output = new(); + using McpStdoutWriter stdoutWriter = new(output); + McpStdioToolListChangedNotifier notifier = new(stdoutWriter); + + notifier.MarkInitialized(); + notifier.MarkInitialized(); + notifier.NotifyToolsListChanged(); + + string[] lines = output.ToString().Split( + Environment.NewLine, + StringSplitOptions.RemoveEmptyEntries); + Assert.AreEqual(1, lines.Length); + } + } +} 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 => From a45bb416774b981cdd00953135ad0f6290c64afb Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Fri, 31 Jul 2026 00:19:54 -0700 Subject: [PATCH 05/21] fix(mcp): initialize tool schemas after metadata --- docs/design/McpToolRegistryHotReload.md | 23 +++++++++++++--- .../Core/McpToolRegistryRefreshService.cs | 7 ++++- .../Mcp/McpToolRegistryRefreshServiceTests.cs | 19 ++++++++++++++ .../UnitTests/McpStdioHelperTests.cs | 26 +++++++++++++++++-- src/Service/Startup.cs | 7 +++++ src/Service/Utilities/McpStdioHelper.cs | 8 ++++++ 6 files changed, 84 insertions(+), 6 deletions(-) diff --git a/docs/design/McpToolRegistryHotReload.md b/docs/design/McpToolRegistryHotReload.md index 9d088ae06b..03fbd1a501 100644 --- a/docs/design/McpToolRegistryHotReload.md +++ b/docs/design/McpToolRegistryHotReload.md @@ -240,13 +240,23 @@ The refresh service exposes one idempotent initialization path. #### HTTP mode -The host starts the same singleton through `IHostedService.StartAsync()`. +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. + +After `Startup.PerformOnConfigChangeAsync()` successfully initializes +`IMetadataProviderFactory`, it invokes the refresh service's idempotent initialization method. +This prevents the initial registry from advertising a configuration-only fallback schema while +database metadata is still being initialized. 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 hosted services do not run. `McpStdioHelper` explicitly resolves the same refresh-service singleton and invokes its idempotent initialization method before starting the stdio loop. +Stdio intentionally does not start the ASP.NET Core host, so `Startup.Configure` does not initialize +database metadata. `McpStdioHelper` explicitly initializes `IMetadataProviderFactory`, then resolves +the same refresh-service singleton and invokes its idempotent initialization method before starting +the stdio loop. This replaces the current duplicate per-tool registration path and gives both transports identical validation and metadata behavior. @@ -354,12 +364,16 @@ The current public `RegisterTool` method is not used by the new production path. ```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() + Host->>Refresh: StartAsync (subscribe only) + Startup->>Metadata: InitializeAsync() + Metadata-->>Startup: DB metadata ready + Startup->>Refresh: EnsureInitialized() Refresh->>Refresh: Capture current RuntimeConfig Refresh->>Factory: Create custom tools(config) Factory-->>Refresh: Fresh custom tools @@ -378,9 +392,12 @@ An invalid or duplicate tool causes startup to fail, preserving current strict s sequenceDiagram participant Helper as McpStdioHelper participant Refresh as McpToolRegistryRefreshService + participant Metadata as IMetadataProviderFactory participant Registry as McpToolRegistry participant Server as McpStdioServer + Helper->>Metadata: InitializeAsync() + Metadata-->>Helper: DB metadata ready Helper->>Refresh: EnsureInitialized() Refresh->>Registry: Build and publish initial snapshot Helper->>Server: RunAsync() diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs index 5ddff8aee6..e05ee9ad90 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs @@ -74,7 +74,12 @@ public void EnsureInitialized() public Task StartAsync(CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); - EnsureInitialized(); + + // 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; } diff --git a/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs index bf56702712..e062905b77 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs @@ -46,6 +46,25 @@ public void EnsureInitialized_IsIdempotentForSameConfig() context.Notifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Never); } + [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() { diff --git a/src/Service.Tests/UnitTests/McpStdioHelperTests.cs b/src/Service.Tests/UnitTests/McpStdioHelperTests.cs index 9d7fe855ac..8a8d0153ba 100644 --- a/src/Service.Tests/UnitTests/McpStdioHelperTests.cs +++ b/src/Service.Tests/UnitTests/McpStdioHelperTests.cs @@ -3,13 +3,16 @@ #nullable enable +using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; +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.VisualStudio.TestTools.UnitTesting; +using Moq; namespace Azure.DataApiBuilder.Service.Tests.UnitTests { @@ -22,8 +25,15 @@ public void RunMcpStdioHost_DoesNotStartWebHost() ServiceCollection services = new(); TestApplicationLifetime lifetime = new(); TestMcpStdioServer stdioServer = new(); - TestMcpToolRegistryRefreshService refreshService = new(); - + List initializationOrder = new(); + Mock metadataProviderFactory = new(); + metadataProviderFactory + .Setup(factory => factory.InitializeAsync()) + .Callback(() => initializationOrder.Add("metadata")) + .Returns(Task.CompletedTask); + TestMcpToolRegistryRefreshService refreshService = new(initializationOrder); + + services.AddSingleton(metadataProviderFactory.Object); services.AddSingleton(refreshService); services.AddSingleton(lifetime); services.AddSingleton(stdioServer); @@ -42,6 +52,10 @@ public void RunMcpStdioHost_DoesNotStartWebHost() "MCP stdio mode should still run the stdio JSON-RPC loop."); Assert.AreEqual(1, refreshService.EnsureInitializedCallCount, "MCP stdio mode should initialize the shared tool registry before running the loop."); + CollectionAssert.AreEqual( + new[] { "metadata", "registry" }, + 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, @@ -50,11 +64,19 @@ public void RunMcpStdioHost_DoesNotStartWebHost() 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"); } } diff --git a/src/Service/Startup.cs b/src/Service/Startup.cs index bcbaa235f8..a5a257a2ef 100644 --- a/src/Service/Startup.cs +++ b/src/Service/Startup.cs @@ -1439,6 +1439,13 @@ private async Task PerformOnConfigChangeAsync(IApplicationBuilder app) app.ApplicationServices.GetRequiredService(); await sqlMetadataProviderFactory.InitializeAsync(); + // Hosted services start before this metadata initialization. Publish the initial + // MCP registry only now so custom tool schemas use this initialized metadata + // generation. MCP services are absent when MCP was disabled at startup. + IMcpToolRegistryRefreshService? mcpToolRegistryRefreshService = + app.ApplicationServices.GetService(); + mcpToolRegistryRefreshService?.EnsureInitialized(); + // Manually trigger DI service instantiation of GraphQLSchemaCreator and RestService // to attempt to reduce chances that the first received client request // triggers instantiation and encounters undesired instantiation latency. diff --git a/src/Service/Utilities/McpStdioHelper.cs b/src/Service/Utilities/McpStdioHelper.cs index 799f344e80..b40a43caa4 100644 --- a/src/Service/Utilities/McpStdioHelper.cs +++ b/src/Service/Utilities/McpStdioHelper.cs @@ -4,6 +4,7 @@ using System; using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; +using Azure.DataApiBuilder.Core.Services.MetadataProviders; using Microsoft.Extensions.Configuration; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Hosting; @@ -78,6 +79,13 @@ public static bool RunMcpStdioHost(IHost host) { try { + // Stdio deliberately does not start the web host, so Startup.Configure does not + // initialize metadata. Do that explicitly before publishing the initial registry + // to keep custom tool schemas identical across transports. + IMetadataProviderFactory metadataProviderFactory = + host.Services.GetRequiredService(); + metadataProviderFactory.InitializeAsync().GetAwaiter().GetResult(); + Mcp.Core.IMcpToolRegistryRefreshService refreshService = host.Services.GetRequiredService(); refreshService.EnsureInitialized(); From 7d9b1e45110f182e5ebe9019d1d88d248af9560b Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Fri, 31 Jul 2026 00:31:33 -0700 Subject: [PATCH 06/21] fix(mcp): enforce registry discovery invariants --- docs/design/McpToolRegistryHotReload.md | 20 ++- .../Core/McpToolRegistry.cs | 88 +++++++++- .../Core/McpToolRegistryInitializer.cs | 23 ++- .../Mcp/McpToolRegistryInitializerTests.cs | 120 +++++++++++++ src/Service.Tests/Mcp/McpToolRegistryTests.cs | 159 ++++++++++++++++-- 5 files changed, 378 insertions(+), 32 deletions(-) create mode 100644 src/Service.Tests/Mcp/McpToolRegistryInitializerTests.cs diff --git a/docs/design/McpToolRegistryHotReload.md b/docs/design/McpToolRegistryHotReload.md index 03fbd1a501..cd0573561b 100644 --- a/docs/design/McpToolRegistryHotReload.md +++ b/docs/design/McpToolRegistryHotReload.md @@ -141,7 +141,14 @@ internal sealed record McpToolRegistrySnapshot( 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 stored in the snapshot are treated as immutable after construction. A list handler may create a new list container, but it must not mutate snapshot metadata. +Protocol `Tool` objects are mutable SDK models, so the registry defensively clones metadata during +candidate construction and again when returning public discovery results. Neither a tool retaining +its source metadata object nor a caller mutating a returned object can modify a published snapshot +or invalidate its fingerprint. + +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 @@ -324,7 +331,11 @@ The installed MCP SDK can send a notification through an individual `McpServer` 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. The comparison covers the complete serialized tool metadata, including name, description, input schema, and any future advertised fields. +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 +preserving array order. The comparison therefore ignores semantically irrelevant object insertion +order while still covering the complete tool metadata, including name, description, input schema, +and any future advertised fields. The swap still occurs when advertised metadata is equal, but `notifications/tools/list_changed` is emitted only when discovery metadata differs. @@ -357,6 +368,11 @@ McpToolRegistryUpdateResult ReplaceAll( The current public `RegisterTool` method is not used by the new production path. Because it is public, implementation should preserve it unless API review explicitly approves removal. If retained, it must use copy-on-write under the writer gate and must never mutate a published dictionary in place. +The obsolete `McpToolRegistryInitializer` compatibility fallback collects all tools and invokes +`ReplaceAll` with the current `RuntimeConfig`; it does not incrementally advertise disabled tools. +Manually assembled service providers using this fallback must register `RuntimeConfigProvider`, +because accurate snapshot discovery cannot be constructed without a configuration generation. + ## Detailed Flows ### Initial HTTP startup diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs index 7f2bc50117..356cdfd008 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs @@ -3,6 +3,7 @@ using System.Collections.Immutable; using System.Net; +using System.Text; using System.Text.Json; using System.Text.Json.Serialization; using Azure.DataApiBuilder.Config.ObjectModel; @@ -36,7 +37,7 @@ public void RegisterTool(IMcpTool tool) { ArgumentNullException.ThrowIfNull(tool); - Tool metadata = tool.GetToolMetadata(); + Tool metadata = CloneMetadata(tool.GetToolMetadata()); string toolName = ValidateToolName(metadata); lock (_writerLock) @@ -94,7 +95,7 @@ internal static McpToolRegistryCandidate CreateCandidate( { ArgumentNullException.ThrowIfNull(tool); - Tool metadata = tool.GetToolMetadata(); + Tool metadata = CloneMetadata(tool.GetToolMetadata()); string toolName = ValidateToolName(metadata); if (toolBuilder.TryGetValue(toolName, out IMcpTool? existingTool)) @@ -155,7 +156,10 @@ internal McpToolRegistryUpdateResult PublishCandidate(McpToolRegistryCandidate c /// public IReadOnlyList GetAdvertisedTools() { - return Volatile.Read(ref _snapshot).AdvertisedTools; + McpToolRegistrySnapshot snapshot = Volatile.Read(ref _snapshot); + return snapshot.AdvertisedTools + .Select(CloneMetadata) + .ToArray(); } /// @@ -169,7 +173,7 @@ public IEnumerable GetEnabledTools(RuntimeConfig config) McpToolRegistrySnapshot snapshot = Volatile.Read(ref _snapshot); return snapshot.Tools.Values .Where(t => t.IsEnabled(config)) - .Select(t => t.GetToolMetadata()); + .Select(t => CloneMetadata(t.GetToolMetadata())); } /// @@ -182,7 +186,7 @@ public bool TryGetTool(string toolName, out IMcpTool? tool) private static string ValidateToolName(Tool metadata) { - string toolName = metadata.Name?.Trim() ?? string.Empty; + string toolName = metadata.Name ?? string.Empty; if (string.IsNullOrWhiteSpace(toolName)) { throw new DataApiBuilderException( @@ -191,6 +195,14 @@ private static string ValidateToolName(Tool metadata) subStatusCode: DataApiBuilderException.SubStatusCodes.ErrorInInitialization); } + if (!string.Equals(toolName, toolName.Trim(), StringComparison.Ordinal)) + { + throw new DataApiBuilderException( + message: "MCP tool name cannot contain leading or trailing whitespace.", + statusCode: HttpStatusCode.ServiceUnavailable, + subStatusCode: DataApiBuilderException.SubStatusCodes.ErrorInInitialization); + } + return toolName; } @@ -221,7 +233,71 @@ private static ImmutableArray SortMetadata(IEnumerable metadata) private static string CreateDiscoveryFingerprint(ImmutableArray metadata) { - return JsonSerializer.Serialize(metadata.ToArray(), _discoveryJsonOptions); + JsonElement serializedMetadata = JsonSerializer.SerializeToElement( + metadata.ToArray(), + _discoveryJsonOptions); + using MemoryStream canonicalJson = new(); + using (Utf8JsonWriter writer = new(canonicalJson)) + { + WriteCanonicalJson(writer, serializedMetadata); + } + + 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) + { + 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); + } + + writer.WriteEndObject(); + break; + + case JsonValueKind.Array: + writer.WriteStartArray(); + foreach (JsonElement item in element.EnumerateArray()) + { + WriteCanonicalJson(writer, item); + } + + 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 sealed record McpToolRegistrySnapshot( diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs index b37e824929..fa9741d139 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +using Azure.DataApiBuilder.Core.Configurations; using Azure.DataApiBuilder.Mcp.Model; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Hosting; @@ -37,18 +38,22 @@ public Task StartAsync(CancellationToken cancellationToken) return Task.CompletedTask; } - // Preserve the legacy behavior for manually assembled service providers that have not - // registered the refresh service. - foreach (IMcpTool tool in _serviceProvider.GetServices()) + // Preserve compatibility for manually assembled service providers without publishing + // disabled tools. Snapshot discovery requires the configuration used to evaluate each + // tool's visibility, so this fallback now requires RuntimeConfigProvider as well. + RuntimeConfigProvider runtimeConfigProvider = + _serviceProvider.GetService() + ?? throw new InvalidOperationException( + $"{nameof(RuntimeConfigProvider)} must be registered when using the legacy " + + $"{nameof(McpToolRegistryInitializer)} fallback."); + IMcpTool[] tools = _serviceProvider.GetServices().ToArray(); + foreach (DynamicCustomTool customTool in tools.OfType()) { - if (tool is DynamicCustomTool customTool) - { - customTool.InitializeMetadata(_serviceProvider); - } - - _toolRegistry.RegisterTool(tool); + customTool.InitializeMetadata(_serviceProvider); } + _toolRegistry.ReplaceAll(tools, runtimeConfigProvider.GetConfig()); + return Task.CompletedTask; } diff --git a/src/Service.Tests/Mcp/McpToolRegistryInitializerTests.cs b/src/Service.Tests/Mcp/McpToolRegistryInitializerTests.cs new file mode 100644 index 0000000000..99ec895db4 --- /dev/null +++ b/src/Service.Tests/Mcp/McpToolRegistryInitializerTests.cs @@ -0,0 +1,120 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Azure.DataApiBuilder.Config; +using Azure.DataApiBuilder.Config.ObjectModel; +using Azure.DataApiBuilder.Core.Configurations; +using Azure.DataApiBuilder.Mcp.Core; +using Azure.DataApiBuilder.Mcp.Model; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.VisualStudio.TestTools.UnitTesting; +using ModelContextProtocol.Protocol; +using Moq; +using static Azure.DataApiBuilder.Mcp.Model.McpEnums; + +namespace Azure.DataApiBuilder.Service.Tests.Mcp +{ + [TestClass] + public class McpToolRegistryInitializerTests + { + [TestMethod] + public async Task LegacyFallback_UsesConfigAwareBulkReplacement() + { + RuntimeConfig config = CreateRuntimeConfig(); + Mock configLoader = new(null, null); + Mock configProvider = new(configLoader.Object); + configProvider.Setup(provider => provider.GetConfig()).Returns(config); + + ServiceCollection services = new(); + services.AddSingleton(configProvider.Object); + services.AddSingleton( + new TestMcpTool("enabled_tool", isEnabled: true)); + services.AddSingleton( + new TestMcpTool("disabled_tool", isEnabled: false)); + using ServiceProvider serviceProvider = services.BuildServiceProvider(); + McpToolRegistry registry = new(); +#pragma warning disable CS0618 // Explicitly exercises the documented compatibility fallback. + McpToolRegistryInitializer initializer = new(serviceProvider, registry); +#pragma warning restore CS0618 + + await initializer.StartAsync(CancellationToken.None); + + CollectionAssert.AreEqual( + new[] { "enabled_tool" }, + registry.GetAdvertisedTools().Select(tool => tool.Name).ToArray()); + Assert.IsTrue(registry.TryGetTool("enabled_tool", out _)); + Assert.IsTrue( + registry.TryGetTool("disabled_tool", out _), + "Disabled built-ins remain registered so execution can return a structured disabled response."); + } + + [TestMethod] + public async Task LegacyFallback_WithoutRuntimeConfigProvider_Throws() + { + ServiceCollection services = new(); + services.AddSingleton(new TestMcpTool("test_tool", isEnabled: true)); + using ServiceProvider serviceProvider = services.BuildServiceProvider(); +#pragma warning disable CS0618 // Explicitly exercises the documented compatibility fallback. + McpToolRegistryInitializer initializer = new(serviceProvider, new McpToolRegistry()); +#pragma warning restore CS0618 + + InvalidOperationException exception = await Assert.ThrowsExceptionAsync( + () => initializer.StartAsync(CancellationToken.None)); + + StringAssert.Contains(exception.Message, nameof(RuntimeConfigProvider)); + } + + 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 TestMcpTool : IMcpTool + { + private readonly string _name; + private readonly bool _isEnabled; + + public TestMcpTool(string name, bool isEnabled) + { + _name = name; + _isEnabled = isEnabled; + } + + public ToolType ToolType => ToolType.BuiltIn; + + public bool IsEnabled(RuntimeConfig config) => _isEnabled; + + public Tool GetToolMetadata() + { + return new Tool + { + Name = _name, + Description = "Test tool", + InputSchema = JsonSerializer.Deserialize("{\"type\":\"object\"}") + }; + } + + public Task ExecuteAsync( + JsonDocument? arguments, + IServiceProvider serviceProvider, + CancellationToken cancellationToken = default) + { + throw new NotImplementedException(); + } + } + } +} diff --git a/src/Service.Tests/Mcp/McpToolRegistryTests.cs b/src/Service.Tests/Mcp/McpToolRegistryTests.cs index 68d25a4146..6a2c88b73c 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryTests.cs @@ -251,24 +251,20 @@ public void RegisterTool_WithRealisticBuiltInToolNames_DetectsDuplicates() } /// - /// 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. + /// Test that leading/trailing whitespace is rejected rather than producing a lookup key + /// that differs from the advertised tool name. /// [TestMethod] - public void RegisterTool_WithLeadingTrailingWhitespace_DetectsDuplicate() + public void RegisterTool_WithLeadingTrailingWhitespace_ThrowsException() { - // Arrange McpToolRegistry registry = new(); - IMcpTool tool1 = new MockMcpTool("my_tool", ToolType.BuiltIn); - IMcpTool tool2 = new MockMcpTool(" my_tool ", ToolType.Custom); + IMcpTool tool = new MockMcpTool(" my_tool ", ToolType.Custom); - // Act - registry.RegisterTool(tool1); + DataApiBuilderException exception = Assert.ThrowsException( + () => registry.RegisterTool(tool)); - // Assert - trimmed name should collide - Assert.ThrowsException( - () => registry.RegisterTool(tool2) - ); + StringAssert.Contains(exception.Message, "leading or trailing whitespace"); + Assert.IsFalse(registry.TryGetTool("my_tool", out _)); } /// @@ -398,6 +394,30 @@ public void ReplaceAll_PublishesCompleteOrderedSnapshot() Assert.AreEqual(2, result.AdvertisedToolCount); } + /// + /// 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()); + + 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); + } + } + /// /// A candidate containing a duplicate name is rejected before publication, leaving the /// complete previous snapshot active. @@ -447,6 +467,89 @@ public void ReplaceAll_WithEquivalentMetadata_DoesNotReportDiscoveryChange() Assert.IsFalse(result.DiscoveryChanged); } + /// + /// 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 ReplaceAll_WithEquivalentSchemaPropertyOrder_DoesNotReportDiscoveryChange() + { + 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(); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, inputSchemaJson: SCHEMA_AB) }, + config); + + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, inputSchemaJson: SCHEMA_BA) }, + config); + + Assert.IsFalse(result.DiscoveryChanged); + } + + /// + /// A real input-schema change remains client-visible after canonicalization. + /// + [TestMethod] + public void ReplaceAll_WithChangedInputSchema_ReportsDiscoveryChange() + { + McpToolRegistry registry = new(); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll( + new[] + { + new MockMcpTool( + "same_tool", + ToolType.Custom, + inputSchemaJson: "{\"type\":\"object\",\"properties\":{\"id\":{\"type\":\"string\"}}}") + }, + config); + + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] + { + new MockMcpTool( + "same_tool", + ToolType.Custom, + inputSchemaJson: "{\"type\":\"object\",\"properties\":{\"id\":{\"type\":\"integer\"}}}") + }, + config); + + Assert.IsTrue(result.DiscoveryChanged); + } + + /// + /// Published metadata is isolated both from the tool-owned source object and from callers + /// mutating a value returned by the public snapshot accessor. + /// + [TestMethod] + public void ReplaceAll_DefensivelyClonesPublishedMetadata() + { + Tool retainedMetadata = new() + { + 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. /// @@ -647,17 +750,20 @@ private class MockMcpTool : IMcpTool private readonly string _toolName; private readonly Func? _isEnabledFunc; private readonly string _description; + private readonly string _inputSchemaJson; public MockMcpTool( string toolName, ToolType toolType, Func? isEnabledFunc = null, - string? description = 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; } @@ -669,8 +775,7 @@ 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, @@ -689,6 +794,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. /// From 3724bf96e8103a631c23830d4867efe628cab7a1 Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Fri, 31 Jul 2026 00:46:51 -0700 Subject: [PATCH 07/21] test(mcp): cover registry reload transports and failures --- ...tpToolRegistryHotReloadIntegrationTests.cs | 406 ++++++++++++++++++ ...ioToolRegistryHotReloadIntegrationTests.cs | 199 +++++++++ .../Mcp/McpToolRegistryRefreshServiceTests.cs | 186 +++++++- 3 files changed, 788 insertions(+), 3 deletions(-) create mode 100644 src/Service.Tests/Mcp/McpHttpToolRegistryHotReloadIntegrationTests.cs create mode 100644 src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs diff --git a/src/Service.Tests/Mcp/McpHttpToolRegistryHotReloadIntegrationTests.cs b/src/Service.Tests/Mcp/McpHttpToolRegistryHotReloadIntegrationTests.cs new file mode 100644 index 0000000000..a2e17909d8 --- /dev/null +++ b/src/Service.Tests/Mcp/McpHttpToolRegistryHotReloadIntegrationTests.cs @@ -0,0 +1,406 @@ +// 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.TestHost; +using Microsoft.Data.SqlClient; +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 TestServer server = new(Program.CreateWebHostBuilder(args)); + 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); + + 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 Task.Delay(500); + 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) + { + Dictionary entities = tools.ToDictionary( + tool => tool.EntityName, + tool => new Entity( + Source: new( + Object: "get_book_by_id", + 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(false)), + 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 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); + } +} diff --git a/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs b/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs new file mode 100644 index 0000000000..32b78012fb --- /dev/null +++ b/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs @@ -0,0 +1,199 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#nullable enable + +using System; +using System.Collections.Generic; +using System.IO; +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; +using static Azure.DataApiBuilder.Config.DabConfigEvents; + +namespace Azure.DataApiBuilder.Service.Tests.Mcp +{ + [TestClass] + public class McpStdioToolRegistryHotReloadIntegrationTests + { + [TestMethod] + public async Task InitializedClient_RefreshesRegistry_ReceivesNotificationAndUpdatedList() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(); + Mock configLoader = new(null, null); + 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); + + McpToolRegistry registry = new(); + HotReloadEventHandler hotReloadEventHandler = 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."); + + currentConfig = CreateRuntimeConfig(("GetBook", "Gets one book")); + hotReloadEventHandler.OnConfigChangedEvent( + hotReloadEventHandler, + new HotReloadEventArgs(MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED, string.Empty)); + + 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); + 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()); + await serverTask; + } + + private static async Task ReadJsonLineAsync( + ChannelTextWriter output, + CancellationToken cancellationToken) + { + string line = await output.ReadLineAsync(cancellationToken); + return JsonDocument.Parse(line); + } + + 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 index e062905b77..101e99cf1a 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs @@ -147,6 +147,115 @@ public void HotReload_WithEquivalentDiscoveryMetadata_DoesNotNotify() 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()); + } + + [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 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() { @@ -218,6 +327,20 @@ public void RuntimeConfigLoader_RaisesMcpEventAfterDependenciesAndBeforeGraphQL( 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) { Mock configLoader = new(null, null); Mock configProvider = new(configLoader.Object); @@ -232,7 +355,6 @@ private static TestContext CreateContext( .Setup(factory => factory.GetMetadataProvider(It.IsAny())) .Returns(sqlMetadataProvider.Object); - Mock notifier = new(); McpToolRegistry registry = new(); HotReloadEventHandler hotReloadEventHandler = new(); McpToolRegistryRefreshService service = new( @@ -240,14 +362,14 @@ private static TestContext CreateContext( builtInTools, registry, metadataProviderFactory.Object, - new[] { notifier.Object }, + notifiers, NullLogger.Instance, hotReloadEventHandler); return new TestContext( service, registry, - notifier, + primaryNotifier, hotReloadEventHandler); } @@ -286,6 +408,53 @@ private static RuntimeConfig CreateRuntimeConfig( 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)); @@ -350,6 +519,17 @@ public Task ExecuteAsync( } } + private sealed class ThrowingNotifier : IMcpToolListChangedNotifier + { + public int CallCount { get; private set; } + + public void NotifyToolsListChanged() + { + CallCount++; + throw new InvalidOperationException("Expected notification failure."); + } + } + private sealed class TestRuntimeConfigLoader : RuntimeConfigLoader { public TestRuntimeConfigLoader(HotReloadEventHandler handler) From 5dc1e8341b2a1242c902a5ee86518b2419dbb41d Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Fri, 31 Jul 2026 02:15:43 -0700 Subject: [PATCH 08/21] fix(config): serialize hot reload generations --- docs/design/McpToolRegistryHotReload.md | 21 ++- .../Core/DynamicCustomTool.cs | 29 +++- .../Core/McpToolRegistry.cs | 14 +- .../Core/McpToolRegistryRefreshService.cs | 21 ++- .../Utils/McpMetadataHelper.cs | 3 +- src/Config/FileSystemRuntimeConfigLoader.cs | 81 +++++++--- .../Mcp/McpToolRegistryRefreshServiceTests.cs | 36 ++++- .../UnitTests/ConfigFileWatcherUnitTests.cs | 140 ++++++++++++++++++ 8 files changed, 296 insertions(+), 49 deletions(-) diff --git a/docs/design/McpToolRegistryHotReload.md b/docs/design/McpToolRegistryHotReload.md index cd0573561b..80ba248e3c 100644 --- a/docs/design/McpToolRegistryHotReload.md +++ b/docs/design/McpToolRegistryHotReload.md @@ -271,7 +271,13 @@ This replaces the current duplicate per-tool registration path and gives both tr Distinct file edits can produce overlapping hot-reload callbacks even though duplicate notifications for one file content are suppressed by `ConfigFileWatcher`. -The refresh service serializes its own rebuilds and uses a stale-generation guard: +`FileSystemRuntimeConfigLoader` serializes the complete reload operation per loader instance. The +gate is acquired before loading the new configuration and remains held until every synchronous +`SignalConfigChanged()` handler returns. Consequently, configuration, metadata, authorization, MCP, +GraphQL, and logging handlers for generation A complete before generation B can replace the active +configuration or begin updating dependencies. + +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()`. @@ -281,7 +287,10 @@ The refresh service serializes its own rebuilds and uses a stale-generation guar 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. -This guard prevents an older, slower registry rebuild from overwriting a newer registry generation. It does not claim to make the full DAB hot-reload pipeline transactional; that remains separate work. +The loader gate prevents mixed dependency generations. 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. ### 11. Existing tool-call safety is preserved @@ -479,9 +488,13 @@ An in-flight `tools/call` retains the resolved tool instance. A later swap does ### Multiple configuration changes -The stale-generation guard prevents publication when the active `RuntimeConfig` changed during candidate construction. The latest callback eventually publishes the latest generation. +The per-loader reload gate ensures one file generation completes all ordered handlers before the next +file 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. -The design intentionally does not lock the entire DAB hot-reload pipeline. Global serialization and transactional rollback are tracked separately. +Serialization is scoped to each `FileSystemRuntimeConfigLoader`; independent loaders do not block one +another. Transactional rollback is still tracked separately. ## Failure Semantics diff --git a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs index d991e00277..be73fd7cc9 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs @@ -72,7 +72,7 @@ public DynamicCustomTool(string entityName, Entity entity) /// 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. + /// Falls back to a config-based schema if DB metadata is unavailable. /// /// The application service provider with initialized metadata providers. public void InitializeMetadata(IServiceProvider serviceProvider) @@ -97,10 +97,26 @@ public void InitializeMetadata(IServiceProvider serviceProvider) 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. + /// + public bool InitializeMetadata( + RuntimeConfig config, + IMetadataProviderFactory metadataProviderFactory, + out string fallbackReason) { ArgumentNullException.ThrowIfNull(config); ArgumentNullException.ThrowIfNull(metadataProviderFactory); - _cachedInputSchema = BuildInputSchemaFromDbMetadata(config, metadataProviderFactory); + _cachedInputSchema = BuildInputSchemaFromDbMetadata( + config, + metadataProviderFactory, + out fallbackReason); return _cachedInputSchema.HasValue; } @@ -343,7 +359,8 @@ private JsonElement BuildInputSchema() /// private JsonElement? BuildInputSchemaFromDbMetadata( RuntimeConfig config, - IMetadataProviderFactory metadataProviderFactory) + IMetadataProviderFactory metadataProviderFactory, + out string fallbackReason) { if (!McpMetadataHelper.TryResolveMetadata( EntityName, @@ -352,16 +369,20 @@ private JsonElement BuildInputSchema() 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/McpToolRegistry.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs index 356cdfd008..81506094be 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs @@ -195,13 +195,13 @@ private static string ValidateToolName(Tool metadata) subStatusCode: DataApiBuilderException.SubStatusCodes.ErrorInInitialization); } - if (!string.Equals(toolName, toolName.Trim(), StringComparison.Ordinal)) - { - throw new DataApiBuilderException( - message: "MCP tool name cannot contain leading or trailing whitespace.", - statusCode: HttpStatusCode.ServiceUnavailable, - subStatusCode: DataApiBuilderException.SubStatusCodes.ErrorInInitialization); - } + if (!string.Equals(toolName, toolName.Trim(), StringComparison.Ordinal)) + { + throw new DataApiBuilderException( + message: "MCP tool name cannot contain leading or trailing whitespace.", + statusCode: HttpStatusCode.ServiceUnavailable, + subStatusCode: DataApiBuilderException.SubStatusCodes.ErrorInInitialization); + } return toolName; } diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs index e05ee9ad90..e2465e0e19 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs @@ -122,14 +122,16 @@ private void RefreshRegistry() { bool initializedFromDatabase = customTool.InitializeMetadata( config, - _metadataProviderFactory); + _metadataProviderFactory, + out string fallbackReason); if (!initializedFromDatabase) { _logger.LogWarning( - "Database metadata was unavailable for custom MCP tool '{ToolName}' " + - "on entity '{EntityName}'. Using configuration-derived input schema.", + "Using configuration-derived input schema for custom MCP tool " + + "'{ToolName}' on entity '{EntityName}'. Reason: {FallbackReason}", customTool.GetToolMetadata().Name, - customTool.EntityName); + customTool.EntityName, + fallbackReason); } } @@ -150,11 +152,16 @@ private void RefreshRegistry() _lastAppliedConfig = config; _logger.LogInformation( - "Published MCP tool registry version {Version} with {RegisteredToolCount} " + - "registered tools and {AdvertisedToolCount} advertised tools.", + "Published MCP tool registry version {Version} with {BuiltInToolCount} " + + "built-in tools, {CustomToolCount} custom tools, {RegisteredToolCount} " + + "registered tools, and {AdvertisedToolCount} advertised tools. " + + "Discovery changed: {DiscoveryChanged}.", result.Version, + _builtInTools.Count, + customTools.Count, result.RegisteredToolCount, - result.AdvertisedToolCount); + result.AdvertisedToolCount, + result.DiscoveryChanged); if (!isInitialGeneration && result.DiscoveryChanged) { diff --git a/src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs b/src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs index 2ca7e116e9..fe05f50446 100644 --- a/src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs +++ b/src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs @@ -150,7 +150,8 @@ 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; } diff --git a/src/Config/FileSystemRuntimeConfigLoader.cs b/src/Config/FileSystemRuntimeConfigLoader.cs index 016931bf62..82f2f5af1b 100644 --- a/src/Config/FileSystemRuntimeConfigLoader.cs +++ b/src/Config/FileSystemRuntimeConfigLoader.cs @@ -32,6 +32,7 @@ namespace Azure.DataApiBuilder.Config; /// public class FileSystemRuntimeConfigLoader : RuntimeConfigLoader, IDisposable { + private readonly object _hotReloadGate = new(); private bool _disposed; /// /// This stores either the default config name e.g. dab-config.json @@ -111,19 +112,28 @@ public FileSystemRuntimeConfigLoader( /// public void Dispose() { - if (_disposed) + ConfigFileWatcher? configFileWatcher; + lock (_hotReloadGate) { - return; - } - - _disposed = true; + if (_disposed) + { + return; + } - if (_configFileWatcher is not null) - { - _configFileWatcher.NewFileContentsDetected -= OnNewFileContentsDetected; - _configFileWatcher.Dispose(); + _disposed = true; + configFileWatcher = _configFileWatcher; _configFileWatcher = null; + + if (configFileWatcher is not null) + { + configFileWatcher.NewFileContentsDetected -= OnNewFileContentsDetected; + } } + + // FileSystemWatcher disposal can block while an OS callback completes. Do not hold the + // reload gate during that external operation. Any callback already queued will observe + // _disposed after it acquires the gate and return without loading another generation. + configFileWatcher?.Dispose(); } /// @@ -198,18 +208,43 @@ private bool TrySetupConfigFileWatcher() /// private void OnNewFileContentsDetected(object? sender, EventArgs e) { - try + 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(); + + lock (_hotReloadGate) { - if (RuntimeConfig is not null) + if (_disposed) { - HotReloadConfig(RuntimeConfig.IsDevelopmentMode()); + return; + } + + try + { + if (RuntimeConfig is not null) + { + HotReloadConfig(RuntimeConfig.IsDevelopmentMode()); + } + } + catch (Exception ex) + { + SendLogToBufferOrLogger( + LogLevel.Error, + $"Unable to hot reload configuration file due to {ex.Message}"); } - } - 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); } } @@ -341,14 +376,16 @@ 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) { - logger?.LogInformation(message: "Starting hot-reload process for config: {ConfigFilePath}", ConfigFilePath); + 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 _, isDevMode: isDevMode, replacementSettings: replacementSettings)) { throw new DataApiBuilderException( message: "Deserialization of the configuration file failed.", @@ -360,7 +397,7 @@ private void HotReloadConfig(bool isDevMode, ILogger? logger = null) IsNewConfigValidated = false; SignalConfigChanged(); - logger?.LogInformation("Hot-reload process finished."); + SendLogToBufferOrLogger(LogLevel.Information, "Hot-reload process finished."); } /// diff --git a/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs index 101e99cf1a..afee7b8cc0 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs @@ -17,7 +17,7 @@ using Azure.DataApiBuilder.Mcp.Core; using Azure.DataApiBuilder.Mcp.Model; using Azure.DataApiBuilder.Service.Exceptions; -using Microsoft.Extensions.Logging.Abstractions; +using Microsoft.Extensions.Logging; using Microsoft.VisualStudio.TestTools.UnitTesting; using ModelContextProtocol.Protocol; using Moq; @@ -169,6 +169,15 @@ public void EnsureInitialized_WhenDatabaseMetadataUnavailable_PublishesConfigFal .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, 1 custom tools, 1 registered tools, and 1 advertised tools. " + + "Discovery changed: True."); } [TestMethod] @@ -357,20 +366,22 @@ private static TestContext CreateContextWithNotifiers( McpToolRegistry registry = new(); HotReloadEventHandler hotReloadEventHandler = new(); + Mock> logger = new(); McpToolRegistryRefreshService service = new( configProvider.Object, builtInTools, registry, metadataProviderFactory.Object, notifiers, - NullLogger.Instance, + logger.Object, hotReloadEventHandler); return new TestContext( service, registry, primaryNotifier, - hotReloadEventHandler); + hotReloadEventHandler, + logger); } private static RuntimeConfig CreateRuntimeConfig( @@ -480,11 +491,28 @@ private static Tool CreateMetadata(string name, string description) }; } + 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); + HotReloadEventHandler HotReloadEventHandler, + Mock> Logger); private sealed class TestMcpTool : IMcpTool { diff --git a/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs b/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs index b0ae580828..077c4b23a5 100644 --- a/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs +++ b/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs @@ -1,15 +1,20 @@ // 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 Microsoft.VisualStudio.TestTools.UnitTesting; using Moq; +using static Azure.DataApiBuilder.Config.DabConfigEvents; namespace Azure.DataApiBuilder.Service.Tests.UnitTests; @@ -159,6 +164,141 @@ 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); + } + } + #region ConfigFileWatcher NewFileContentsDetected event invocation tests private const string UNEXPECTED_INVOCATION_COUNT_ERR = "Unexpected number of invocations of the NewFileContentsDetected event."; From 4c702c26fe8ce7e17110c3d67372fc0178efd614 Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Fri, 31 Jul 2026 03:55:04 -0700 Subject: [PATCH 09/21] fix(mcp): serialize initial registry construction --- docs/design/McpToolRegistryHotReload.md | 38 ++- src/Config/FileSystemRuntimeConfigLoader.cs | 43 ++- .../McpInitialHotReloadSerializationTests.cs | 287 ++++++++++++++++++ .../UnitTests/McpStdioHelperTests.cs | 18 ++ src/Service/Startup.cs | 19 +- src/Service/Utilities/McpStdioHelper.cs | 16 +- .../Utilities/RuntimeInitializationHelper.cs | 58 ++++ 7 files changed, 434 insertions(+), 45 deletions(-) create mode 100644 src/Service.Tests/Mcp/McpInitialHotReloadSerializationTests.cs create mode 100644 src/Service/Utilities/RuntimeInitializationHelper.cs diff --git a/docs/design/McpToolRegistryHotReload.md b/docs/design/McpToolRegistryHotReload.md index 80ba248e3c..bda42c0aed 100644 --- a/docs/design/McpToolRegistryHotReload.md +++ b/docs/design/McpToolRegistryHotReload.md @@ -251,19 +251,19 @@ The host resolves the singleton through `IHostedService` early enough to subscri hot-reload events, but `StartAsync()` does not publish the initial snapshot. ASP.NET Core starts hosted services before `Startup.Configure` finishes initializing database metadata. -After `Startup.PerformOnConfigChangeAsync()` successfully initializes -`IMetadataProviderFactory`, it invokes the refresh service's idempotent initialization method. -This prevents the initial registry from advertising a configuration-only fallback schema while -database metadata is still being initialized. +`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` explicitly initializes `IMetadataProviderFactory`, then resolves -the same refresh-service singleton and invokes its idempotent initialization method before starting -the stdio loop. +database metadata. `McpStdioHelper` invokes the same serialized initial dependency operation used by +HTTP startup before starting the stdio loop. This replaces the current duplicate per-tool registration path and gives both transports identical validation and metadata behavior. @@ -271,11 +271,12 @@ This replaces the current duplicate per-tool registration path and gives both tr Distinct file edits can produce overlapping hot-reload callbacks even though duplicate notifications for one file content are suppressed by `ConfigFileWatcher`. -`FileSystemRuntimeConfigLoader` serializes the complete reload operation per loader instance. The -gate is acquired before loading the new configuration and remains held until every synchronous -`SignalConfigChanged()` handler returns. Consequently, configuration, metadata, authorization, MCP, -GraphQL, and logging handlers for generation A complete before generation B can replace the active -configuration or begin updating dependencies. +`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. The refresh service retains its own writer gate and stale-generation guard as defense in depth: @@ -488,10 +489,11 @@ An in-flight `tools/call` retains the resolved tool instance. A later swap does ### Multiple configuration changes -The per-loader reload gate ensures one file generation completes all ordered handlers before the next -file 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. +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. @@ -571,6 +573,7 @@ The exact file split may change during implementation, but the expected touchpoi - [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 @@ -586,6 +589,7 @@ The exact file split may change during implementation, but the expected touchpoi ### 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. @@ -649,6 +653,8 @@ The exact file split may change during implementation, but the expected touchpoi 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. 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. diff --git a/src/Config/FileSystemRuntimeConfigLoader.cs b/src/Config/FileSystemRuntimeConfigLoader.cs index 82f2f5af1b..005009538c 100644 --- a/src/Config/FileSystemRuntimeConfigLoader.cs +++ b/src/Config/FileSystemRuntimeConfigLoader.cs @@ -32,7 +32,7 @@ namespace Azure.DataApiBuilder.Config; /// public class FileSystemRuntimeConfigLoader : RuntimeConfigLoader, IDisposable { - private readonly object _hotReloadGate = new(); + private readonly SemaphoreSlim _hotReloadGate = new(initialCount: 1, maxCount: 1); private bool _disposed; /// /// This stores either the default config name e.g. dab-config.json @@ -113,7 +113,8 @@ public FileSystemRuntimeConfigLoader( public void Dispose() { ConfigFileWatcher? configFileWatcher; - lock (_hotReloadGate) + _hotReloadGate.Wait(); + try { if (_disposed) { @@ -129,6 +130,10 @@ public void Dispose() configFileWatcher.NewFileContentsDetected -= OnNewFileContentsDetected; } } + finally + { + _hotReloadGate.Release(); + } // FileSystemWatcher disposal can block while an OS callback completes. Do not hold the // reload gate during that external operation. Any callback already queued will observe @@ -225,7 +230,8 @@ internal void ProcessHotReloadNotification(Action? beforeEnteringGate = null) { beforeEnteringGate?.Invoke(); - lock (_hotReloadGate) + _hotReloadGate.Wait(); + try { if (_disposed) { @@ -246,6 +252,37 @@ internal void ProcessHotReloadNotification(Action? beforeEnteringGate = null) $"Unable to hot reload configuration file due to {ex.Message}"); } } + finally + { + _hotReloadGate.Release(); + } + } + + /// + /// 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 async Task ExecuteWithHotReloadSerializationAsync(Func operation) + { + ArgumentNullException.ThrowIfNull(operation); + + await _hotReloadGate.WaitAsync().ConfigureAwait(false); + try + { + if (_disposed) + { + throw new ObjectDisposedException(nameof(FileSystemRuntimeConfigLoader)); + } + + await operation().ConfigureAwait(false); + } + finally + { + _hotReloadGate.Release(); + } } /// diff --git a/src/Service.Tests/Mcp/McpInitialHotReloadSerializationTests.cs b/src/Service.Tests/Mcp/McpInitialHotReloadSerializationTests.cs new file mode 100644 index 0000000000..49d8ea8756 --- /dev/null +++ b/src/Service.Tests/Mcp/McpInitialHotReloadSerializationTests.cs @@ -0,0 +1,287 @@ +// 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()) + .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(), 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/UnitTests/McpStdioHelperTests.cs b/src/Service.Tests/UnitTests/McpStdioHelperTests.cs index 8a8d0153ba..87f64095b2 100644 --- a/src/Service.Tests/UnitTests/McpStdioHelperTests.cs +++ b/src/Service.Tests/UnitTests/McpStdioHelperTests.cs @@ -4,13 +4,17 @@ #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; @@ -26,6 +30,17 @@ public void RunMcpStdioHost_DoesNotStartWebHost() TestApplicationLifetime lifetime = new(); TestMcpStdioServer stdioServer = new(); 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()) @@ -33,6 +48,9 @@ public void RunMcpStdioHost_DoesNotStartWebHost() .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); diff --git a/src/Service/Startup.cs b/src/Service/Startup.cs index a5a257a2ef..36aa725f50 100644 --- a/src/Service/Startup.cs +++ b/src/Service/Startup.cs @@ -1426,25 +1426,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(); - - // Hosted services start before this metadata initialization. Publish the initial - // MCP registry only now so custom tool schemas use this initialized metadata - // generation. MCP services are absent when MCP was disabled at startup. - IMcpToolRegistryRefreshService? mcpToolRegistryRefreshService = - app.ApplicationServices.GetService(); - mcpToolRegistryRefreshService?.EnsureInitialized(); // 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/Utilities/McpStdioHelper.cs b/src/Service/Utilities/McpStdioHelper.cs index b40a43caa4..4c969acbf7 100644 --- a/src/Service/Utilities/McpStdioHelper.cs +++ b/src/Service/Utilities/McpStdioHelper.cs @@ -4,7 +4,6 @@ using System; using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; -using Azure.DataApiBuilder.Core.Services.MetadataProviders; using Microsoft.Extensions.Configuration; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Hosting; @@ -80,15 +79,12 @@ public static bool RunMcpStdioHost(IHost host) try { // Stdio deliberately does not start the web host, so Startup.Configure does not - // initialize metadata. Do that explicitly before publishing the initial registry - // to keep custom tool schemas identical across transports. - IMetadataProviderFactory metadataProviderFactory = - host.Services.GetRequiredService(); - metadataProviderFactory.InitializeAsync().GetAwaiter().GetResult(); - - Mcp.Core.IMcpToolRegistryRefreshService refreshService = - host.Services.GetRequiredService(); - refreshService.EnsureInitialized(); + // 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(); diff --git a/src/Service/Utilities/RuntimeInitializationHelper.cs b/src/Service/Utilities/RuntimeInitializationHelper.cs new file mode 100644 index 0000000000..4833c3fe84 --- /dev/null +++ b/src/Service/Utilities/RuntimeInitializationHelper.cs @@ -0,0 +1,58 @@ +// 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 () => + { + RuntimeConfigProvider runtimeConfigProvider = + serviceProvider.GetRequiredService(); + initializedConfig = runtimeConfigProvider.GetConfig(); + + RuntimeConfigValidator runtimeConfigValidator = + serviceProvider.GetRequiredService(); + runtimeConfigValidator.ValidateConfigProperties(); + + IMetadataProviderFactory metadataProviderFactory = + serviceProvider.GetRequiredService(); + await metadataProviderFactory.InitializeAsync().ConfigureAwait(false); + + // MCP services are absent when MCP was disabled at startup. + IMcpToolRegistryRefreshService? mcpToolRegistryRefreshService = + serviceProvider.GetService(); + mcpToolRegistryRefreshService?.EnsureInitialized(); + }).ConfigureAwait(false); + + return initializedConfig!; + } + } +} From 35e6325cadeb223a8473772c6d4d3801ddd13709 Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Fri, 31 Jul 2026 05:13:07 -0700 Subject: [PATCH 10/21] fix(mcp): harden registry notifications and APIs --- docs/design/McpToolRegistryHotReload.md | 19 ++- .../Core/CustomMcpToolFactory.cs | 9 +- .../Core/McpServerConfiguration.cs | 4 + .../Core/McpStdioToolListChangedNotifier.cs | 72 ++++++++- .../Core/McpToolRegistry.cs | 28 ++-- .../Core/McpToolRegistryRefreshService.cs | 35 +++-- .../Mcp/McpToolRegistryRefreshServiceTests.cs | 65 ++++++++ src/Service.Tests/Mcp/McpToolRegistryTests.cs | 2 + .../UnitTests/McpServerConfigurationTests.cs | 36 +++++ .../McpStdioToolListChangedNotifierTests.cs | 141 +++++++++++++++++- src/Service/Utilities/McpStdioHelper.cs | 4 + 11 files changed, 376 insertions(+), 39 deletions(-) create mode 100644 src/Service.Tests/UnitTests/McpServerConfigurationTests.cs diff --git a/docs/design/McpToolRegistryHotReload.md b/docs/design/McpToolRegistryHotReload.md index bda42c0aed..177934140b 100644 --- a/docs/design/McpToolRegistryHotReload.md +++ b/docs/design/McpToolRegistryHotReload.md @@ -144,7 +144,7 @@ Keeping lookup state and advertised metadata in the same snapshot prevents a req Protocol `Tool` objects are mutable SDK models, so the registry defensively clones metadata during candidate construction and again when returning public discovery results. Neither a tool retaining its source metadata object nor a caller mutating a returned object can modify a published snapshot -or invalidate its fingerprint. +or invalidate its canonical discovery representation. 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 @@ -324,16 +324,24 @@ Notification rules: - 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 will own initialization state and frame writing. `McpStdioServer` marks it initialized when handling `notifications/initialized`. The refresh service depends on zero or more tool-list notifiers; HTTP mode has no notifier registered in this iteration. +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` marks the notifier initialized when +handling `notifications/initialized`. The refresh service depends on zero or more tool-list +notifiers; HTTP mode has no notifier registered in this iteration. ### 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 must not be advertised until HTTP notification delivery is implemented. +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. @@ -378,6 +386,10 @@ McpToolRegistryUpdateResult ReplaceAll( The current public `RegisterTool` method is not used by the new production path. Because it is public, implementation should preserve it unless API review explicitly approves removal. If retained, it must use copy-on-write under the writer gate and must never mutate a published dictionary in place. +The compatibility `GetEnabledTools(RuntimeConfig)` accessor is obsolete because combining a +published tool map with caller-supplied configuration can cross generations. Production discovery +must use `GetAdvertisedTools()`, whose visibility and metadata were captured with the snapshot. + The obsolete `McpToolRegistryInitializer` compatibility fallback collects all tools and invokes `ReplaceAll` with the current `RuntimeConfig`; it does not incrementally advertise disabled tools. Manually assembled service providers using this fallback must register `RuntimeConfigProvider`, @@ -641,6 +653,7 @@ The exact file split may change during implementation, but the expected touchpoi 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. ### Hot-reload integration tests diff --git a/src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs b/src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs index 67303f6dbf..6f3cb306d2 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) { diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs index 8b5c943863..f40497a0bb 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs @@ -91,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/McpStdioToolListChangedNotifier.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs index 827a82855c..3af2fc2a72 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs @@ -3,6 +3,8 @@ 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 @@ -24,12 +26,25 @@ public interface IMcpStdioToolListChangedNotifier : IMcpToolListChangedNotifier /// 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 int _isInitialized; + private int _pendingNotificationCount; + private int _notificationWorkerScheduled; - public McpStdioToolListChangedNotifier(McpStdoutWriter stdoutWriter) + public McpStdioToolListChangedNotifier( + McpStdoutWriter stdoutWriter, + ILogger? logger = null) { _stdoutWriter = stdoutWriter; + _logger = logger ?? NullLogger.Instance; } /// @@ -46,14 +61,57 @@ public void NotifyToolsListChanged() return; } - var notification = new + Interlocked.Increment(ref _pendingNotificationCount); + ScheduleNotificationWorker(); + } + + private void ScheduleNotificationWorker() + { + if (Interlocked.CompareExchange(ref _notificationWorkerScheduled, 1, 0) != 0) { - jsonrpc = McpStdioJsonRpcErrorCodes.JSON_RPC_VERSION, - method = NotificationMethods.ToolListChangedNotification, - @params = new { } - }; + return; + } - _stdoutWriter.WriteLine(JsonSerializer.Serialize(notification)); + if (!ThreadPool.QueueUserWorkItem( + static notifier => notifier.ProcessPendingNotifications(), + this, + preferLocal: false)) + { + Interlocked.Exchange(ref _pendingNotificationCount, 0); + Volatile.Write(ref _notificationWorkerScheduled, 0); + _logger.LogError("Failed to queue an MCP tool-list change notification."); + } + } + + private void ProcessPendingNotifications() + { + try + { + while (Interlocked.Exchange(ref _pendingNotificationCount, 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-count + // exchange. Reschedule so that invalidation is never lost in that window. + if (Volatile.Read(ref _pendingNotificationCount) > 0) + { + ScheduleNotificationWorker(); + } + } } } } diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs index 81506094be..02f2591e6f 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs @@ -56,7 +56,7 @@ public void RegisterTool(IMcpTool tool) ImmutableDictionary tools = current.Tools.Add(toolName, tool); ImmutableArray advertisedTools = SortMetadata(current.AdvertisedTools.Add(metadata)); - string fingerprint = CreateDiscoveryFingerprint(advertisedTools); + string discoveryCanonicalJson = CreateDiscoveryCanonicalJson(advertisedTools); Interlocked.Exchange( ref _snapshot, @@ -64,7 +64,7 @@ public void RegisterTool(IMcpTool tool) Version: current.Version + 1, Tools: tools, AdvertisedTools: advertisedTools, - DiscoveryFingerprint: fingerprint)); + DiscoveryCanonicalJson: discoveryCanonicalJson)); } } @@ -119,7 +119,7 @@ internal static McpToolRegistryCandidate CreateCandidate( return new McpToolRegistryCandidate( Tools: toolBuilder.ToImmutable(), AdvertisedTools: advertisedTools, - DiscoveryFingerprint: CreateDiscoveryFingerprint(advertisedTools)); + DiscoveryCanonicalJson: CreateDiscoveryCanonicalJson(advertisedTools)); } /// @@ -136,15 +136,15 @@ internal McpToolRegistryUpdateResult PublishCandidate(McpToolRegistryCandidate c Version: current.Version + 1, Tools: candidate.Tools, AdvertisedTools: candidate.AdvertisedTools, - DiscoveryFingerprint: candidate.DiscoveryFingerprint); + DiscoveryCanonicalJson: candidate.DiscoveryCanonicalJson); Interlocked.Exchange(ref _snapshot, replacement); return new McpToolRegistryUpdateResult( Version: replacement.Version, DiscoveryChanged: !string.Equals( - current.DiscoveryFingerprint, - replacement.DiscoveryFingerprint, + current.DiscoveryCanonicalJson, + replacement.DiscoveryCanonicalJson, StringComparison.Ordinal), RegisteredToolCount: replacement.Tools.Count, AdvertisedToolCount: replacement.AdvertisedTools.Length); @@ -154,6 +154,11 @@ internal McpToolRegistryUpdateResult PublishCandidate(McpToolRegistryCandidate c /// /// Gets the metadata snapshot advertised by tools/list. /// + /// + /// Returns defensive deep clones so callers cannot mutate the private snapshot shared by + /// concurrent readers. The JSON round trip is intentional and occurs only for discovery + /// requests, not for tool lookup or execution. + /// public IReadOnlyList GetAdvertisedTools() { McpToolRegistrySnapshot snapshot = Volatile.Read(ref _snapshot); @@ -167,6 +172,9 @@ public IReadOnlyList GetAdvertisedTools() /// Retained for compatibility; MCP handlers should use so /// lookup and discovery come from the same registry generation. /// + [Obsolete( + "GetEnabledTools combines a registry snapshot with caller-supplied configuration. " + + "Use GetAdvertisedTools so discovery comes from one published generation.")] public IEnumerable GetEnabledTools(RuntimeConfig config) { ArgumentNullException.ThrowIfNull(config); @@ -231,7 +239,7 @@ private static ImmutableArray SortMetadata(IEnumerable metadata) .ToImmutableArray(); } - private static string CreateDiscoveryFingerprint(ImmutableArray metadata) + private static string CreateDiscoveryCanonicalJson(ImmutableArray metadata) { JsonElement serializedMetadata = JsonSerializer.SerializeToElement( metadata.ToArray(), @@ -304,13 +312,13 @@ private sealed record McpToolRegistrySnapshot( long Version, ImmutableDictionary Tools, ImmutableArray AdvertisedTools, - string DiscoveryFingerprint) + string DiscoveryCanonicalJson) { public static McpToolRegistrySnapshot Empty { get; } = new( Version: 0, Tools: ImmutableDictionary.Create(StringComparer.OrdinalIgnoreCase), AdvertisedTools: ImmutableArray.Empty, - DiscoveryFingerprint: "[]"); + DiscoveryCanonicalJson: "[]"); } } @@ -329,5 +337,5 @@ public readonly record struct McpToolRegistryUpdateResult( internal sealed record McpToolRegistryCandidate( ImmutableDictionary Tools, ImmutableArray AdvertisedTools, - string DiscoveryFingerprint); + string DiscoveryCanonicalJson); } diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs index e2465e0e19..f7990b09f7 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs @@ -67,7 +67,10 @@ public McpToolRegistryRefreshService( /// public void EnsureInitialized() { - RefreshRegistry(); + if (RefreshRegistry()) + { + NotifyToolsListChanged(); + } } /// @@ -93,7 +96,13 @@ private void OnConfigChanged(object? sender, HotReloadEventArgs args) { try { - RefreshRegistry(); + if (RefreshRegistry()) + { + // 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 (Exception ex) { @@ -104,21 +113,25 @@ private void OnConfigChanged(object? sender, HotReloadEventArgs args) } } - private void RefreshRegistry() + /// + /// when an initialized client should be notified after the writer + /// lock is released; otherwise . + /// + private bool RefreshRegistry() { lock (_refreshLock) { RuntimeConfig config = _runtimeConfigProvider.GetConfig(); if (ReferenceEquals(config, _lastAppliedConfig)) { - return; + return false; } - List customTools = CustomMcpToolFactory + List customTools = CustomMcpToolFactory .CreateCustomTools(config, _logger) .ToList(); - foreach (DynamicCustomTool customTool in customTools.Cast()) + foreach (DynamicCustomTool customTool in customTools) { bool initializedFromDatabase = customTool.InitializeMetadata( config, @@ -144,7 +157,7 @@ private void RefreshRegistry() _logger.LogWarning( "Discarded a stale MCP tool registry candidate because a newer runtime " + "configuration became active during the rebuild."); - return; + return false; } bool isInitialGeneration = _lastAppliedConfig is null; @@ -163,10 +176,7 @@ private void RefreshRegistry() result.AdvertisedToolCount, result.DiscoveryChanged); - if (!isInitialGeneration && result.DiscoveryChanged) - { - NotifyToolsListChanged(); - } + return !isInitialGeneration && result.DiscoveryChanged; } } @@ -194,7 +204,8 @@ private void NotifyToolsListChanged() public interface IMcpToolListChangedNotifier { /// - /// Notifies a connected, initialized client that it should refresh tools/list. + /// 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/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs index afee7b8cc0..5c18d951f7 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs @@ -222,6 +222,44 @@ public void HotReload_WhenNotifierThrows_PreservesPublicationAndContinuesNotifyi healthyNotifier.Verify(notifier => notifier.NotifyToolsListChanged(), Times.Once); } + [TestMethod] + public async Task HotReload_WhenNotifierBlocks_DoesNotRetainRefreshLock() + { + RuntimeConfig currentConfig = CreateRuntimeConfig(); + BlockingFirstNotifier blockingNotifier = new(); + Mock healthyNotifier = new(); + TestContext context = CreateContextWithNotifiers( + () => currentConfig, + healthyNotifier, + new IMcpToolListChangedNotifier[] { blockingNotifier, healthyNotifier.Object }); + context.Service.EnsureInitialized(); + + currentConfig = CreateRuntimeConfig(("FirstTool", "First generation")); + Task firstRefresh = Task.Run(() => RaiseRegistryChanged(context.HotReloadEventHandler)); + Task secondRefresh = Task.CompletedTask; + + try + { + Assert.IsTrue( + blockingNotifier.FirstCallEntered.Wait(TimeSpan.FromSeconds(5)), + "The first notification did not reach the blocking transport."); + + currentConfig = CreateRuntimeConfig(("LatestTool", "Latest generation")); + secondRefresh = Task.Run(() => RaiseRegistryChanged(context.HotReloadEventHandler)); + + Assert.IsTrue( + blockingNotifier.SecondCallEntered.Wait(TimeSpan.FromSeconds(5)), + "A blocked notifier must not retain the registry refresh writer lock."); + Assert.IsTrue(context.Registry.TryGetTool("latest_tool", out _)); + Assert.IsFalse(context.Registry.TryGetTool("first_tool", out _)); + } + finally + { + blockingNotifier.ReleaseFirstCall.Set(); + await Task.WhenAll(firstRefresh, secondRefresh).WaitAsync(TimeSpan.FromSeconds(5)); + } + } + [TestMethod] public void HotReload_AfterRejectedCandidate_RecoversOnNextConfig() { @@ -558,6 +596,33 @@ public void NotifyToolsListChanged() } } + private sealed class BlockingFirstNotifier : IMcpToolListChangedNotifier + { + private int _callCount; + + public ManualResetEventSlim FirstCallEntered { get; } = new(); + + public ManualResetEventSlim SecondCallEntered { get; } = new(); + + public ManualResetEventSlim ReleaseFirstCall { get; } = new(); + + public void NotifyToolsListChanged() + { + if (Interlocked.Increment(ref _callCount) == 1) + { + FirstCallEntered.Set(); + if (!ReleaseFirstCall.Wait(TimeSpan.FromSeconds(10))) + { + throw new TimeoutException("Timed out waiting to release the first notifier call."); + } + + return; + } + + SecondCallEntered.Set(); + } + } + private sealed class TestRuntimeConfigLoader : RuntimeConfigLoader { public TestRuntimeConfigLoader(HotReloadEventHandler handler) diff --git a/src/Service.Tests/Mcp/McpToolRegistryTests.cs b/src/Service.Tests/Mcp/McpToolRegistryTests.cs index 6a2c88b73c..c1f14c7165 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryTests.cs @@ -270,6 +270,7 @@ public void RegisterTool_WithLeadingTrailingWhitespace_ThrowsException() /// /// Parameterized test verifying GetEnabledTools returns only enabled tools. /// +#pragma warning disable CS0618 // Compatibility API behavior remains covered until removal. [DataTestMethod] [DataRow(1, 1, DisplayName = "Mixed: 1 enabled, 1 disabled → returns 1")] [DataRow(3, 0, DisplayName = "All enabled → returns all")] @@ -360,6 +361,7 @@ public void GetEnabledTools_MixedBuiltInAndCustomTools() Assert.IsFalse(enabledTools.Any(t => t.Name == "create_record")); Assert.IsFalse(enabledTools.Any(t => t.Name == "delete_record")); } +#pragma warning restore CS0618 /// /// Replacing the registry publishes a complete, deterministically ordered snapshot and diff --git a/src/Service.Tests/UnitTests/McpServerConfigurationTests.cs b/src/Service.Tests/UnitTests/McpServerConfigurationTests.cs new file mode 100644 index 0000000000..ac902192bf --- /dev/null +++ b/src/Service.Tests/UnitTests/McpServerConfigurationTests.cs @@ -0,0 +1,36 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#nullable enable + +using Azure.DataApiBuilder.Mcp.Core; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Options; +using Microsoft.VisualStudio.TestTools.UnitTesting; +using ModelContextProtocol.Server; + +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."); + } + } +} diff --git a/src/Service.Tests/UnitTests/McpStdioToolListChangedNotifierTests.cs b/src/Service.Tests/UnitTests/McpStdioToolListChangedNotifierTests.cs index f4266074e0..a5c732bf88 100644 --- a/src/Service.Tests/UnitTests/McpStdioToolListChangedNotifierTests.cs +++ b/src/Service.Tests/UnitTests/McpStdioToolListChangedNotifierTests.cs @@ -7,8 +7,12 @@ 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 { @@ -30,12 +34,15 @@ public void NotifyToolsListChanged_BeforeInitialized_DoesNotWrite() [TestMethod] public void NotifyToolsListChanged_AfterInitialized_WritesProtocolFrame() { - StringWriter output = new(); + 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, @@ -57,18 +64,148 @@ public void NotifyToolsListChanged_AfterInitialized_WritesProtocolFrame() [TestMethod] public void MarkInitialized_IsIdempotent() { - StringWriter output = new(); + 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_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 + { + public ManualResetEventSlim LineWritten { get; } = new(); + + public override void WriteLine(string? value) + { + base.WriteLine(value); + 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/Utilities/McpStdioHelper.cs b/src/Service/Utilities/McpStdioHelper.cs index 4c969acbf7..4a19ea7b6c 100644 --- a/src/Service/Utilities/McpStdioHelper.cs +++ b/src/Service/Utilities/McpStdioHelper.cs @@ -78,6 +78,10 @@ public static bool RunMcpStdioHost(IHost host) { try { + // 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. From 95f0ca6d48c77939d2934d95a63fc28fb6925f0f Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Fri, 31 Jul 2026 06:46:13 -0700 Subject: [PATCH 11/21] improve code quality, eliminate unused methods --- docs/design/McpToolRegistryHotReload.md | 73 +++-- .../Core/CustomMcpToolFactory.cs | 7 +- .../Core/McpServiceCollectionExtensions.cs | 5 +- .../Core/McpStdioServer.cs | 12 +- .../Core/McpToolRegistry.cs | 71 +---- .../Core/McpToolRegistryInitializer.cs | 65 ---- .../Core/McpToolRegistryRefreshService.cs | 19 +- ...tpToolRegistryHotReloadIntegrationTests.cs | 98 +++++- ...ioToolRegistryHotReloadIntegrationTests.cs | 214 ++++++++----- .../Mcp/McpToolRegistryInitializerTests.cs | 120 ------- .../Mcp/McpToolRegistryRefreshServiceTests.cs | 52 +++- src/Service.Tests/Mcp/McpToolRegistryTests.cs | 294 +----------------- .../UnitTests/McpServerConfigurationTests.cs | 163 ++++++++++ .../UnitTests/McpStdioServerRunAsyncTests.cs | 21 +- 14 files changed, 540 insertions(+), 674 deletions(-) delete mode 100644 src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs delete mode 100644 src/Service.Tests/Mcp/McpToolRegistryInitializerTests.cs diff --git a/docs/design/McpToolRegistryHotReload.md b/docs/design/McpToolRegistryHotReload.md index 177934140b..80a0bca13d 100644 --- a/docs/design/McpToolRegistryHotReload.md +++ b/docs/design/McpToolRegistryHotReload.md @@ -31,24 +31,24 @@ Consequently, a configuration hot-reload can leave MCP discovery stale in severa Built-in tool visibility already evaluates the current configuration during each `tools/list` request, but it is combined with a fixed startup registry. This avoids some stale built-in visibility, but does not solve stale custom tools or provide a single consistent registry generation. -## Current Implementation +## Previous Implementation ### Registry construction -[McpServiceCollectionExtensions.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpServiceCollectionExtensions.cs) currently: +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. -[McpToolRegistryInitializer.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs) resolves every `IMcpTool` and registers it once when the host starts. +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) stores tools in a mutable, case-insensitive `Dictionary`. It supports individual registration, lookup by name, and filtering enabled tools using a supplied `RuntimeConfig`. +[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 is safe under current startup-only mutation, but it cannot be modified concurrently with MCP requests. @@ -95,6 +95,7 @@ The MCP registry needs refreshed database metadata and must not publish newly ca 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 @@ -135,16 +136,18 @@ internal sealed record McpToolRegistrySnapshot( `Tools` contains: - Every built-in tool, including built-ins currently disabled by DML tool configuration. -- Every custom tool enabled in the configuration used to build the snapshot. +- Every independently DI-registered `IMcpTool` implementation. +- Every configuration-generated custom tool enabled in the configuration used to build the snapshot. `AdvertisedTools` contains precomputed metadata for tools whose `IsEnabled(config)` result was true for that same configuration generation. It is sorted deterministically by tool name. 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. Neither a tool retaining -its source metadata object nor a caller mutating a returned object can modify a published snapshot -or invalidate its canonical discovery representation. +candidate construction and again when returning public discovery results. Candidate publication +pre-serializes the complete canonical discovery representation, which the accessor deserializes to +produce caller-owned clones without reserializing every tool per request. 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 @@ -161,11 +164,11 @@ Readers capture the current snapshot once per operation: Readers do not acquire the rebuild lock. They observe either the complete previous snapshot or the complete replacement snapshot. -### 4. Built-in and custom tool lifetimes differ +### 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. +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. -Custom tools are removed from DI registration. They are configuration-generation objects and are recreated for every registry candidate. +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. @@ -176,8 +179,8 @@ A singleton `McpToolRegistryRefreshService` will coordinate initialization and h Its responsibilities are: 1. Capture the current `RuntimeConfig` generation. -2. Obtain the DI-owned built-in tools. -3. Create fresh custom tools from the captured configuration. +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. @@ -320,7 +323,8 @@ After a successful noninitial refresh, send: Notification rules: - Do not notify for initial registry construction. -- Do not notify before the client sends `notifications/initialized`. +- 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. @@ -332,9 +336,10 @@ Notification rules: 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` marks the notifier initialized when -handling `notifications/initialized`. The refresh service depends on zero or more tool-list -notifiers; HTTP mode has no notifier registered in this iteration. +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. ### 13. HTTP reads are immediately current, but HTTP push is deferred @@ -384,16 +389,12 @@ McpToolRegistryUpdateResult ReplaceAll( 7. Atomically publishes the candidate. 8. Returns the new version and whether discovery metadata changed. -The current public `RegisterTool` method is not used by the new production path. Because it is public, implementation should preserve it unless API review explicitly approves removal. If retained, it must use copy-on-write under the writer gate and must never mutate a published dictionary in place. - -The compatibility `GetEnabledTools(RuntimeConfig)` accessor is obsolete because combining a -published tool map with caller-supplied configuration can cross generations. Production discovery -must use `GetAdvertisedTools()`, whose visibility and metadata were captured with the snapshot. - -The obsolete `McpToolRegistryInitializer` compatibility fallback collects all tools and invokes -`ReplaceAll` with the current `RuntimeConfig`; it does not incrementally advertise disabled tools. -Manually assembled service providers using this fallback must register `RuntimeConfigProvider`, -because accurate snapshot discovery cannot be constructed without a configuration generation. +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. ## Detailed Flows @@ -569,12 +570,13 @@ 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. -- Custom tools: not registered in DI. +- Configuration-generated `DynamicCustomTool` instances: not registered in DI. - Stdio tool-list notifier: singleton, registered only in stdio mode. -The refresh service receives `IEnumerable` containing built-ins only. Reflection-based built-in discovery remains unchanged except for continuing to exclude `DynamicCustomTool`. +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. ## Anticipated Source Changes @@ -590,7 +592,7 @@ The exact file split may change during implementation, but the expected touchpoi ### MCP project - [McpToolRegistry.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs): immutable snapshots, bulk replacement, and atomic reads/publication. -- [McpToolRegistryInitializer.cs](../../src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs): replace with or evolve into the refresh service. +- 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. @@ -630,7 +632,7 @@ The exact file split may change during implementation, but the expected touchpoi ### Refresh-service unit tests -1. Initial construction uses DI-owned built-ins and newly created custom tools. +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. @@ -641,19 +643,22 @@ The exact file split may change during implementation, but the expected touchpoi 9. A stale candidate is discarded. 10. Repeated callbacks for an already successfully applied configuration do not publish duplicate generations unnecessarily. 11. Notifications occur only after a successful noninitial semantic discovery change. +12. 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 does not notify before `notifications/initialized`. +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. ### Hot-reload integration tests @@ -668,6 +673,8 @@ The exact file split may change during implementation, but the expected touchpoi 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. 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. @@ -736,7 +743,6 @@ Deferred because registry correctness does not require it and the current SDK ex | 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. | -| Public `RegisterTool` API conflicts with snapshot design | Preserve via safe copy-on-write unless API review approves removal. | ## Acceptance Criteria @@ -756,6 +762,7 @@ The implementation is complete when: 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. ## Follow-Up Work diff --git a/src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs b/src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs index 6f3cb306d2..2d23c6585c 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/CustomMcpToolFactory.cs @@ -47,11 +47,8 @@ public static IEnumerable CreateCustomTools(RuntimeConfig con } catch (Exception ex) { - logger?.LogError( - ex, - "Failed to create custom MCP tool for entity '{EntityName}'.", - 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/McpServiceCollectionExtensions.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpServiceCollectionExtensions.cs index 5bc5bcd278..e22490d110 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpServiceCollectionExtensions.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpServiceCollectionExtensions.cs @@ -40,8 +40,9 @@ public static IServiceCollection AddDabMcpServer(this IServiceCollection service services.AddSingleton(serviceProvider => serviceProvider.GetRequiredService()); - // Auto-discover and register built-in MCP tools. Custom tools are configuration- - // generation objects and are created by McpToolRegistryRefreshService. + // 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); // Configure MCP server and propagate runtime description to MCP initialize instructions. diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs index fcf34ca2ba..8ff80829ac 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs @@ -68,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 initializeResponseCompleted = false; while (!cancellationToken.IsCancellationRequested) { @@ -130,10 +131,19 @@ public async Task RunAsync(CancellationToken cancellationToken) { case "initialize": HandleInitialize(id, root); + initializeResponseCompleted = true; break; case "notifications/initialized": - _toolListChangedNotifier?.MarkInitialized(); + // 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 (initializeResponseCompleted) + { + _toolListChangedNotifier?.MarkInitialized(); + } + break; case "tools/list": diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs index 02f2591e6f..dd7cedc564 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs @@ -27,47 +27,6 @@ public class McpToolRegistry private readonly object _writerLock = new(); private McpToolRegistrySnapshot _snapshot = McpToolRegistrySnapshot.Empty; - /// - /// Registers a tool in the registry using copy-on-write publication. - /// This compatibility API is retained for callers that incrementally construct a registry; - /// production initialization and hot-reload use . - /// - /// Thrown when tool name is invalid or duplicate - public void RegisterTool(IMcpTool tool) - { - ArgumentNullException.ThrowIfNull(tool); - - Tool metadata = CloneMetadata(tool.GetToolMetadata()); - string toolName = ValidateToolName(metadata); - - lock (_writerLock) - { - McpToolRegistrySnapshot current = _snapshot; - - if (current.Tools.TryGetValue(toolName, out IMcpTool? existingTool)) - { - if (ReferenceEquals(existingTool, tool)) - { - return; - } - - throw CreateDuplicateToolException(toolName, existingTool, tool); - } - - ImmutableDictionary tools = current.Tools.Add(toolName, tool); - ImmutableArray advertisedTools = SortMetadata(current.AdvertisedTools.Add(metadata)); - string discoveryCanonicalJson = CreateDiscoveryCanonicalJson(advertisedTools); - - Interlocked.Exchange( - ref _snapshot, - new McpToolRegistrySnapshot( - Version: current.Version + 1, - Tools: tools, - AdvertisedTools: advertisedTools, - DiscoveryCanonicalJson: discoveryCanonicalJson)); - } - } - /// /// Replaces the complete registry with a snapshot built for . /// The candidate is validated and materialized before it is atomically published. @@ -156,32 +115,18 @@ internal McpToolRegistryUpdateResult PublishCandidate(McpToolRegistryCandidate c /// /// /// Returns defensive deep clones so callers cannot mutate the private snapshot shared by - /// concurrent readers. The JSON round trip is intentional and occurs only for discovery - /// requests, not for tool lookup or execution. + /// concurrent readers. Candidate construction already serializes the complete metadata + /// array for semantic comparison, so discovery deserializes that representation instead + /// of serializing every tool again on each request. /// public IReadOnlyList GetAdvertisedTools() { McpToolRegistrySnapshot snapshot = Volatile.Read(ref _snapshot); - return snapshot.AdvertisedTools - .Select(CloneMetadata) - .ToArray(); - } - - /// - /// Gets metadata for all registered tools that are enabled in the given runtime configuration. - /// Retained for compatibility; MCP handlers should use so - /// lookup and discovery come from the same registry generation. - /// - [Obsolete( - "GetEnabledTools combines a registry snapshot with caller-supplied configuration. " + - "Use GetAdvertisedTools so discovery comes from one published generation.")] - public IEnumerable GetEnabledTools(RuntimeConfig config) - { - ArgumentNullException.ThrowIfNull(config); - McpToolRegistrySnapshot snapshot = Volatile.Read(ref _snapshot); - return snapshot.Tools.Values - .Where(t => t.IsEnabled(config)) - .Select(t => CloneMetadata(t.GetToolMetadata())); + return JsonSerializer.Deserialize( + snapshot.DiscoveryCanonicalJson, + _discoveryJsonOptions) + ?? throw new InvalidOperationException( + "Failed to clone advertised MCP tool metadata."); } /// diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs deleted file mode 100644 index fa9741d139..0000000000 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryInitializer.cs +++ /dev/null @@ -1,65 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. - -using Azure.DataApiBuilder.Core.Configurations; -using Azure.DataApiBuilder.Mcp.Model; -using Microsoft.Extensions.DependencyInjection; -using Microsoft.Extensions.Hosting; - -namespace Azure.DataApiBuilder.Mcp.Core -{ - /// - /// Compatibility hosted service for callers that previously constructed the registry initializer - /// directly. DAB startup uses . - /// - [Obsolete($"Use {nameof(McpToolRegistryRefreshService)} instead.")] - 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) - { - cancellationToken.ThrowIfCancellationRequested(); - - IMcpToolRegistryRefreshService? refreshService = - _serviceProvider.GetService(); - if (refreshService is not null) - { - refreshService.EnsureInitialized(); - return Task.CompletedTask; - } - - // Preserve compatibility for manually assembled service providers without publishing - // disabled tools. Snapshot discovery requires the configuration used to evaluate each - // tool's visibility, so this fallback now requires RuntimeConfigProvider as well. - RuntimeConfigProvider runtimeConfigProvider = - _serviceProvider.GetService() - ?? throw new InvalidOperationException( - $"{nameof(RuntimeConfigProvider)} must be registered when using the legacy " + - $"{nameof(McpToolRegistryInitializer)} fallback."); - IMcpTool[] tools = _serviceProvider.GetServices().ToArray(); - foreach (DynamicCustomTool customTool in tools.OfType()) - { - customTool.InitializeMetadata(_serviceProvider); - } - - _toolRegistry.ReplaceAll(tools, runtimeConfigProvider.GetConfig()); - - 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 index f7990b09f7..4539f97e25 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs @@ -33,12 +33,14 @@ public sealed class McpToolRegistryRefreshService : IHostedService { private readonly RuntimeConfigProvider _runtimeConfigProvider; - private readonly IReadOnlyList _builtInTools; + 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( @@ -51,9 +53,10 @@ public McpToolRegistryRefreshService( HotReloadEventHandler? hotReloadEventHandler = null) { _runtimeConfigProvider = runtimeConfigProvider; - _builtInTools = tools - .Where(tool => tool.ToolType == ToolType.BuiltIn) - .ToArray(); + // 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(); @@ -149,7 +152,7 @@ private bool RefreshRegistry() } McpToolRegistryCandidate candidate = McpToolRegistry.CreateCandidate( - _builtInTools.Concat(customTools), + _registeredTools.Concat(customTools), config); if (!ReferenceEquals(config, _runtimeConfigProvider.GetConfig())) @@ -166,11 +169,13 @@ private bool RefreshRegistry() _logger.LogInformation( "Published MCP tool registry version {Version} with {BuiltInToolCount} " + - "built-in tools, {CustomToolCount} custom tools, {RegisteredToolCount} " + + "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, - _builtInTools.Count, + _registeredTools.Count(tool => tool.ToolType == ToolType.BuiltIn), + _registeredTools.Count(tool => tool.ToolType != ToolType.BuiltIn), customTools.Count, result.RegisteredToolCount, result.AdvertisedToolCount, diff --git a/src/Service.Tests/Mcp/McpHttpToolRegistryHotReloadIntegrationTests.cs b/src/Service.Tests/Mcp/McpHttpToolRegistryHotReloadIntegrationTests.cs index a2e17909d8..c4663a30ee 100644 --- a/src/Service.Tests/Mcp/McpHttpToolRegistryHotReloadIntegrationTests.cs +++ b/src/Service.Tests/Mcp/McpHttpToolRegistryHotReloadIntegrationTests.cs @@ -105,6 +105,57 @@ await SendMcpAsync( "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"))); @@ -166,12 +217,25 @@ await WriteConfigAsync( 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: "get_book_by_id", + Object: storedProcedure, Type: EntitySourceType.StoredProcedure, Parameters: null, KeyFields: null), @@ -208,7 +272,7 @@ private static RuntimeConfig CreateConfig( Mcp: new( Enabled: true, Path: MCP_PATH, - DmlTools: DmlToolsConfig.FromBoolean(false)), + DmlTools: DmlToolsConfig.FromBoolean(dmlToolsEnabled)), Host: new( Cors: null, Authentication: new( @@ -305,6 +369,36 @@ private static async Task WaitForToolSetAsync( 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, diff --git a/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs b/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs index 32b78012fb..8c6baeb0c1 100644 --- a/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs +++ b/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs @@ -12,6 +12,7 @@ using System.Threading; using System.Threading.Channels; using System.Threading.Tasks; +using System.IO.Abstractions; using Azure.DataApiBuilder.Config; using Azure.DataApiBuilder.Config.DatabasePrimitives; using Azure.DataApiBuilder.Config.ObjectModel; @@ -32,92 +33,122 @@ namespace Azure.DataApiBuilder.Service.Tests.Mcp public class McpStdioToolRegistryHotReloadIntegrationTests { [TestMethod] - public async Task InitializedClient_RefreshesRegistry_ReceivesNotificationAndUpdatedList() + public async Task InitializedClient_FileReload_EmitsOneNotificationAndReturnsUpdatedList() { - RuntimeConfig currentConfig = CreateRuntimeConfig(); - Mock configLoader = new(null, null); - 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); - - McpToolRegistry registry = new(); - HotReloadEventHandler hotReloadEventHandler = 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 + 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("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."); - - currentConfig = CreateRuntimeConfig(("GetBook", "Gets one book")); - hotReloadEventHandler.OnConfigChangedEvent( - hotReloadEventHandler, - new HotReloadEventArgs(MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED, string.Empty)); - - 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); - 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()); - await serverTask; + .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( @@ -128,6 +159,23 @@ private static async Task ReadJsonLineAsync( 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) { diff --git a/src/Service.Tests/Mcp/McpToolRegistryInitializerTests.cs b/src/Service.Tests/Mcp/McpToolRegistryInitializerTests.cs deleted file mode 100644 index 99ec895db4..0000000000 --- a/src/Service.Tests/Mcp/McpToolRegistryInitializerTests.cs +++ /dev/null @@ -1,120 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. - -using System; -using System.Collections.Generic; -using System.Linq; -using System.Text.Json; -using System.Threading; -using System.Threading.Tasks; -using Azure.DataApiBuilder.Config; -using Azure.DataApiBuilder.Config.ObjectModel; -using Azure.DataApiBuilder.Core.Configurations; -using Azure.DataApiBuilder.Mcp.Core; -using Azure.DataApiBuilder.Mcp.Model; -using Microsoft.Extensions.DependencyInjection; -using Microsoft.VisualStudio.TestTools.UnitTesting; -using ModelContextProtocol.Protocol; -using Moq; -using static Azure.DataApiBuilder.Mcp.Model.McpEnums; - -namespace Azure.DataApiBuilder.Service.Tests.Mcp -{ - [TestClass] - public class McpToolRegistryInitializerTests - { - [TestMethod] - public async Task LegacyFallback_UsesConfigAwareBulkReplacement() - { - RuntimeConfig config = CreateRuntimeConfig(); - Mock configLoader = new(null, null); - Mock configProvider = new(configLoader.Object); - configProvider.Setup(provider => provider.GetConfig()).Returns(config); - - ServiceCollection services = new(); - services.AddSingleton(configProvider.Object); - services.AddSingleton( - new TestMcpTool("enabled_tool", isEnabled: true)); - services.AddSingleton( - new TestMcpTool("disabled_tool", isEnabled: false)); - using ServiceProvider serviceProvider = services.BuildServiceProvider(); - McpToolRegistry registry = new(); -#pragma warning disable CS0618 // Explicitly exercises the documented compatibility fallback. - McpToolRegistryInitializer initializer = new(serviceProvider, registry); -#pragma warning restore CS0618 - - await initializer.StartAsync(CancellationToken.None); - - CollectionAssert.AreEqual( - new[] { "enabled_tool" }, - registry.GetAdvertisedTools().Select(tool => tool.Name).ToArray()); - Assert.IsTrue(registry.TryGetTool("enabled_tool", out _)); - Assert.IsTrue( - registry.TryGetTool("disabled_tool", out _), - "Disabled built-ins remain registered so execution can return a structured disabled response."); - } - - [TestMethod] - public async Task LegacyFallback_WithoutRuntimeConfigProvider_Throws() - { - ServiceCollection services = new(); - services.AddSingleton(new TestMcpTool("test_tool", isEnabled: true)); - using ServiceProvider serviceProvider = services.BuildServiceProvider(); -#pragma warning disable CS0618 // Explicitly exercises the documented compatibility fallback. - McpToolRegistryInitializer initializer = new(serviceProvider, new McpToolRegistry()); -#pragma warning restore CS0618 - - InvalidOperationException exception = await Assert.ThrowsExceptionAsync( - () => initializer.StartAsync(CancellationToken.None)); - - StringAssert.Contains(exception.Message, nameof(RuntimeConfigProvider)); - } - - 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 TestMcpTool : IMcpTool - { - private readonly string _name; - private readonly bool _isEnabled; - - public TestMcpTool(string name, bool isEnabled) - { - _name = name; - _isEnabled = isEnabled; - } - - public ToolType ToolType => ToolType.BuiltIn; - - public bool IsEnabled(RuntimeConfig config) => _isEnabled; - - public Tool GetToolMetadata() - { - return new Tool - { - Name = _name, - Description = "Test tool", - InputSchema = JsonSerializer.Deserialize("{\"type\":\"object\"}") - }; - } - - public Task ExecuteAsync( - JsonDocument? arguments, - IServiceProvider serviceProvider, - CancellationToken cancellationToken = default) - { - throw new NotImplementedException(); - } - } - } -} diff --git a/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs index 5c18d951f7..38c8078b14 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs @@ -17,6 +17,7 @@ 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; @@ -85,6 +86,54 @@ public void HotReload_AddsFreshCustomToolAndNotifiesClient() 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() { @@ -176,7 +225,8 @@ public void EnsureInitialized_WhenDatabaseMetadataUnavailable_PublishesConfigFal VerifyLogContains( context.Logger, LogLevel.Information, - "with 0 built-in tools, 1 custom tools, 1 registered tools, and 1 advertised tools. " + + "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."); } diff --git a/src/Service.Tests/Mcp/McpToolRegistryTests.cs b/src/Service.Tests/Mcp/McpToolRegistryTests.cs index c1f14c7165..a43b4825a4 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryTests.cs @@ -5,7 +5,6 @@ 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,166 +24,6 @@ namespace Azure.DataApiBuilder.Service.Tests.Mcp [TestClass] public class McpToolRegistryTests { - /// - /// Test that registering multiple tools with unique names succeeds. - /// - [TestMethod] - public void RegisterTool_WithMultipleUniqueNames_Succeeds() - { - // 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 _)); - } - - /// - /// 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. - /// - [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) - { - // Arrange - McpToolRegistry registry = new(); - IMcpTool tool1 = new MockMcpTool(toolName, toolType); - IMcpTool tool2 = new MockMcpTool(toolName, toolType); - - // Act - Register first tool - registry.RegisterTool(tool1); - - // Assert - Second registration should throw - DataApiBuilderException exception = Assert.ThrowsException( - () => registry.RegisterTool(tool2) - ); - - // 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.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). - /// - [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) - { - // Arrange - McpToolRegistry registry = new(); - IMcpTool existingTool = new MockMcpTool(toolName, firstToolType); - IMcpTool conflictingTool = new MockMcpTool(toolName, secondToolType); - - // Act - Register first tool - registry.RegisterTool(existingTool); - - // Assert - Second tool registration should throw - DataApiBuilderException exception = Assert.ThrowsException( - () => registry.RegisterTool(conflictingTool) - ); - - // 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); - } - - /// - /// Test that tool name comparison is case-sensitive. - /// Tools with different casing should not be allowed. - /// - [TestMethod] - public void RegisterTool_WithDifferentCasing_ThrowsException() - { - // Arrange - McpToolRegistry registry = new(); - IMcpTool tool1 = new MockMcpTool("my_tool", ToolType.BuiltIn); - IMcpTool tool2 = new MockMcpTool("My_Tool", ToolType.Custom); - - // Act - Register first tool - registry.RegisterTool(tool1); - - // Assert - Case-insensitive duplicate should throw - DataApiBuilderException exception = Assert.ThrowsException( - () => registry.RegisterTool(tool2) - ); - - Assert.IsTrue(exception.Message.Contains("Duplicate MCP tool name")); - Assert.AreEqual(DataApiBuilderException.SubStatusCodes.ErrorInInitialization, exception.SubStatusCode); - } - - /// - /// Test that registering the same tool instance twice is silently ignored (idempotent). - /// This preserves the compatibility registration API's idempotent behavior. - /// - [TestMethod] - public void RegisterTool_SameInstanceTwice_IsIdempotent() - { - // Arrange - McpToolRegistry registry = new(); - IMcpTool tool = new MockMcpTool("my_tool", ToolType.BuiltIn); - - // Act - Register the same instance twice - registry.RegisterTool(tool); - registry.RegisterTool(tool); - - // Assert - Tool should be registered only once - Assert.IsTrue(registry.TryGetTool("my_tool", out _)); - } - - /// - /// Test that registering a different instance with the same name throws an exception, - /// even though a same-instance re-registration would be allowed. - /// - [TestMethod] - public void RegisterTool_DifferentInstanceSameName_ThrowsException() - { - // 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); - - // Assert - Different instance with same name should throw - DataApiBuilderException exception = Assert.ThrowsException( - () => registry.RegisterTool(tool2) - ); - - Assert.IsTrue(exception.Message.Contains("Duplicate MCP tool name 'my_tool' detected")); - Assert.AreEqual(DataApiBuilderException.SubStatusCodes.ErrorInInitialization, exception.SubStatusCode); - } - /// /// Test that TryGetTool returns false for non-existent tool. /// @@ -193,7 +32,6 @@ public void TryGetTool_WithNonExistentName_ReturnsFalse() { // Arrange McpToolRegistry registry = new(); - registry.RegisterTool(new MockMcpTool("existing_tool", ToolType.BuiltIn)); // Act bool found = registry.TryGetTool("non_existent_tool", out IMcpTool? tool); @@ -207,7 +45,7 @@ public void TryGetTool_WithNonExistentName_ReturnsFalse() /// Test edge case: empty tool name should throw exception. /// [TestMethod] - public void RegisterTool_WithEmptyToolName_ThrowsException() + public void ReplaceAll_WithEmptyToolName_ThrowsException() { // Arrange McpToolRegistry registry = new(); @@ -215,154 +53,30 @@ public void RegisterTool_WithEmptyToolName_ThrowsException() // Assert - Empty tool names should be rejected DataApiBuilderException exception = Assert.ThrowsException( - () => registry.RegisterTool(tool) + () => registry.ReplaceAll(new[] { tool }, CreateRuntimeConfig()) ); Assert.IsTrue(exception.Message.Contains("cannot be null, empty, or whitespace")); Assert.AreEqual(DataApiBuilderException.SubStatusCodes.ErrorInInitialization, exception.SubStatusCode); } - /// - /// Test realistic scenario with actual built-in tool names. - /// - [TestMethod] - public void RegisterTool_WithRealisticBuiltInToolNames_DetectsDuplicates() - { - // Arrange - McpToolRegistry registry = new(); - - // 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) - ); - - Assert.IsTrue(exception.Message.Contains("read_records")); - Assert.IsTrue(exception.Message.Contains("built-in tool")); - } - /// /// Test that leading/trailing whitespace is rejected rather than producing a lookup key /// that differs from the advertised tool name. /// [TestMethod] - public void RegisterTool_WithLeadingTrailingWhitespace_ThrowsException() + public void ReplaceAll_WithLeadingTrailingWhitespace_ThrowsException() { McpToolRegistry registry = new(); IMcpTool tool = new MockMcpTool(" my_tool ", ToolType.Custom); DataApiBuilderException exception = Assert.ThrowsException( - () => registry.RegisterTool(tool)); + () => registry.ReplaceAll(new[] { tool }, CreateRuntimeConfig())); StringAssert.Contains(exception.Message, "leading or trailing whitespace"); Assert.IsFalse(registry.TryGetTool("my_tool", out _)); } - /// - /// Parameterized test verifying GetEnabledTools returns only enabled tools. - /// -#pragma warning disable CS0618 // Compatibility API behavior remains covered until removal. - [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) - { - // Arrange - McpToolRegistry registry = new(); - for (int i = 0; i < enabledCount; i++) - { - registry.RegisterTool(new MockMcpTool($"enabled_{i}", ToolType.BuiltIn, isEnabledFunc: _ => true)); - } - - for (int i = 0; i < disabledCount; i++) - { - registry.RegisterTool(new MockMcpTool($"disabled_{i}", ToolType.BuiltIn, isEnabledFunc: _ => false)); - } - - RuntimeConfig config = CreateRuntimeConfig(); - - // Act - List result = registry.GetEnabledTools(config).ToList(); - - // Assert - Assert.AreEqual(enabledCount, result.Count); - } - - /// - /// Test that GetEnabledTools passes the RuntimeConfig to IsEnabled so tools - /// can check DmlToolsConfig flags. - /// - [TestMethod] - public void GetEnabledTools_PassesConfigToIsEnabled() - { - // Arrange - McpToolRegistry registry = new(); - - // This tool checks config.McpDmlTools?.CreateRecord - IMcpTool configAwareTool = new MockMcpTool( - "create_record", ToolType.BuiltIn, - isEnabledFunc: config => config.McpDmlTools?.CreateRecord == true); - - registry.RegisterTool(configAwareTool); - - // Config with create-record disabled - DmlToolsConfig disabledConfig = new(createRecord: false); - RuntimeConfig configDisabled = CreateRuntimeConfig(disabledConfig); - - // Config with create-record enabled - DmlToolsConfig enabledConfig = new(createRecord: true); - RuntimeConfig configEnabled = CreateRuntimeConfig(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); - } - - /// - /// 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. - /// - [TestMethod] - public void GetEnabledTools_MixedBuiltInAndCustomTools() - { - // 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(); - - // 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")); - } -#pragma warning restore CS0618 - /// /// Replacing the registry publishes a complete, deterministically ordered snapshot and /// removes tools that belonged only to the previous generation. diff --git a/src/Service.Tests/UnitTests/McpServerConfigurationTests.cs b/src/Service.Tests/UnitTests/McpServerConfigurationTests.cs index ac902192bf..199279b001 100644 --- a/src/Service.Tests/UnitTests/McpServerConfigurationTests.cs +++ b/src/Service.Tests/UnitTests/McpServerConfigurationTests.cs @@ -3,11 +3,27 @@ #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 { @@ -32,5 +48,152 @@ public void ConfigureMcpServer_HttpDoesNotAdvertiseToolListChanges() 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/McpStdioServerRunAsyncTests.cs b/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs index 61328f4c0f..222b1f0ae4 100644 --- a/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs +++ b/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs @@ -60,12 +60,13 @@ public async Task RunAsync_BlankLineThenShutdown_IgnoresBlankLineAndHandlesShutd } [TestMethod] - public async Task RunAsync_InitializedNotification_MarksToolListNotifierReady() + 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\":1,\"method\":\"shutdown\"}" + Environment.NewLine; + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"shutdown\"}" + Environment.NewLine; (McpStdioServer server, _) = CreateServerWithCapturedOutput( new StringReader(input), notifier.Object); @@ -75,6 +76,22 @@ public async Task RunAsync_InitializedNotification_MarksToolListNotifierReady() 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) From 3c450ce3a87ba14cfbb74949dfcba8e192a91a50 Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Fri, 31 Jul 2026 07:33:30 -0700 Subject: [PATCH 12/21] cleanup code remove more unused code --- .../Core/DynamicCustomTool.cs | 30 ++------ .../Core/McpStdioServer.cs | 6 +- .../Core/McpStdioToolListChangedNotifier.cs | 14 ++-- .../Core/McpToolRegistry.cs | 20 +++--- .../Core/McpToolRegistryRefreshService.cs | 2 +- .../Utils/McpMetadataHelper.cs | 20 +++++- .../DynamicCustomToolMsSqlIntegrationTests.cs | 25 +++++-- .../Mcp/DynamicCustomToolTests.cs | 45 ++++++++---- ...tpToolRegistryHotReloadIntegrationTests.cs | 57 ++++++++++++++- .../Mcp/McpMetadataHelperTests.cs | 50 +++++++++++++ .../Mcp/McpToolRegistryRefreshServiceTests.cs | 72 ++++++++++--------- .../McpStdioServerInitializeTests.cs | 5 +- 12 files changed, 245 insertions(+), 101 deletions(-) create mode 100644 src/Service.Tests/Mcp/McpMetadataHelperTests.cs diff --git a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs index be73fd7cc9..906219e861 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs @@ -35,6 +35,7 @@ namespace Azure.DataApiBuilder.Mcp.Core public class DynamicCustomTool : IMcpTool { private readonly Entity _entity; + private readonly string _toolName; private JsonElement? _cachedInputSchema; /// @@ -46,6 +47,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) @@ -69,26 +71,9 @@ 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 to a config-based schema if DB metadata is unavailable. + /// Gets the normalized MCP tool name without materializing the complete metadata schema. /// - /// The application service provider with initialized metadata providers. - public void InitializeMetadata(IServiceProvider serviceProvider) - { - ArgumentNullException.ThrowIfNull(serviceProvider); - - RuntimeConfigProvider? configProvider = serviceProvider.GetService(); - IMetadataProviderFactory? metadataProviderFactory = serviceProvider.GetService(); - if (configProvider is null || metadataProviderFactory is null) - { - _cachedInputSchema = null; - return; - } - - _ = InitializeMetadata(configProvider.GetConfig(), metadataProviderFactory); - } + internal string ToolName => _toolName; /// /// Initializes the input schema using an explicit configuration and metadata-provider @@ -125,15 +110,14 @@ public bool InitializeMetadata( /// 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 }; @@ -148,7 +132,7 @@ public async Task ExecuteAsync( CancellationToken cancellationToken = default) { ILogger? logger = serviceProvider.GetService>(); - string toolName = GetToolMetadata().Name; + string toolName = _toolName; try { diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs index 8ff80829ac..a588650a5f 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs @@ -130,8 +130,7 @@ public async Task RunAsync(CancellationToken cancellationToken) switch (method) { case "initialize": - HandleInitialize(id, root); - initializeResponseCompleted = true; + initializeResponseCompleted = HandleInitialize(id, root); break; case "notifications/initialized": @@ -191,7 +190,7 @@ public async Task RunAsync(CancellationToken cancellationToken) /// server-supported version and client-requested version, and includes supported capabilities and server information. No notifications /// are sent here; the server waits for the client to send "notifications/initialized" before sending any notifications. /// - private void HandleInitialize(JsonElement? id, JsonElement root) + private bool HandleInitialize(JsonElement? id, JsonElement root) { string? clientRequestedProtocolVersion = GetClientProtocolVersion(root); string negotiatedProtocolVersion = @@ -273,6 +272,7 @@ private void HandleInitialize(JsonElement? id, JsonElement root) } WriteResult(id, result); + return true; } private static string? GetClientProtocolVersion(JsonElement root) diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs index 3af2fc2a72..f3a4a47c82 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs @@ -36,7 +36,7 @@ public sealed class McpStdioToolListChangedNotifier : IMcpStdioToolListChangedNo private readonly McpStdoutWriter _stdoutWriter; private readonly ILogger _logger; private int _isInitialized; - private int _pendingNotificationCount; + private int _notificationPending; private int _notificationWorkerScheduled; public McpStdioToolListChangedNotifier( @@ -61,7 +61,9 @@ public void NotifyToolsListChanged() return; } - Interlocked.Increment(ref _pendingNotificationCount); + // 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(); } @@ -77,7 +79,7 @@ private void ScheduleNotificationWorker() this, preferLocal: false)) { - Interlocked.Exchange(ref _pendingNotificationCount, 0); + Interlocked.Exchange(ref _notificationPending, 0); Volatile.Write(ref _notificationWorkerScheduled, 0); _logger.LogError("Failed to queue an MCP tool-list change notification."); } @@ -87,7 +89,7 @@ private void ProcessPendingNotifications() { try { - while (Interlocked.Exchange(ref _pendingNotificationCount, 0) > 0) + while (Interlocked.Exchange(ref _notificationPending, 0) != 0) { try { @@ -105,9 +107,9 @@ private void ProcessPendingNotifications() { Volatile.Write(ref _notificationWorkerScheduled, 0); - // A publication can race with worker shutdown after the final pending-count + // 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 _pendingNotificationCount) > 0) + if (Volatile.Read(ref _notificationPending) != 0) { ScheduleNotificationWorker(); } diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs index dd7cedc564..3654323aaf 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs @@ -31,7 +31,7 @@ public class McpToolRegistry /// Replaces the complete registry with a snapshot built for . /// The candidate is validated and materialized before it is atomically published. /// - public McpToolRegistryUpdateResult ReplaceAll(IEnumerable tools, RuntimeConfig config) + internal McpToolRegistryUpdateResult ReplaceAll(IEnumerable tools, RuntimeConfig config) { return PublishCandidate(CreateCandidate(tools, config)); } @@ -77,7 +77,7 @@ internal static McpToolRegistryCandidate CreateCandidate( ImmutableArray advertisedTools = SortMetadata(advertisedMetadata); return new McpToolRegistryCandidate( Tools: toolBuilder.ToImmutable(), - AdvertisedTools: advertisedTools, + AdvertisedToolCount: advertisedTools.Length, DiscoveryCanonicalJson: CreateDiscoveryCanonicalJson(advertisedTools)); } @@ -94,7 +94,7 @@ internal McpToolRegistryUpdateResult PublishCandidate(McpToolRegistryCandidate c McpToolRegistrySnapshot replacement = new( Version: current.Version + 1, Tools: candidate.Tools, - AdvertisedTools: candidate.AdvertisedTools, + AdvertisedToolCount: candidate.AdvertisedToolCount, DiscoveryCanonicalJson: candidate.DiscoveryCanonicalJson); Interlocked.Exchange(ref _snapshot, replacement); @@ -106,7 +106,7 @@ internal McpToolRegistryUpdateResult PublishCandidate(McpToolRegistryCandidate c replacement.DiscoveryCanonicalJson, StringComparison.Ordinal), RegisteredToolCount: replacement.Tools.Count, - AdvertisedToolCount: replacement.AdvertisedTools.Length); + AdvertisedToolCount: replacement.AdvertisedToolCount); } } @@ -117,7 +117,9 @@ internal McpToolRegistryUpdateResult PublishCandidate(McpToolRegistryCandidate c /// Returns defensive deep clones so callers cannot mutate the private snapshot shared by /// concurrent readers. Candidate construction already serializes the complete metadata /// array for semantic comparison, so discovery deserializes that representation instead - /// of serializing every tool again on each request. + /// of serializing every tool again on each request. This still allocates caller-owned + /// objects per request; canonical object-property ordering is not semantically observable + /// in JSON or JSON Schema. /// public IReadOnlyList GetAdvertisedTools() { @@ -256,13 +258,13 @@ private static void WriteCanonicalJson(Utf8JsonWriter writer, JsonElement elemen private sealed record McpToolRegistrySnapshot( long Version, ImmutableDictionary Tools, - ImmutableArray AdvertisedTools, + int AdvertisedToolCount, string DiscoveryCanonicalJson) { public static McpToolRegistrySnapshot Empty { get; } = new( Version: 0, Tools: ImmutableDictionary.Create(StringComparer.OrdinalIgnoreCase), - AdvertisedTools: ImmutableArray.Empty, + AdvertisedToolCount: 0, DiscoveryCanonicalJson: "[]"); } } @@ -270,7 +272,7 @@ private sealed record McpToolRegistrySnapshot( /// /// Describes the result of atomically replacing an MCP registry snapshot. /// - public readonly record struct McpToolRegistryUpdateResult( + internal readonly record struct McpToolRegistryUpdateResult( long Version, bool DiscoveryChanged, int RegisteredToolCount, @@ -281,6 +283,6 @@ public readonly record struct McpToolRegistryUpdateResult( /// internal sealed record McpToolRegistryCandidate( ImmutableDictionary Tools, - ImmutableArray AdvertisedTools, + int AdvertisedToolCount, string DiscoveryCanonicalJson); } diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs index 4539f97e25..9a42b879f6 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs @@ -145,7 +145,7 @@ private bool RefreshRegistry() _logger.LogWarning( "Using configuration-derived input schema for custom MCP tool " + "'{ToolName}' on entity '{EntityName}'. Reason: {FallbackReason}", - customTool.GetToolMetadata().Name, + customTool.ToolName, customTool.EntityName, fallbackReason); } diff --git a/src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs b/src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs index fe05f50446..6ad9323ff6 100644 --- a/src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs +++ b/src/Azure.DataApiBuilder.Mcp/Utils/McpMetadataHelper.cs @@ -60,9 +60,8 @@ 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; } @@ -109,6 +108,11 @@ public static bool TryResolveMetadata( dataSourceName = string.Empty; error = string.Empty; + if (!TryValidateEntityName(entityName, out error)) + { + return false; + } + // Resolve datasource name for the entity. try { @@ -158,5 +162,17 @@ public static bool TryResolveMetadata( 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/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 index c4663a30ee..e475f51252 100644 --- a/src/Service.Tests/Mcp/McpHttpToolRegistryHotReloadIntegrationTests.cs +++ b/src/Service.Tests/Mcp/McpHttpToolRegistryHotReloadIntegrationTests.cs @@ -14,8 +14,10 @@ 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 @@ -52,7 +54,11 @@ await WriteConfigAsync( $"--ConfigFileName={configPath}", "--no-https-redirect" }; - using TestServer server = new(Program.CreateWebHostBuilder(args)); + 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( @@ -193,7 +199,9 @@ await WriteConfigAsync( connectionString.ConnectionString, ("DuplicateTool", "First duplicate"), ("duplicate_tool", "Second duplicate"))); - await Task.Delay(500); + 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")); @@ -496,5 +504,50 @@ private static void AssertTool(JsonElement response, string name, string descrip } 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/McpMetadataHelperTests.cs b/src/Service.Tests/Mcp/McpMetadataHelperTests.cs new file mode 100644 index 0000000000..983fa50201 --- /dev/null +++ b/src/Service.Tests/Mcp/McpMetadataHelperTests.cs @@ -0,0 +1,50 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#nullable enable + +using System.Collections.Generic; +using Azure.DataApiBuilder.Config.DatabasePrimitives; +using Azure.DataApiBuilder.Config.ObjectModel; +using Azure.DataApiBuilder.Core.Services; +using Azure.DataApiBuilder.Core.Services.MetadataProviders; +using Azure.DataApiBuilder.Mcp.Utils; +using Microsoft.VisualStudio.TestTools.UnitTesting; +using Moq; + +namespace Azure.DataApiBuilder.Service.Tests.Mcp +{ + [TestClass] + public class McpMetadataHelperTests + { + [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); + } + } +} diff --git a/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs index 38c8078b14..fb63b62578 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs @@ -4,6 +4,7 @@ using System; using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; +using System.IO; using System.Linq; using System.Text.Json; using System.Threading; @@ -273,40 +274,47 @@ public void HotReload_WhenNotifierThrows_PreservesPublicationAndContinuesNotifyi } [TestMethod] - public async Task HotReload_WhenNotifierBlocks_DoesNotRetainRefreshLock() + public async Task HotReload_WhenNotifierBlocks_DoesNotBlockPublicationOrLaterHandlers() { RuntimeConfig currentConfig = CreateRuntimeConfig(); - BlockingFirstNotifier blockingNotifier = new(); - Mock healthyNotifier = new(); + BlockingStringWriter output = new(); + using McpStdoutWriter stdoutWriter = new(output); + McpStdioToolListChangedNotifier notifier = new(stdoutWriter); + notifier.MarkInitialized(); + ManualResetEventSlim laterHandlerCalled = new(); TestContext context = CreateContextWithNotifiers( () => currentConfig, - healthyNotifier, - new IMcpToolListChangedNotifier[] { blockingNotifier, healthyNotifier.Object }); + new Mock(), + new IMcpToolListChangedNotifier[] { notifier }); + context.HotReloadEventHandler.Subscribe( + GRAPHQL_SCHEMA_EVICTION_ON_CONFIG_CHANGED, + (_, _) => laterHandlerCalled.Set()); context.Service.EnsureInitialized(); currentConfig = CreateRuntimeConfig(("FirstTool", "First generation")); - Task firstRefresh = Task.Run(() => RaiseRegistryChanged(context.HotReloadEventHandler)); - Task secondRefresh = Task.CompletedTask; + TestRuntimeConfigLoader loader = new(context.HotReloadEventHandler) + { + RuntimeConfig = currentConfig + }; + Task firstRefresh = Task.Run(loader.RaiseConfigChanged); try { Assert.IsTrue( - blockingNotifier.FirstCallEntered.Wait(TimeSpan.FromSeconds(5)), - "The first notification did not reach the blocking transport."); - - currentConfig = CreateRuntimeConfig(("LatestTool", "Latest generation")); - secondRefresh = Task.Run(() => RaiseRegistryChanged(context.HotReloadEventHandler)); - + output.WriteEntered.Wait(TimeSpan.FromSeconds(5)), + "The stdio notification worker did not reach the blocking writer."); Assert.IsTrue( - blockingNotifier.SecondCallEntered.Wait(TimeSpan.FromSeconds(5)), - "A blocked notifier must not retain the registry refresh writer lock."); - Assert.IsTrue(context.Registry.TryGetTool("latest_tool", out _)); - Assert.IsFalse(context.Registry.TryGetTool("first_tool", out _)); + 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 { - blockingNotifier.ReleaseFirstCall.Set(); - await Task.WhenAll(firstRefresh, secondRefresh).WaitAsync(TimeSpan.FromSeconds(5)); + 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."); } } @@ -646,30 +654,24 @@ public void NotifyToolsListChanged() } } - private sealed class BlockingFirstNotifier : IMcpToolListChangedNotifier + private sealed class BlockingStringWriter : StringWriter { - private int _callCount; - - public ManualResetEventSlim FirstCallEntered { get; } = new(); + public ManualResetEventSlim WriteEntered { get; } = new(); - public ManualResetEventSlim SecondCallEntered { get; } = new(); + public ManualResetEventSlim ReleaseWrite { get; } = new(); - public ManualResetEventSlim ReleaseFirstCall { get; } = new(); + public ManualResetEventSlim LineWritten { get; } = new(); - public void NotifyToolsListChanged() + public override void WriteLine(string? value) { - if (Interlocked.Increment(ref _callCount) == 1) + WriteEntered.Set(); + if (!ReleaseWrite.Wait(TimeSpan.FromSeconds(10))) { - FirstCallEntered.Set(); - if (!ReleaseFirstCall.Wait(TimeSpan.FromSeconds(10))) - { - throw new TimeoutException("Timed out waiting to release the first notifier call."); - } - - return; + throw new TimeoutException("Timed out waiting to release the stdout write."); } - SecondCallEntered.Set(); + base.WriteLine(value); + LineWritten.Set(); } } diff --git a/src/Service.Tests/UnitTests/McpStdioServerInitializeTests.cs b/src/Service.Tests/UnitTests/McpStdioServerInitializeTests.cs index 53cff20188..214798916d 100644 --- a/src/Service.Tests/UnitTests/McpStdioServerInitializeTests.cs +++ b/src/Service.Tests/UnitTests/McpStdioServerInitializeTests.cs @@ -120,7 +120,10 @@ 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 }); + object? initializeCompleted = handleInitialize.Invoke( + server, + new object?[] { id, requestRoot }); + Assert.AreEqual(true, initializeCompleted); string output = ExtractSingleOutputLine(stdoutCapture); using JsonDocument response = JsonDocument.Parse(output); From 57d08c5db6d6e12eced811a3db28d026fc618397 Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Fri, 31 Jul 2026 08:44:44 -0700 Subject: [PATCH 13/21] fix(mcp): address hot-reload review findings --- docs/design/McpToolRegistryHotReload.md | 86 ++++++++++---- .../Core/DynamicCustomTool.cs | 6 + .../Core/McpStdioServer.cs | 18 +-- .../Core/McpStdioToolListChangedNotifier.cs | 59 ++++++++-- .../Core/McpToolRegistry.cs | 30 +++-- .../Core/McpToolRegistryRefreshService.cs | 12 +- .../Mcp/McpToolRegistryRefreshServiceTests.cs | 107 +++++++++++++++++- src/Service.Tests/Mcp/McpToolRegistryTests.cs | 28 +++++ .../McpStdioServerInitializeTests.cs | 56 +++++++-- .../McpStdioToolListChangedNotifierTests.cs | 61 ++++++++++ 10 files changed, 403 insertions(+), 60 deletions(-) diff --git a/docs/design/McpToolRegistryHotReload.md b/docs/design/McpToolRegistryHotReload.md index 80a0bca13d..0277af7f93 100644 --- a/docs/design/McpToolRegistryHotReload.md +++ b/docs/design/McpToolRegistryHotReload.md @@ -130,7 +130,9 @@ The registry snapshot will conceptually contain: internal sealed record McpToolRegistrySnapshot( long Version, ImmutableDictionary Tools, - ImmutableArray AdvertisedTools); + int AdvertisedToolCount, + string DiscoveryJson, + string DiscoveryCanonicalJson); ``` `Tools` contains: @@ -139,15 +141,21 @@ internal sealed record McpToolRegistrySnapshot( - Every independently DI-registered `IMcpTool` implementation. - Every configuration-generated custom tool enabled in the configuration used to build the snapshot. -`AdvertisedTools` contains precomputed metadata for tools whose `IsEnabled(config)` result was true for that same configuration generation. It is sorted deterministically by tool name. +`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 complete canonical discovery representation, which the accessor deserializes to -produce caller-owned clones without reserializing every tool per request. Neither a tool retaining -its source metadata object nor a caller mutating a returned object can modify a published snapshot. +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 @@ -159,7 +167,7 @@ A candidate snapshot is built completely before the live registry is changed. Th Readers capture the current snapshot once per operation: -- `tools/list` reads `AdvertisedTools` from one snapshot. +- `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. @@ -188,7 +196,7 @@ Its responsibilities are: 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. +`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 @@ -206,9 +214,7 @@ Database metadata unavailability is a deliberate exception to this strict behavi ### 7. Metadata initialization uses explicit dependencies -`DynamicCustomTool.InitializeMetadata(IServiceProvider)` is a service-locator pattern and allows metadata initialization to retrieve a `RuntimeConfig` different from the generation being built. - -The initialization path will instead receive explicit dependencies, conceptually: +The metadata initialization path receives explicit dependencies: ```csharp void InitializeMetadata( @@ -248,6 +254,12 @@ The refresh callback catches and logs hot-reload failures so an MCP candidate fa 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 @@ -291,10 +303,14 @@ The refresh service retains its own writer gate and stale-generation guard as de 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. 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 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. ### 11. Existing tool-call safety is preserved @@ -308,7 +324,9 @@ A request that resolved a tool immediately before a swap may finish with that to ### 12. Stdio sends tool-list change notifications -The stdio initialize response already advertises `tools.listChanged = true`, so the server must implement the corresponding notification. +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: @@ -340,6 +358,9 @@ client to request the latest complete snapshot. `McpStdioServer` tracks successf 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. ### 13. HTTP reads are immediately current, but HTTP push is deferred @@ -360,6 +381,10 @@ preserving array order. The comparison therefore ignores semantically irrelevant 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 @@ -373,7 +398,7 @@ IReadOnlyList GetAdvertisedTools(); bool TryGetTool(string toolName, out IMcpTool? tool); -McpToolRegistryUpdateResult ReplaceAll( +internal McpToolRegistryUpdateResult ReplaceAll( IEnumerable tools, RuntimeConfig config); ``` @@ -396,6 +421,21 @@ the refresh service builds and atomically publishes a complete configuration-awa 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()`, and the former +`McpToolRegistryInitializer`; bulk candidate construction/publication is now `internal`. This is an +intentional source/API-surface change, not an accidental compatibility omission. + +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. + ## Detailed Flows ### Initial HTTP startup @@ -629,6 +669,7 @@ The exact file split may change during implementation, but the expected touchpoi 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. ### Refresh-service unit tests @@ -641,9 +682,10 @@ The exact file split may change during implementation, but the expected touchpoi 7. Startup failure propagates. 8. Hot-reload failure is caught and logged. 9. A stale candidate is discarded. -10. Repeated callbacks for an already successfully applied configuration do not publish duplicate generations unnecessarily. -11. Notifications occur only after a successful noninitial semantic discovery change. -12. Independently DI-registered custom implementations remain published after refresh. +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 @@ -659,6 +701,11 @@ The exact file split may change during implementation, but the expected touchpoi 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. ### Hot-reload integration tests @@ -689,6 +736,7 @@ Use structured logs from the refresh service for: - 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. diff --git a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs index 906219e861..e60140cdac 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs @@ -63,6 +63,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; /// diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs index a588650a5f..4cd49bc48d 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs @@ -68,7 +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 initializeResponseCompleted = false; + bool initializeResponseWritten = false; while (!cancellationToken.IsCancellationRequested) { @@ -130,7 +130,9 @@ public async Task RunAsync(CancellationToken cancellationToken) switch (method) { case "initialize": - initializeResponseCompleted = HandleInitialize(id, root); + HandleInitialize(id, root); + // This assignment is reached only after WriteResult succeeds. + initializeResponseWritten = true; break; case "notifications/initialized": @@ -138,7 +140,7 @@ public async Task RunAsync(CancellationToken cancellationToken) // server successfully wrote its initialize response. Ignore an // out-of-order notification rather than enabling capabilities the // client has not negotiated. - if (initializeResponseCompleted) + if (initializeResponseWritten) { _toolListChangedNotifier?.MarkInitialized(); } @@ -190,11 +192,12 @@ public async Task RunAsync(CancellationToken cancellationToken) /// server-supported version and client-requested version, and includes supported capabilities and server information. No notifications /// are sent here; the server waits for the client to send "notifications/initialized" before sending any notifications. /// - private bool HandleInitialize(JsonElement? id, JsonElement root) + 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; @@ -224,7 +227,7 @@ private bool HandleInitialize(JsonElement? id, JsonElement root) protocolVersion = negotiatedProtocolVersion, capabilities = new { - tools = new { listChanged = true }, + tools = new { listChanged = supportsToolListChanged }, logging = new { } }, serverInfo = new @@ -242,7 +245,7 @@ private bool HandleInitialize(JsonElement? id, JsonElement root) protocolVersion = negotiatedProtocolVersion, capabilities = new { - tools = new { listChanged = true }, + tools = new { listChanged = supportsToolListChanged }, logging = new { } }, serverInfo = new @@ -260,7 +263,7 @@ private bool HandleInitialize(JsonElement? id, JsonElement root) protocolVersion = negotiatedProtocolVersion, capabilities = new { - tools = new { listChanged = true }, + tools = new { listChanged = supportsToolListChanged }, logging = new { } }, serverInfo = new @@ -272,7 +275,6 @@ private bool HandleInitialize(JsonElement? id, JsonElement root) } WriteResult(id, result); - return true; } private static string? GetClientProtocolVersion(JsonElement root) diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs index f3a4a47c82..dc71ed3f02 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioToolListChangedNotifier.cs @@ -35,6 +35,7 @@ public sealed class McpStdioToolListChangedNotifier : IMcpStdioToolListChangedNo private readonly McpStdoutWriter _stdoutWriter; private readonly ILogger _logger; + private readonly Func _tryScheduleWorker; private int _isInitialized; private int _notificationPending; private int _notificationWorkerScheduled; @@ -42,9 +43,19 @@ public sealed class McpStdioToolListChangedNotifier : IMcpStdioToolListChangedNo public McpStdioToolListChangedNotifier( McpStdoutWriter stdoutWriter, ILogger? logger = null) + : this(stdoutWriter, logger, TryScheduleOnThreadPool) { - _stdoutWriter = stdoutWriter; + } + + 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)); } /// @@ -74,14 +85,48 @@ private void ScheduleNotificationWorker() return; } - if (!ThreadPool.QueueUserWorkItem( - static notifier => notifier.ProcessPendingNotifications(), - this, - preferLocal: false)) + 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) { - Interlocked.Exchange(ref _notificationPending, 0); + // 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("Failed to queue an MCP tool-list change notification."); + _logger.LogError( + ex, + "Failed to start the MCP tool-list notification fallback worker. " + + "The notification remains pending."); } } diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs index 3654323aaf..f87a25bd85 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs @@ -75,10 +75,12 @@ internal static McpToolRegistryCandidate CreateCandidate( } ImmutableArray advertisedTools = SortMetadata(advertisedMetadata); + string discoveryJson = CreateDiscoveryJson(advertisedTools); return new McpToolRegistryCandidate( Tools: toolBuilder.ToImmutable(), AdvertisedToolCount: advertisedTools.Length, - DiscoveryCanonicalJson: CreateDiscoveryCanonicalJson(advertisedTools)); + DiscoveryJson: discoveryJson, + DiscoveryCanonicalJson: CreateDiscoveryCanonicalJson(discoveryJson)); } /// @@ -95,6 +97,7 @@ internal McpToolRegistryUpdateResult PublishCandidate(McpToolRegistryCandidate c Version: current.Version + 1, Tools: candidate.Tools, AdvertisedToolCount: candidate.AdvertisedToolCount, + DiscoveryJson: candidate.DiscoveryJson, DiscoveryCanonicalJson: candidate.DiscoveryCanonicalJson); Interlocked.Exchange(ref _snapshot, replacement); @@ -115,17 +118,16 @@ internal McpToolRegistryUpdateResult PublishCandidate(McpToolRegistryCandidate c /// /// /// Returns defensive deep clones so callers cannot mutate the private snapshot shared by - /// concurrent readers. Candidate construction already serializes the complete metadata - /// array for semantic comparison, so discovery deserializes that representation instead - /// of serializing every tool again on each request. This still allocates caller-owned - /// objects per request; canonical object-property ordering is not semantically observable - /// in JSON or JSON Schema. + /// 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.DiscoveryCanonicalJson, + snapshot.DiscoveryJson, _discoveryJsonOptions) ?? throw new InvalidOperationException( "Failed to clone advertised MCP tool metadata."); @@ -186,15 +188,20 @@ private static ImmutableArray SortMetadata(IEnumerable metadata) .ToImmutableArray(); } - private static string CreateDiscoveryCanonicalJson(ImmutableArray metadata) + private static string CreateDiscoveryJson(ImmutableArray metadata) { - JsonElement serializedMetadata = JsonSerializer.SerializeToElement( + 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); + WriteCanonicalJson(writer, serializedMetadata.RootElement); } return Encoding.UTF8.GetString(canonicalJson.ToArray()); @@ -259,12 +266,14 @@ 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: "[]"); } } @@ -284,5 +293,6 @@ internal readonly record struct McpToolRegistryUpdateResult( internal sealed record McpToolRegistryCandidate( ImmutableDictionary Tools, int AdvertisedToolCount, + string DiscoveryJson, string DiscoveryCanonicalJson); } diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs index 9a42b879f6..02783948fd 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs @@ -70,7 +70,7 @@ public McpToolRegistryRefreshService( /// public void EnsureInitialized() { - if (RefreshRegistry()) + if (RefreshRegistry(forceRebuildForCurrentConfig: false)) { NotifyToolsListChanged(); } @@ -99,7 +99,11 @@ private void OnConfigChanged(object? sender, HotReloadEventArgs args) { try { - if (RefreshRegistry()) + // 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)) { // Transport notification is deliberately outside _refreshLock. Implementations // must enqueue any potentially blocking I/O so the ordered reload pipeline can @@ -120,12 +124,12 @@ private void OnConfigChanged(object? sender, HotReloadEventArgs args) /// when an initialized client should be notified after the writer /// lock is released; otherwise . /// - private bool RefreshRegistry() + private bool RefreshRegistry(bool forceRebuildForCurrentConfig) { lock (_refreshLock) { RuntimeConfig config = _runtimeConfigProvider.GetConfig(); - if (ReferenceEquals(config, _lastAppliedConfig)) + if (!forceRebuildForCurrentConfig && ReferenceEquals(config, _lastAppliedConfig)) { return false; } diff --git a/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs index fb63b62578..fe375aae7d 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryRefreshServiceTests.cs @@ -48,6 +48,42 @@ public void EnsureInitialized_IsIdempotentForSameConfig() 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() { @@ -446,6 +482,34 @@ private static TestContext CreateContextWithNotifiers( 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); @@ -454,7 +518,7 @@ private static TestContext CreateContextWithNotifiers( Mock sqlMetadataProvider = new(); sqlMetadataProvider .SetupGet(provider => provider.EntityToDatabaseObject) - .Returns(new Dictionary()); + .Returns(getMetadata); Mock metadataProviderFactory = new(); metadataProviderFactory .Setup(factory => factory.GetMetadataProvider(It.IsAny())) @@ -465,7 +529,7 @@ private static TestContext CreateContextWithNotifiers( Mock> logger = new(); McpToolRegistryRefreshService service = new( configProvider.Object, - builtInTools, + registeredTools, registry, metadataProviderFactory.Object, notifiers, @@ -480,6 +544,45 @@ private static TestContext CreateContextWithNotifiers( 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) { diff --git a/src/Service.Tests/Mcp/McpToolRegistryTests.cs b/src/Service.Tests/Mcp/McpToolRegistryTests.cs index a43b4825a4..be7241e305 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryTests.cs @@ -207,6 +207,34 @@ public void ReplaceAll_WithEquivalentSchemaPropertyOrder_DoesNotReportDiscoveryC Assert.IsFalse(result.DiscoveryChanged); } + /// + /// 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 GetAdvertisedTools_PreservesInputSchemaPropertyOrder() + { + const string SCHEMA = + "{\"type\":\"object\",\"properties\":{" + + "\"second\":{\"type\":\"string\"}," + + "\"first\":{\"type\":\"integer\"}}}"; + McpToolRegistry registry = new(); + 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); + } + /// /// A real input-schema change remains client-visible after canonicalization. /// diff --git a/src/Service.Tests/UnitTests/McpStdioServerInitializeTests.cs b/src/Service.Tests/UnitTests/McpStdioServerInitializeTests.cs index 214798916d..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,17 +151,20 @@ private static JsonElement InvokeHandleInitialize(McpStdioServer server, StringW JsonElement requestRoot = request.RootElement; JsonElement? id = requestRoot.TryGetProperty("id", out JsonElement idElement) ? idElement : null; - object? initializeCompleted = handleInitialize.Invoke( + handleInitialize.Invoke( server, new object?[] { id, requestRoot }); - Assert.AreEqual(true, initializeCompleted); 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) @@ -146,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/McpStdioToolListChangedNotifierTests.cs b/src/Service.Tests/UnitTests/McpStdioToolListChangedNotifierTests.cs index a5c732bf88..fdd10c486f 100644 --- a/src/Service.Tests/UnitTests/McpStdioToolListChangedNotifierTests.cs +++ b/src/Service.Tests/UnitTests/McpStdioToolListChangedNotifierTests.cs @@ -110,6 +110,62 @@ await Task.WhenAny(notificationCall, Task.Delay(TimeSpan.FromSeconds(1))) == not "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() { @@ -142,11 +198,16 @@ public void NotifyToolsListChanged_WhenQueuedWriteFails_LogsError() 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(); } From 77436d66e8234dd0281d0810016bcf1a41d3f80a Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Fri, 31 Jul 2026 09:35:24 -0700 Subject: [PATCH 14/21] dont block on shutdown --- docs/design/McpToolRegistryHotReload.md | 9 + src/Config/ConfigFileWatcher.cs | 61 +++++- src/Config/FileSystemRuntimeConfigLoader.cs | 205 ++++++++++++++---- .../UnitTests/ConfigFileWatcherUnitTests.cs | 175 +++++++++++++++ 4 files changed, 398 insertions(+), 52 deletions(-) diff --git a/docs/design/McpToolRegistryHotReload.md b/docs/design/McpToolRegistryHotReload.md index 0277af7f93..2a830a4c4c 100644 --- a/docs/design/McpToolRegistryHotReload.md +++ b/docs/design/McpToolRegistryHotReload.md @@ -312,6 +312,13 @@ prevents an older, slower registry rebuild initiated outside the file-loader pip overwriting a newer registry generation. Neither mechanism provides transactional rollback after a handler failure; that remains separate work. +Loader shutdown does not wait to acquire this gate. `FileSystemRuntimeConfigLoader.Dispose()` first +atomically marks the loader disposed, detaches and disables its file watcher under a separate +watcher-lifecycle lock, and then schedules potentially blocking OS watcher resource disposal on a +background worker. An active reload may finish, but cancellation releases callbacks waiting to +enter the gate. This prevents host shutdown from waiting indefinitely when a reload is stalled in +external database metadata initialization. + ### 11. Existing tool-call safety is preserved After a successful swap: @@ -722,6 +729,8 @@ The exact file split may change during implementation, but the expected touchpoi 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. Disposing the file loader while a reload handler is blocked returns without waiting for the + pipeline, and callbacks queued on the loader gate exit after disposal begins. 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. 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/FileSystemRuntimeConfigLoader.cs b/src/Config/FileSystemRuntimeConfigLoader.cs index 005009538c..903ea0f081 100644 --- a/src/Config/FileSystemRuntimeConfigLoader.cs +++ b/src/Config/FileSystemRuntimeConfigLoader.cs @@ -33,7 +33,10 @@ namespace Azure.DataApiBuilder.Config; public class FileSystemRuntimeConfigLoader : RuntimeConfigLoader, IDisposable { private readonly SemaphoreSlim _hotReloadGate = new(initialCount: 1, maxCount: 1); - private bool _disposed; + private readonly CancellationTokenSource _disposeCancellation = new(); + private readonly object _watcherLock = new(); + private readonly Func _configFileWatcherFactory; + private int _disposed; /// /// This stores either the default config name e.g. dab-config.json /// or user provided config file which could be a relative file path, @@ -54,7 +57,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. @@ -97,9 +100,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; @@ -108,20 +136,21 @@ public FileSystemRuntimeConfigLoader( /// /// Disposes the config file watcher to release file handles and stop - /// monitoring the config file for changes. + /// monitoring the config file for changes. Disposal deliberately does not wait for an active + /// hot-reload pipeline, which can be blocked in external database metadata initialization. /// public void Dispose() { - ConfigFileWatcher? configFileWatcher; - _hotReloadGate.Wait(); - try + if (Interlocked.Exchange(ref _disposed, 1) != 0) { - if (_disposed) - { - return; - } + return; + } - _disposed = true; + _disposeCancellation.Cancel(); + + IConfigFileWatcher? configFileWatcher; + lock (_watcherLock) + { configFileWatcher = _configFileWatcher; _configFileWatcher = null; @@ -130,15 +159,24 @@ public void Dispose() configFileWatcher.NewFileContentsDetected -= OnNewFileContentsDetected; } } - finally + + if (configFileWatcher is not null) { - _hotReloadGate.Release(); - } + try + { + configFileWatcher.StopWatching(); + } + catch (Exception ex) + { + SendLogToBufferOrLogger( + LogLevel.Warning, + $"Unable to disable the configuration file watcher during shutdown due to {ex.Message}"); + } - // FileSystemWatcher disposal can block while an OS callback completes. Do not hold the - // reload gate during that external operation. Any callback already queued will observe - // _disposed after it acquires the gate and return without loading another generation. - configFileWatcher?.Dispose(); + // 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); + } } /// @@ -174,36 +212,43 @@ public string GetConfigFileName() /// private bool TrySetupConfigFileWatcher() { - // File watching / hot-reload isn't used for the CLI. - if (_isCliLoader) - { - return false; - } - - // If the file watcher is already set up, we don't need to do it again. - if (_configFileWatcher is not null) + lock (_watcherLock) { - return false; - } + // 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 (RuntimeConfig is not null) - { - try + // If the file watcher is already set up, we don't need to do it again. + if (_configFileWatcher is not null) { - _configFileWatcher = new(new FileSystemWatcherWrapper(_fileSystem), GetConfigDirectoryName(), GetConfigFileName()); - _configFileWatcher.NewFileContentsDetected += OnNewFileContentsDetected; + return false; } - catch (Exception ex) + + if (RuntimeConfig is not null) { - // 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}."); + 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 _configFileWatcher is not null; + return false; } - - return false; } /// @@ -230,10 +275,23 @@ internal void ProcessHotReloadNotification(Action? beforeEnteringGate = null) { beforeEnteringGate?.Invoke(); - _hotReloadGate.Wait(); + if (IsDisposed) + { + return; + } + try { - if (_disposed) + _hotReloadGate.Wait(_disposeCancellation.Token); + } + catch (OperationCanceledException) when (IsDisposed) + { + return; + } + + try + { + if (IsDisposed) { return; } @@ -269,10 +327,23 @@ public async Task ExecuteWithHotReloadSerializationAsync(Func operation) { ArgumentNullException.ThrowIfNull(operation); - await _hotReloadGate.WaitAsync().ConfigureAwait(false); + if (IsDisposed) + { + throw new ObjectDisposedException(nameof(FileSystemRuntimeConfigLoader)); + } + + try + { + await _hotReloadGate.WaitAsync(_disposeCancellation.Token).ConfigureAwait(false); + } + catch (OperationCanceledException) when (IsDisposed) + { + throw new ObjectDisposedException(nameof(FileSystemRuntimeConfigLoader)); + } + try { - if (_disposed) + if (IsDisposed) { throw new ObjectDisposedException(nameof(FileSystemRuntimeConfigLoader)); } @@ -285,6 +356,50 @@ public async Task ExecuteWithHotReloadSerializationAsync(Func operation) } } + private bool IsDisposed => Volatile.Read(ref _disposed) != 0; + + private void ScheduleConfigFileWatcherDisposal(IConfigFileWatcher configFileWatcher) + { + Action disposeWatcher = () => + { + try + { + configFileWatcher.Dispose(); + } + catch (Exception ex) + { + 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) + { + SendLogToBufferOrLogger( + LogLevel.Warning, + $"Unable to schedule configuration file watcher disposal due to {ex.Message}"); + } + } + /// /// Load the runtime config from the specified path. /// diff --git a/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs b/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs index 077c4b23a5..2a556df1e3 100644 --- a/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs +++ b/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs @@ -299,6 +299,181 @@ public async Task ConcurrentHotReloadNotifications_SerializeCompletePipelines() } } + [TestMethod] + public async Task Dispose_WhileHotReloadPipelineIsBlocked_DoesNotWaitForPipeline() + { + 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 releaseReloadHandler = new(); + using ManualResetEventSlim queuedReloadReachedGate = new(); + hotReloadEventHandler.Subscribe( + METADATA_PROVIDER_FACTORY_ON_CONFIG_CHANGED, + (_, _) => + { + reloadHandlerEntered.Set(); + if (!releaseReloadHandler.Wait(TimeSpan.FromSeconds(10))) + { + throw new TimeoutException("Timed out waiting to release the simulated metadata refresh."); + } + }); + + Task activeReload = Task.CompletedTask; + Task queuedReload = 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."); + + Task disposeTask = Task.Run(configLoader.Dispose); + Assert.AreSame( + disposeTask, + await Task.WhenAny(disposeTask, Task.Delay(TimeSpan.FromSeconds(1))), + "Dispose must not wait for a hot-reload pipeline blocked in external metadata work."); + Assert.IsTrue( + configFileWatcher.StopWatchingCalled.Wait(TimeSpan.FromSeconds(1)), + "Dispose must synchronously disable the watcher before returning."); + 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 disposal."); + Assert.IsFalse( + activeReload.IsCompleted, + "Disposal must not require the active external metadata operation to finish."); + + releaseReloadHandler.Set(); + await Task.WhenAll(activeReload, queuedReload).WaitAsync(TimeSpan.FromSeconds(5)); + Assert.AreEqual( + "/blocked", + configLoader.RuntimeConfig!.Runtime!.Rest!.Path, + "A callback queued before disposal must exit without loading another generation."); + } + finally + { + releaseReloadHandler.Set(); + configFileWatcher.ReleaseDispose.Set(); + configLoader.Dispose(); + await Task.WhenAll(activeReload, queuedReload).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); + } + } + + 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."; From a15cdb5990a316b257a710a9b62e36ede88ff495 Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Fri, 31 Jul 2026 11:24:45 -0700 Subject: [PATCH 15/21] add a proper shutdown mechanism --- .../Core/McpToolRegistryRefreshService.cs | 20 +- src/Config/FileSystemRuntimeConfigLoader.cs | 152 +++++++++++-- src/Config/HotReloadEventArgs.cs | 11 +- src/Config/Properties/AssemblyInfo.cs | 1 + src/Config/RuntimeConfigLoader.cs | 37 ++- src/Core/Resolvers/IQueryExecutor.cs | 24 +- src/Core/Resolvers/MsSqlQueryExecutor.cs | 26 ++- src/Core/Resolvers/MySqlQueryExecutor.cs | 15 +- src/Core/Resolvers/PostgreSqlExecutor.cs | 19 +- src/Core/Resolvers/QueryExecutor.cs | 131 ++++++++++- .../CosmosSqlMetadataProvider.cs | 6 + .../IMetadataProviderFactory.cs | 5 + .../MetadataProviders/ISqlMetadataProvider.cs | 5 + .../MetadataProviderFactory.cs | 12 +- .../MsSqlMetadataProvider.cs | 48 +++- .../MySqlMetadataProvider.cs | 14 +- .../MetadataProviders/SqlMetadataProvider.cs | 210 ++++++++++++++---- .../UnitTests/ConfigFileWatcherUnitTests.cs | 199 +++++++++++++++-- src/Service/Startup.cs | 7 + src/Service/Utilities/McpStdioHelper.cs | 7 + .../RuntimeConfigLoaderShutdownService.cs | 30 +++ .../Utilities/RuntimeInitializationHelper.cs | 8 +- 22 files changed, 848 insertions(+), 139 deletions(-) create mode 100644 src/Service/Utilities/RuntimeConfigLoaderShutdownService.cs diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs index 02783948fd..e2f9e5e724 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs @@ -70,7 +70,9 @@ public McpToolRegistryRefreshService( /// public void EnsureInitialized() { - if (RefreshRegistry(forceRebuildForCurrentConfig: false)) + if (RefreshRegistry( + forceRebuildForCurrentConfig: false, + CancellationToken.None)) { NotifyToolsListChanged(); } @@ -99,11 +101,14 @@ 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)) + 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 @@ -111,6 +116,10 @@ private void OnConfigChanged(object? sender, HotReloadEventArgs args) NotifyToolsListChanged(); } } + catch (OperationCanceledException) when (args.CancellationToken.IsCancellationRequested) + { + // Host shutdown canceled this generation before publication. + } catch (Exception ex) { _logger.LogError( @@ -124,10 +133,13 @@ private void OnConfigChanged(object? sender, HotReloadEventArgs args) /// when an initialized client should be notified after the writer /// lock is released; otherwise . /// - private bool RefreshRegistry(bool forceRebuildForCurrentConfig) + private bool RefreshRegistry( + bool forceRebuildForCurrentConfig, + CancellationToken cancellationToken) { lock (_refreshLock) { + cancellationToken.ThrowIfCancellationRequested(); RuntimeConfig config = _runtimeConfigProvider.GetConfig(); if (!forceRebuildForCurrentConfig && ReferenceEquals(config, _lastAppliedConfig)) { @@ -140,6 +152,7 @@ private bool RefreshRegistry(bool forceRebuildForCurrentConfig) foreach (DynamicCustomTool customTool in customTools) { + cancellationToken.ThrowIfCancellationRequested(); bool initializedFromDatabase = customTool.InitializeMetadata( config, _metadataProviderFactory, @@ -159,6 +172,7 @@ private bool RefreshRegistry(bool forceRebuildForCurrentConfig) _registeredTools.Concat(customTools), config); + cancellationToken.ThrowIfCancellationRequested(); if (!ReferenceEquals(config, _runtimeConfigProvider.GetConfig())) { _logger.LogWarning( diff --git a/src/Config/FileSystemRuntimeConfigLoader.cs b/src/Config/FileSystemRuntimeConfigLoader.cs index 903ea0f081..053871fc10 100644 --- a/src/Config/FileSystemRuntimeConfigLoader.cs +++ b/src/Config/FileSystemRuntimeConfigLoader.cs @@ -34,8 +34,11 @@ public class FileSystemRuntimeConfigLoader : RuntimeConfigLoader, IDisposable { 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 int _activeOperationCount; private int _disposed; /// /// This stores either the default config name e.g. dab-config.json @@ -136,17 +139,41 @@ internal FileSystemRuntimeConfigLoader( /// /// Disposes the config file watcher to release file handles and stop - /// monitoring the config file for changes. Disposal deliberately does not wait for an active - /// hot-reload pipeline, which can be blocked in external database metadata initialization. + /// monitoring the config file for changes. Active serialized work is canceled and drained + /// before this method returns so it cannot outlive host-owned dependencies. /// public void Dispose() { - if (Interlocked.Exchange(ref _disposed, 1) != 0) + StopAsync(CancellationToken.None).GetAwaiter().GetResult(); + } + + /// + /// 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. + /// + internal async Task StopAsync(CancellationToken cancellationToken) + { + Task activeOperationsDrained = BeginShutdown(); + await activeOperationsDrained.WaitAsync(cancellationToken).ConfigureAwait(false); + } + + private Task BeginShutdown() + { + bool firstShutdownRequest = Interlocked.Exchange(ref _disposed, 1) == 0; + if (firstShutdownRequest) { - return; + _disposeCancellation.Cancel(); + StopAndDisposeConfigFileWatcher(); } - _disposeCancellation.Cancel(); + lock (_operationLock) + { + return _activeOperationsDrained.Task; + } + } + + private void StopAndDisposeConfigFileWatcher() + { IConfigFileWatcher? configFileWatcher; lock (_watcherLock) @@ -289,20 +316,27 @@ internal void ProcessHotReloadNotification(Action? beforeEnteringGate = null) return; } - try + if (!TryBeginSerializedOperation()) { - if (IsDisposed) - { - return; - } + _hotReloadGate.Release(); + return; + } + try + { try { if (RuntimeConfig is not null) { - HotReloadConfig(RuntimeConfig.IsDevelopmentMode()); + HotReloadConfig( + RuntimeConfig.IsDevelopmentMode(), + _disposeCancellation.Token); } } + catch (OperationCanceledException) when (IsDisposed) + { + // Host shutdown canceled this generation before it could finish publication. + } catch (Exception ex) { SendLogToBufferOrLogger( @@ -312,6 +346,7 @@ internal void ProcessHotReloadNotification(Action? beforeEnteringGate = null) } finally { + EndSerializedOperation(); _hotReloadGate.Release(); } } @@ -323,7 +358,22 @@ internal void ProcessHotReloadNotification(Action? beforeEnteringGate = null) /// dependent component publication. /// /// The complete initial configuration operation to serialize. - public async Task ExecuteWithHotReloadSerializationAsync(Func operation) + public Task ExecuteWithHotReloadSerializationAsync(Func operation) + { + ArgumentNullException.ThrowIfNull(operation); + return ExecuteWithHotReloadSerializationAsync(_ => operation()); + } + + /// + /// Executes initial runtime dependency construction under the same per-loader gate used by + /// file-triggered hot reload, with cooperative shutdown cancellation. + /// + /// + /// 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); @@ -341,22 +391,65 @@ public async Task ExecuteWithHotReloadSerializationAsync(Func operation) throw new ObjectDisposedException(nameof(FileSystemRuntimeConfigLoader)); } + if (!TryBeginSerializedOperation()) + { + _hotReloadGate.Release(); + throw new ObjectDisposedException(nameof(FileSystemRuntimeConfigLoader)); + } + try + { + await operation(_disposeCancellation.Token).ConfigureAwait(false); + } + finally + { + EndSerializedOperation(); + _hotReloadGate.Release(); + } + } + + private bool IsDisposed => Volatile.Read(ref _disposed) != 0; + + private bool TryBeginSerializedOperation() + { + lock (_operationLock) { if (IsDisposed) { - throw new ObjectDisposedException(nameof(FileSystemRuntimeConfigLoader)); + return false; } - await operation().ConfigureAwait(false); + if (_activeOperationCount++ == 0) + { + _activeOperationsDrained = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + } + + return true; } - finally + } + + private void EndSerializedOperation() + { + TaskCompletionSource? drainedSignal = null; + lock (_operationLock) { - _hotReloadGate.Release(); + if (--_activeOperationCount == 0) + { + drainedSignal = _activeOperationsDrained; + } } + + drainedSignal?.TrySetResult(); } - private bool IsDisposed => Volatile.Read(ref _disposed) != 0; + private static TaskCompletionSource CreateCompletedDrainSignal() + { + TaskCompletionSource signal = new( + TaskCreationOptions.RunContinuationsAsynchronously); + signal.SetResult(); + return signal; + } private void ScheduleConfigFileWatcherDisposal(IConfigFileWatcher configFileWatcher) { @@ -414,8 +507,10 @@ public bool TryLoadConfig( [NotNullWhen(true)] out RuntimeConfig? config, ILogger? logger = null, bool? isDevMode = null, - DeserializationVariableReplacementSettings? replacementSettings = null) + DeserializationVariableReplacementSettings? replacementSettings = null, + CancellationToken cancellationToken = default) { + cancellationToken.ThrowIfCancellationRequested(); IsParseErrorEmitted = false; if (_fileSystem.File.Exists(path)) { @@ -431,6 +526,7 @@ public bool TryLoadConfig( string json = string.Empty; while (runCount <= FileUtilities.RunLimit) { + cancellationToken.ThrowIfCancellationRequested(); try { json = _fileSystem.File.ReadAllText(path); @@ -445,7 +541,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++; } } @@ -528,8 +630,9 @@ 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) + private void HotReloadConfig(bool isDevMode, CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); SendLogToBufferOrLogger( LogLevel.Information, $"Starting hot-reload process for config: {ConfigFilePath}"); @@ -537,7 +640,12 @@ private void HotReloadConfig(bool isDevMode) // Use default replacement settings for hot reload DeserializationVariableReplacementSettings replacementSettings = new(azureKeyVaultOptions: null, doReplaceEnvVar: true, doReplaceAkvVar: true); - if (!TryLoadConfig(ConfigFilePath, out _, isDevMode: isDevMode, replacementSettings: replacementSettings)) + if (!TryLoadConfig( + ConfigFilePath, + out _, + isDevMode: isDevMode, + replacementSettings: replacementSettings, + cancellationToken: cancellationToken)) { throw new DataApiBuilderException( message: "Deserialization of the configuration file failed.", @@ -547,7 +655,7 @@ private void HotReloadConfig(bool isDevMode) IsNewConfigDetected = true; IsNewConfigValidated = false; - SignalConfigChanged(); + SignalConfigChanged(cancellationToken: cancellationToken); SendLogToBufferOrLogger(LogLevel.Information, "Hot-reload process finished."); } diff --git a/src/Config/HotReloadEventArgs.cs b/src/Config/HotReloadEventArgs.cs index 5fa20e8d8d..2737709252 100644 --- a/src/Config/HotReloadEventArgs.cs +++ b/src/Config/HotReloadEventArgs.cs @@ -9,9 +9,18 @@ public class HotReloadEventArgs : EventArgs public string Message { get; set; } - public HotReloadEventArgs(string eventName, string message) + /// + /// Cancels the current ordered hot-reload generation during loader shutdown. + /// + public CancellationToken CancellationToken { get; } + + public HotReloadEventArgs( + string eventName, + string message, + CancellationToken cancellationToken = default) { EventName = eventName; Message = message; + CancellationToken = cancellationToken; } } diff --git a/src/Config/Properties/AssemblyInfo.cs b/src/Config/Properties/AssemblyInfo.cs index a2e622a838..3057d9a28b 100644 --- a/src/Config/Properties/AssemblyInfo.cs +++ b/src/Config/Properties/AssemblyInfo.cs @@ -4,3 +4,4 @@ using System.Runtime.CompilerServices; [assembly: InternalsVisibleTo("Azure.DataApiBuilder.Service.Tests")] +[assembly: InternalsVisibleTo("Azure.DataApiBuilder.Service")] diff --git a/src/Config/RuntimeConfigLoader.cs b/src/Config/RuntimeConfigLoader.cs index 2e05e208b4..698ff59c0b 100644 --- a/src/Config/RuntimeConfigLoader.cs +++ b/src/Config/RuntimeConfigLoader.cs @@ -85,38 +85,51 @@ protected virtual void OnConfigChangedEvent(HotReloadEventArgs args) /// been refreshed by previously called event triggers. /// /// - protected void SignalConfigChanged(string message = "") + protected void SignalConfigChanged( + string message = "", + CancellationToken cancellationToken = default) { + 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. - OnConfigChangedEvent(new HotReloadEventArgs(MCP_TOOL_REGISTRY_ON_CONFIG_CHANGED, message)); + 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/Resolvers/IQueryExecutor.cs b/src/Core/Resolvers/IQueryExecutor.cs index 2eac7242de..aacff00c34 100644 --- a/src/Core/Resolvers/IQueryExecutor.cs +++ b/src/Core/Resolvers/IQueryExecutor.cs @@ -34,6 +34,18 @@ public interface IQueryExecutor HttpContext? httpContext = null, List? args = null); + /// + /// Executes SQL text with cooperative cancellation. + /// + public Task ExecuteQueryAsync( + string sqltext, + IDictionary parameters, + Func?, Task>? dataReaderHandler, + string dataSourceName, + CancellationToken cancellationToken, + HttpContext? httpContext = null, + List? args = null); + /// /// Executes sql text with the given parameters and /// uses the function dataReaderHandler to process @@ -152,7 +164,17 @@ 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. + /// + public Task SetManagedIdentityAccessTokenIfAnyAsync( + DbConnection conn, + string dataSourceName, + CancellationToken cancellationToken); /// /// 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 670232b826..e174e2751d 100644 --- a/src/Core/Resolvers/MySqlQueryExecutor.cs +++ b/src/Core/Resolvers/MySqlQueryExecutor.cs @@ -106,8 +106,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)) { @@ -126,7 +130,7 @@ public override async Task SetManagedIdentityAccessTokenIfAnyAsync(DbConnection string? accessToken = accessTokenFromController ?? (IsDefaultAccessTokenValid() ? ((AccessToken)_defaultAccessToken!).Token : - await GetAccessTokenAsync()); + await GetAccessTokenAsync(cancellationToken)); if (accessToken is not null) { @@ -170,11 +174,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..3d9a86dfb3 100644 --- a/src/Core/Resolvers/QueryExecutor.cs +++ b/src/Core/Resolvers/QueryExecutor.cs @@ -172,6 +172,46 @@ 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.ThrowIfCancellationRequested(); int retryAttempt = 0; if (string.IsNullOrEmpty(dataSourceName)) @@ -190,12 +230,16 @@ public QueryExecutor(DbExceptionParser dbExceptionParser, DataApiBuilderException.SubStatusCodes.UnexpectedError); } - await SetManagedIdentityAccessTokenIfAnyAsync(conn, dataSourceName); + await SetManagedIdentityAccessTokenIfAnyAsync( + conn, + dataSourceName, + cancellationToken); TResult? result = default(TResult); - result = await _retryPolicyAsync.ExecuteAsync(async () => + result = await _retryPolicyAsync.ExecuteAsync(async retryCancellationToken => { + retryCancellationToken.ThrowIfCancellationRequested(); retryAttempt++; try { @@ -206,7 +250,24 @@ public QueryExecutor(DbExceptionParser dbExceptionParser, QueryExecutorLogger.LogDebug("{correlationId} Executing query: {queryText}", correlationId, sqltext); } - TResult? result = await ExecuteQueryAgainstDbAsync(conn, sqltext, parameters, dataReaderHandler, httpContext, dataSourceName, args); + 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 +297,7 @@ public QueryExecutor(DbExceptionParser dbExceptionParser, throw DbExceptionParser.Parse(e); } } - }); + }, cancellationToken); return result; } @@ -284,19 +345,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 +531,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/MetadataProviders/CosmosSqlMetadataProvider.cs b/src/Core/Services/MetadataProviders/CosmosSqlMetadataProvider.cs index b6d5dd0111..5374c5cf6c 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..beb970d90a 100644 --- a/src/Core/Services/MetadataProviders/IMetadataProviderFactory.cs +++ b/src/Core/Services/MetadataProviders/IMetadataProviderFactory.cs @@ -28,6 +28,11 @@ public interface IMetadataProviderFactory /// public Task InitializeAsync(); + /// + /// Initializes the metadata providers with cooperative cancellation. + /// + public Task InitializeAsync(CancellationToken cancellationToken); + /// /// 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 83989b645a..2a303c44ed 100644 --- a/src/Core/Services/MetadataProviders/ISqlMetadataProvider.cs +++ b/src/Core/Services/MetadataProviders/ISqlMetadataProvider.cs @@ -22,6 +22,11 @@ public interface ISqlMetadataProvider /// Task InitializeAsync(); + /// + /// Initializes this metadata provider for the runtime with cooperative cancellation. + /// + Task InitializeAsync(CancellationToken cancellationToken); + /// /// 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..64385894aa 100644 --- a/src/Core/Services/MetadataProviders/MsSqlMetadataProvider.cs +++ b/src/Core/Services/MetadataProviders/MsSqlMetadataProvider.cs @@ -60,7 +60,12 @@ public override Type SqlToCLRType(string sqlType) } /// - public override async Task PopulateTriggerMetadataForTable(string entityName, string schemaName, string tableName, SourceDefinition sourceDefinition) + public override async Task PopulateTriggerMetadataForTable( + string entityName, + string schemaName, + string tableName, + SourceDefinition sourceDefinition, + CancellationToken cancellationToken) { string enumerateEnabledTriggers = SqlQueryBuilder.BuildFetchEnabledTriggersQuery(); Dictionary parameters = new() @@ -73,7 +78,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 +164,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 +182,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 +197,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 +325,9 @@ private bool TryResolveDbType(string sqlDbTypeName, out DbType dbType) } /// - protected override async Task GenerateAutoentitiesIntoEntities(IReadOnlyDictionary? autoentities) + protected override async Task GenerateAutoentitiesIntoEntities( + IReadOnlyDictionary? autoentities, + CancellationToken cancellationToken) { if (autoentities is null) { @@ -321,8 +339,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 +452,10 @@ 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 async Task QueryAutoentitiesAsync( + string autoentityName, + Autoentity autoentity, + CancellationToken cancellationToken = default) { string include = string.Join(",", autoentity.Patterns.Include); string exclude = string.Join(",", autoentity.Patterns.Exclude); @@ -452,7 +477,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 afc62b3a38..625bf8e507 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. @@ -446,7 +459,8 @@ protected virtual async Task FillSchemaForStoredProcedureAsync( string entityName, string schemaName, string storedProcedureSourceName, - StoredProcedureDefinition storedProcedureDefinition) + StoredProcedureDefinition storedProcedureDefinition, + CancellationToken cancellationToken) { using ConnectionT conn = new(); conn.ConnectionString = ConnectionString; @@ -455,15 +469,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) { @@ -488,7 +512,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) @@ -552,7 +579,12 @@ 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, + CancellationToken cancellationToken) { throw new NotImplementedException(); } @@ -708,7 +740,9 @@ 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, + CancellationToken cancellationToken) { throw new NotSupportedException($"{GetType().Name} does not support autoentities yet."); } @@ -1177,21 +1211,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) { @@ -1199,10 +1246,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) { @@ -1211,14 +1262,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) @@ -1246,7 +1299,8 @@ await PopulateResultSetDefinitionsForStoredProcedureAsync( DataTable dataTable = await GetTableWithSchemaFromDataSetAsync( entityName, GetSchemaName(entityName), - GetDatabaseObjectName(entityName)); + GetDatabaseObjectName(entityName), + cancellationToken); pkFields = dataTable.PrimaryKey.Select(pk => pk.ColumnName).ToList(); } @@ -1259,7 +1313,8 @@ await PopulateSourceDefinitionAsync( GetSchemaName(entityName), GetDatabaseObjectName(entityName), GetSourceDefinition(entityName), - pkFields); + pkFields, + cancellationToken); } else { @@ -1286,7 +1341,8 @@ await PopulateSourceDefinitionAsync( DataTable dataTable = await GetTableWithSchemaFromDataSetAsync( entityName, GetSchemaName(entityName), - GetDatabaseObjectName(entityName)); + GetDatabaseObjectName(entityName), + cancellationToken); pkFields = dataTable.PrimaryKey.Select(pk => pk.ColumnName).ToList(); } @@ -1297,9 +1353,14 @@ await PopulateSourceDefinitionAsync( GetSchemaName(entityName), GetDatabaseObjectName(entityName), viewDefinition, - pkFields); + pkFields, + cancellationToken); } } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + throw; + } catch (Exception e) { HandleOrRecordException(e); @@ -1313,7 +1374,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}"; @@ -1327,7 +1389,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()); @@ -1487,8 +1550,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) @@ -1502,10 +1567,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(); @@ -1550,7 +1624,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, @@ -1560,7 +1637,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); } } @@ -1571,7 +1652,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); @@ -1587,7 +1672,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) { @@ -1652,7 +1738,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 @@ -1666,7 +1753,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) { @@ -1710,14 +1804,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) { @@ -1737,7 +1839,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 @@ -1761,7 +1864,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) { @@ -1774,7 +1884,7 @@ private async Task FillSchemaForTableAsync( innerException: ex); } - await conn.OpenAsync(); + await conn.OpenAsync(cancellationToken); DataAdapterT adapterForTable = new(); CommandT selectCommand = new() @@ -1787,7 +1897,9 @@ private async Task FillSchemaForTableAsync( = $"SELECT * FROM {tableNameWithSchemaPrefix}"; adapterForTable.SelectCommand = selectCommand; + cancellationToken.ThrowIfCancellationRequested(); DataTable[] dataTable = adapterForTable.FillSchema(EntitiesDataSet, SchemaType.Source, tableNameWithSchemaPrefix); + cancellationToken.ThrowIfCancellationRequested(); return dataTable[0]; } @@ -1826,12 +1938,16 @@ internal string GetTableNameWithSchemaPrefix(string schemaName, string tableName /// column of the table. 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. @@ -1845,7 +1961,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; } @@ -1880,8 +1999,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 @@ -1916,7 +2037,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/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs b/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs index 2a556df1e3..43db9d0351 100644 --- a/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs +++ b/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs @@ -12,6 +12,9 @@ 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; @@ -300,7 +303,7 @@ public async Task ConcurrentHotReloadNotifications_SerializeCompletePipelines() } [TestMethod] - public async Task Dispose_WhileHotReloadPipelineIsBlocked_DoesNotWaitForPipeline() + public async Task StopAsync_CancelsAndDrainsActiveReloadBeforeReturning() { string testDirectory = Path.Combine( Path.GetTempPath(), @@ -331,21 +334,33 @@ public async Task Dispose_WhileHotReloadPipelineIsBlocked_DoesNotWaitForPipeline 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( @@ -368,14 +383,13 @@ public async Task Dispose_WhileHotReloadPipelineIsBlocked_DoesNotWaitForPipeline queuedReloadReachedGate.Wait(TimeSpan.FromSeconds(5)), "The queued callback did not reach the serialization gate."); - Task disposeTask = Task.Run(configLoader.Dispose); - Assert.AreSame( - disposeTask, - await Task.WhenAny(disposeTask, Task.Delay(TimeSpan.FromSeconds(1))), - "Dispose must not wait for a hot-reload pipeline blocked in external metadata work."); + 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)), - "Dispose must synchronously disable the watcher before returning."); + "Shutdown must synchronously disable the watcher."); Assert.IsTrue( configFileWatcher.DisposeEntered.Wait(TimeSpan.FromSeconds(5)), "Watcher resource disposal was not scheduled."); @@ -385,24 +399,32 @@ await Task.WhenAny(disposeTask, Task.Delay(TimeSpan.FromSeconds(1))), Assert.AreSame( queuedReload, await Task.WhenAny(queuedReload, Task.Delay(TimeSpan.FromSeconds(1))), - "A callback waiting on the serialization gate must be canceled during disposal."); + "A callback waiting on the serialization gate must be canceled during shutdown."); Assert.IsFalse( - activeReload.IsCompleted, - "Disposal must not require the active external metadata operation to finish."); + stopTask.IsCompleted, + "Shutdown must drain the active reload before host-owned dependencies can be disposed."); + Assert.AreEqual( + 0, + Volatile.Read(ref laterHandlerInvocationCount), + "Cancellation must prevent later ordered handlers from running."); releaseReloadHandler.Set(); - await Task.WhenAll(activeReload, queuedReload).WaitAsync(TimeSpan.FromSeconds(5)); + await Task.WhenAll(activeReload, queuedReload, stopTask).WaitAsync(TimeSpan.FromSeconds(5)); + 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 disposal must exit without loading another generation."); + "A callback queued before shutdown must exit without loading another generation."); } finally { releaseReloadHandler.Set(); configFileWatcher.ReleaseDispose.Set(); - configLoader.Dispose(); - await Task.WhenAll(activeReload, queuedReload).WaitAsync(TimeSpan.FromSeconds(5)); + 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."); @@ -410,6 +432,153 @@ await Task.WhenAny(queuedReload, Task.Delay(TimeSpan.FromSeconds(1))), } } + [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); + } + } + + 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 diff --git a/src/Service/Startup.cs b/src/Service/Startup.cs index 36aa725f50..9e448bb5b0 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); @@ -605,6 +606,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()); } /// diff --git a/src/Service/Utilities/McpStdioHelper.cs b/src/Service/Utilities/McpStdioHelper.cs index 4a19ea7b6c..de81cdca4c 100644 --- a/src/Service/Utilities/McpStdioHelper.cs +++ b/src/Service/Utilities/McpStdioHelper.cs @@ -4,6 +4,8 @@ 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; @@ -101,6 +103,11 @@ public static bool RunMcpStdioHost(IHost host) } finally { + host.Services + .GetService()? + .StopAsync(CancellationToken.None) + .GetAwaiter() + .GetResult(); host.Dispose(); } } diff --git a/src/Service/Utilities/RuntimeConfigLoaderShutdownService.cs b/src/Service/Utilities/RuntimeConfigLoaderShutdownService.cs new file mode 100644 index 0000000000..9bcc4b9a21 --- /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's own shutdown token cancels supported external I/O. Do not detach an + // active operation on host timeout: doing so would let it use disposed singletons. + return configLoader.StopAsync(CancellationToken.None); + } + } +} \ No newline at end of file diff --git a/src/Service/Utilities/RuntimeInitializationHelper.cs b/src/Service/Utilities/RuntimeInitializationHelper.cs index 4833c3fe84..ce945e64df 100644 --- a/src/Service/Utilities/RuntimeInitializationHelper.cs +++ b/src/Service/Utilities/RuntimeInitializationHelper.cs @@ -32,8 +32,9 @@ public static async Task InitializeRuntimeDependenciesAsync( serviceProvider.GetRequiredService(); RuntimeConfig? initializedConfig = null; - await configLoader.ExecuteWithHotReloadSerializationAsync(async () => + await configLoader.ExecuteWithHotReloadSerializationAsync(async cancellationToken => { + cancellationToken.ThrowIfCancellationRequested(); RuntimeConfigProvider runtimeConfigProvider = serviceProvider.GetRequiredService(); initializedConfig = runtimeConfigProvider.GetConfig(); @@ -44,9 +45,12 @@ await configLoader.ExecuteWithHotReloadSerializationAsync(async () => IMetadataProviderFactory metadataProviderFactory = serviceProvider.GetRequiredService(); - await metadataProviderFactory.InitializeAsync().ConfigureAwait(false); + await metadataProviderFactory + .InitializeAsync(cancellationToken) + .ConfigureAwait(false); // MCP services are absent when MCP was disabled at startup. + cancellationToken.ThrowIfCancellationRequested(); IMcpToolRegistryRefreshService? mcpToolRegistryRefreshService = serviceProvider.GetService(); mcpToolRegistryRefreshService?.EnsureInitialized(); From 4fe5bf3cfaf5699b98ee3455053e8011c87a04f4 Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Sat, 1 Aug 2026 02:11:51 -0700 Subject: [PATCH 16/21] fix: finalize reload and query cancellation cleanup --- docs/design/McpToolRegistryHotReload.md | 17 +++ src/Config/FileSystemRuntimeConfigLoader.cs | 93 ++++++----- src/Core/Resolvers/QueryExecutor.cs | 23 ++- .../UnitTests/ConfigFileWatcherUnitTests.cs | 45 ++++++ .../UnitTests/SqlQueryExecutorUnitTests.cs | 144 ++++++++++++++++++ 5 files changed, 285 insertions(+), 37 deletions(-) diff --git a/docs/design/McpToolRegistryHotReload.md b/docs/design/McpToolRegistryHotReload.md index ce192e82ec..efdecb44b8 100644 --- a/docs/design/McpToolRegistryHotReload.md +++ b/docs/design/McpToolRegistryHotReload.md @@ -37,6 +37,7 @@ likely to look surprising when reviewing the implementation in isolation. | 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. | +| 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. | @@ -402,6 +403,15 @@ tracks callback completion rather than running callbacks inline on the host stop 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 @@ -563,6 +573,11 @@ The cancellation additions intentionally extend rather than replace established 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. @@ -1018,6 +1033,8 @@ stdout handle. | 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. | diff --git a/src/Config/FileSystemRuntimeConfigLoader.cs b/src/Config/FileSystemRuntimeConfigLoader.cs index 8c8aca3f1c..5c487ab4dc 100644 --- a/src/Config/FileSystemRuntimeConfigLoader.cs +++ b/src/Config/FileSystemRuntimeConfigLoader.cs @@ -38,9 +38,10 @@ public class FileSystemRuntimeConfigLoader : RuntimeConfigLoader, IDisposable private readonly object _watcherLock = new(); private readonly Func _configFileWatcherFactory; private TaskCompletionSource _activeOperationsDrained = CreateCompletedDrainSignal(); - private Task _cancellationCallbacksCompleted = Task.CompletedTask; + 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, @@ -154,10 +155,13 @@ public void Dispose() /// public async Task StopAsync(CancellationToken cancellationToken) { - Task activeOperationsDrained = BeginShutdown(); - await activeOperationsDrained.WaitAsync(cancellationToken).ConfigureAwait(false); + Task shutdownCompleted = BeginShutdown(); + await shutdownCompleted.WaitAsync(cancellationToken).ConfigureAwait(false); } + internal bool ShutdownResourcesDisposed => + Volatile.Read(ref _shutdownResourcesDisposed) != 0; + private Task BeginShutdown() { bool firstShutdownRequest; @@ -170,12 +174,14 @@ private Task BeginShutdown() { cancellationCallbacksCompleted = new TaskCompletionSource( TaskCreationOptions.RunContinuationsAsynchronously); - _cancellationCallbacksCompleted = cancellationCallbacksCompleted.Task; + Task drainCompleted = Task.WhenAll( + _activeOperationsDrained.Task, + cancellationCallbacksCompleted.Task); + _shutdownCompleted = DisposeSynchronizationResourcesAfterDrainAsync( + drainCompleted); } - shutdownCompleted = Task.WhenAll( - _activeOperationsDrained.Task, - _cancellationCallbacksCompleted); + shutdownCompleted = _shutdownCompleted; } if (firstShutdownRequest) @@ -189,6 +195,18 @@ private Task BeginShutdown() return shutdownCompleted; } + 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) { @@ -338,28 +356,24 @@ internal void ProcessHotReloadNotification(Action? beforeEnteringGate = null) { beforeEnteringGate?.Invoke(); - if (IsDisposed) - { - return; - } - - try - { - _hotReloadGate.Wait(_disposeCancellation.Token); - } - catch (OperationCanceledException) when (IsDisposed) - { - return; - } - if (!TryBeginSerializedOperation()) { - _hotReloadGate.Release(); return; } + bool gateEntered = false; try { + try + { + _hotReloadGate.Wait(_disposeCancellation.Token); + gateEntered = true; + } + catch (OperationCanceledException) when (IsDisposed) + { + return; + } + try { if (RuntimeConfig is not null) @@ -382,8 +396,14 @@ internal void ProcessHotReloadNotification(Action? beforeEnteringGate = null) } finally { + if (gateEntered) + { + // Release before completing the tracked operation. The final operation can make + // the shutdown continuation dispose the semaphore immediately. + _hotReloadGate.Release(); + } + EndSerializedOperation(); - _hotReloadGate.Release(); } } @@ -418,29 +438,34 @@ public async Task ExecuteWithHotReloadSerializationAsync( throw new ObjectDisposedException(nameof(FileSystemRuntimeConfigLoader)); } - try - { - await _hotReloadGate.WaitAsync(_disposeCancellation.Token).ConfigureAwait(false); - } - catch (OperationCanceledException) when (IsDisposed) - { - throw new ObjectDisposedException(nameof(FileSystemRuntimeConfigLoader)); - } - if (!TryBeginSerializedOperation()) { - _hotReloadGate.Release(); throw new ObjectDisposedException(nameof(FileSystemRuntimeConfigLoader)); } + bool gateEntered = false; try { + 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(); - _hotReloadGate.Release(); } } diff --git a/src/Core/Resolvers/QueryExecutor.cs b/src/Core/Resolvers/QueryExecutor.cs index 3d9a86dfb3..c9b9cb63b5 100644 --- a/src/Core/Resolvers/QueryExecutor.cs +++ b/src/Core/Resolvers/QueryExecutor.cs @@ -211,7 +211,20 @@ public QueryExecutor(DbExceptionParser dbExceptionParser, HttpContext? httpContext, List? args) { - cancellationToken.ThrowIfCancellationRequested(); + 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)) @@ -233,7 +246,7 @@ public QueryExecutor(DbExceptionParser dbExceptionParser, await SetManagedIdentityAccessTokenIfAnyAsync( conn, dataSourceName, - cancellationToken); + operationCancellationToken); TResult? result = default(TResult); @@ -250,6 +263,10 @@ await SetManagedIdentityAccessTokenIfAnyAsync( QueryExecutorLogger.LogDebug("{correlationId} Executing query: {queryText}", correlationId, sqltext); } + // 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, @@ -297,7 +314,7 @@ await SetManagedIdentityAccessTokenIfAnyAsync( throw DbExceptionParser.Parse(e); } } - }, cancellationToken); + }, operationCancellationToken); return result; } diff --git a/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs b/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs index d2b81f835e..15a7660cec 100644 --- a/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs +++ b/src/Service.Tests/UnitTests/ConfigFileWatcherUnitTests.cs @@ -403,6 +403,9 @@ await Task.WhenAny(queuedReload, Task.Delay(TimeSpan.FromSeconds(1))), 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), @@ -410,6 +413,9 @@ await Task.WhenAny(queuedReload, Task.Delay(TimeSpan.FromSeconds(1))), 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), @@ -432,6 +438,45 @@ await Task.WhenAny(queuedReload, Task.Delay(TimeSpan.FromSeconds(1))), } } + [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() { 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 From f34bb4850403c3ee64d73012ff69263751c8a234 Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Sat, 1 Aug 2026 02:22:20 -0700 Subject: [PATCH 17/21] docs: record metadata provider cancellation break --- docs/design/McpToolRegistryHotReload.md | 65 +++++++++++++++++-- .../MetadataProviders/SqlMetadataProvider.cs | 14 ++++ 2 files changed, 75 insertions(+), 4 deletions(-) diff --git a/docs/design/McpToolRegistryHotReload.md b/docs/design/McpToolRegistryHotReload.md index efdecb44b8..71f2e25cde 100644 --- a/docs/design/McpToolRegistryHotReload.md +++ b/docs/design/McpToolRegistryHotReload.md @@ -37,6 +37,7 @@ likely to look surprising when reviewing the implementation in isolation. | 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. | @@ -581,11 +582,12 @@ The cancellation additions intentionally extend rather than replace established - `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. -- Existing public and protected Config/Core members retain their original signatures and virtual - slots. `TryLoadConfig()`, `SignalConfigChanged()`, `PopulateTriggerMetadataForTable()`, +- 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. + 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 @@ -595,6 +597,50 @@ The cancellation additions intentionally extend rather than replace established 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 @@ -995,6 +1041,14 @@ second incremental construction path that can violate the atomic-generation inva 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. @@ -1040,6 +1094,7 @@ stdout handle. | 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 @@ -1063,7 +1118,7 @@ The implementation is complete when: 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. +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 @@ -1078,3 +1133,5 @@ The following work remains intentionally separate: 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/Core/Services/MetadataProviders/SqlMetadataProvider.cs b/src/Core/Services/MetadataProviders/SqlMetadataProvider.cs index a738d2a5ec..2895d6b89d 100644 --- a/src/Core/Services/MetadataProviders/SqlMetadataProvider.cs +++ b/src/Core/Services/MetadataProviders/SqlMetadataProvider.cs @@ -473,6 +473,13 @@ 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, @@ -2015,6 +2022,13 @@ 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, From 6408539d92ee8b8286cecbfc9d0f3fdb054b6950 Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Sat, 1 Aug 2026 03:19:01 -0700 Subject: [PATCH 18/21] fix(mcp): ignore schema set ordering in discovery --- docs/design/McpToolRegistryHotReload.md | 10 ++- .../Core/McpToolRegistry.cs | 65 +++++++++++++++-- src/Service.Tests/Mcp/McpToolRegistryTests.cs | 70 ++++++++++++++++++- 3 files changed, 137 insertions(+), 8 deletions(-) diff --git a/docs/design/McpToolRegistryHotReload.md b/docs/design/McpToolRegistryHotReload.md index 71f2e25cde..dd86c8b4cf 100644 --- a/docs/design/McpToolRegistryHotReload.md +++ b/docs/design/McpToolRegistryHotReload.md @@ -496,9 +496,11 @@ Every applicable configuration hot-reload rebuilds and publishes a generation so 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 -preserving array order. The comparison therefore ignores semantically irrelevant object insertion -order while still covering the complete tool metadata, including name, description, input schema, -and any future advertised fields. +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 @@ -880,6 +882,8 @@ The implementation is split across the following touchpoints: 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 diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs index 131384dbab..0e05e65d30 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistry.cs @@ -221,7 +221,11 @@ private static Tool CloneMetadata(Tool metadata) ?? throw new InvalidOperationException("Failed to clone MCP tool metadata."); } - private static void WriteCanonicalJson(Utf8JsonWriter writer, JsonElement element) + private static void WriteCanonicalJson( + Utf8JsonWriter writer, + JsonElement element, + string? propertyName = null, + bool isWithinJsonSchema = false) { switch (element.ValueKind) { @@ -232,7 +236,13 @@ private static void WriteCanonicalJson(Utf8JsonWriter writer, JsonElement elemen .OrderBy(property => property.Name, StringComparer.Ordinal)) { writer.WritePropertyName(property.Name); - WriteCanonicalJson(writer, property.Value); + WriteCanonicalJson( + writer, + property.Value, + property.Name, + isWithinJsonSchema || + property.NameEquals("inputSchema") || + property.NameEquals("outputSchema")); } writer.WriteEndObject(); @@ -240,9 +250,20 @@ private static void WriteCanonicalJson(Utf8JsonWriter writer, JsonElement elemen case JsonValueKind.Array: writer.WriteStartArray(); - foreach (JsonElement item in element.EnumerateArray()) + if (!TryWriteOrderInsensitiveJsonSchemaStringArray( + writer, + element, + propertyName, + isWithinJsonSchema)) { - WriteCanonicalJson(writer, item); + foreach (JsonElement item in element.EnumerateArray()) + { + WriteCanonicalJson( + writer, + item, + propertyName: null, + isWithinJsonSchema); + } } writer.WriteEndArray(); @@ -265,6 +286,42 @@ private static void WriteCanonicalJson(Utf8JsonWriter writer, JsonElement elemen } } + 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 (item.ValueKind != JsonValueKind.String) + { + return false; + } + + 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, diff --git a/src/Service.Tests/Mcp/McpToolRegistryTests.cs b/src/Service.Tests/Mcp/McpToolRegistryTests.cs index be7241e305..dcd3a0de0d 100644 --- a/src/Service.Tests/Mcp/McpToolRegistryTests.cs +++ b/src/Service.Tests/Mcp/McpToolRegistryTests.cs @@ -207,6 +207,63 @@ public void ReplaceAll_WithEquivalentSchemaPropertyOrder_DoesNotReportDiscoveryC Assert.IsFalse(result.DiscoveryChanged); } + /// + /// 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 ReplaceAll_WithEquivalentSchemaSetArrayOrder_DoesNotReportDiscoveryChange() + { + 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(); + RuntimeConfig config = CreateRuntimeConfig(); + registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, inputSchemaJson: SCHEMA_AB) }, + config); + + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, inputSchemaJson: SCHEMA_BA) }, + config); + + Assert.IsFalse(result.DiscoveryChanged); + } + + /// + /// 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 ReplaceAll_WithReorderedArrayDefault_ReportsDiscoveryChange() + { + 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); + + McpToolRegistryUpdateResult result = registry.ReplaceAll( + new[] { new MockMcpTool("same_tool", ToolType.Custom, inputSchemaJson: SCHEMA_BA) }, + config); + + Assert.IsTrue(result.DiscoveryChanged); + } + /// /// Canonical property sorting is used only for change detection. The discovery payload /// preserves schema-property insertion order for clients that render parameters in wire @@ -218,7 +275,8 @@ public void GetAdvertisedTools_PreservesInputSchemaPropertyOrder() const string SCHEMA = "{\"type\":\"object\",\"properties\":{" + "\"second\":{\"type\":\"string\"}," + - "\"first\":{\"type\":\"integer\"}}}"; + "\"first\":{\"type\":\"integer\"}}," + + "\"required\":[\"second\",\"first\"]}"; McpToolRegistry registry = new(); registry.ReplaceAll( new[] { new MockMcpTool("ordered_tool", ToolType.Custom, inputSchemaJson: SCHEMA) }, @@ -233,6 +291,16 @@ public void GetAdvertisedTools_PreservesInputSchemaPropertyOrder() .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); } /// From 06fbe87bff87a802e9cd1cc94d82a329b7ef5b81 Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Sat, 1 Aug 2026 03:35:43 -0700 Subject: [PATCH 19/21] style(mcp): resolve formatter diagnostics --- .../Core/DynamicCustomTool.cs | 11 +++++------ .../McpStdioToolRegistryHotReloadIntegrationTests.cs | 1 - 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs index 04d21ecccb..1d78e81680 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs @@ -35,7 +35,6 @@ namespace Azure.DataApiBuilder.Mcp.Core public class DynamicCustomTool : IMcpTool { private readonly Entity _entity; - private readonly string _toolName; private JsonElement? _cachedInputSchema; /// @@ -47,7 +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); + ToolName = ConvertToToolName(entityName); // Validate that this is a stored procedure if (_entity.Source.Type != EntitySourceType.StoredProcedure) @@ -79,7 +78,7 @@ public DynamicCustomTool(string entityName, Entity entity) /// /// Gets the normalized MCP tool name without materializing the complete metadata schema. /// - internal string ToolName => _toolName; + internal string ToolName { get; } /// /// Initializes the input schema using an explicit configuration and metadata-provider @@ -116,14 +115,14 @@ public bool InitializeMetadata( /// public Tool GetToolMetadata() { - 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 }; @@ -138,7 +137,7 @@ public async Task ExecuteAsync( CancellationToken cancellationToken = default) { ILogger? logger = serviceProvider.GetService>(); - string toolName = _toolName; + string toolName = ToolName; try { diff --git a/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs b/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs index 8c6baeb0c1..7b5a031c68 100644 --- a/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs +++ b/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs @@ -25,7 +25,6 @@ using Microsoft.Extensions.Logging.Abstractions; using Microsoft.VisualStudio.TestTools.UnitTesting; using Moq; -using static Azure.DataApiBuilder.Config.DabConfigEvents; namespace Azure.DataApiBuilder.Service.Tests.Mcp { From 662aec9e7e2711e72d478ed13cf70e991d54fb08 Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Sat, 1 Aug 2026 03:59:02 -0700 Subject: [PATCH 20/21] style: fix newline and import ordering --- .../Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs | 2 +- src/Service/Utilities/RuntimeConfigLoaderShutdownService.cs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs b/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs index 7b5a031c68..96607551a4 100644 --- a/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs +++ b/src/Service.Tests/Mcp/McpStdioToolRegistryHotReloadIntegrationTests.cs @@ -6,13 +6,13 @@ 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 System.IO.Abstractions; using Azure.DataApiBuilder.Config; using Azure.DataApiBuilder.Config.DatabasePrimitives; using Azure.DataApiBuilder.Config.ObjectModel; diff --git a/src/Service/Utilities/RuntimeConfigLoaderShutdownService.cs b/src/Service/Utilities/RuntimeConfigLoaderShutdownService.cs index 4f63b89485..1c533df119 100644 --- a/src/Service/Utilities/RuntimeConfigLoaderShutdownService.cs +++ b/src/Service/Utilities/RuntimeConfigLoaderShutdownService.cs @@ -27,4 +27,4 @@ public Task StopAsync(CancellationToken cancellationToken) return configLoader.StopAsync(cancellationToken); } } -} \ No newline at end of file +} From 08004e22b98cbb9d6fb8a41ed08aa214ecd962b4 Mon Sep 17 00:00:00 2001 From: Aaron Burtle Date: Sat, 1 Aug 2026 05:35:47 -0700 Subject: [PATCH 21/21] fix(tests): update cancellation and reload assertions --- .../HotReload/ConfigurationHotReloadTests.cs | 112 ++++++++++++++---- .../UnitTests/SqlMetadataProviderUnitTests.cs | 6 +- 2 files changed, 97 insertions(+), 21 deletions(-) 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/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();