Merge branch 'main' into dev/aaronburtle/Fix-Pagination-Column-Response

Author: aaronburtle

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

@@ -8,6 +8,7 @@
 using Azure.DataApiBuilder.Core.Configurations;
 using Azure.DataApiBuilder.Core.Telemetry;
 using Azure.DataApiBuilder.Mcp.Model;
+using Azure.DataApiBuilder.Mcp.Telemetry;
 using Azure.DataApiBuilder.Mcp.Utils;
 using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.Configuration;
@@ -26,6 +27,7 @@ public class McpStdioServer : IMcpStdioServer
     {
         private readonly McpToolRegistry _toolRegistry;
         private readonly IServiceProvider _serviceProvider;
+        private readonly McpStdoutWriter _stdoutWriter;
         private readonly string _protocolVersion;
 
         private const int MAX_LINE_LENGTH = 1024 * 1024; // 1 MB limit for incoming JSON-RPC requests
@@ -35,6 +37,11 @@ public McpStdioServer(McpToolRegistry toolRegistry, IServiceProvider serviceProv
             _toolRegistry = toolRegistry ?? throw new ArgumentNullException(nameof(toolRegistry));
             _serviceProvider = serviceProvider ?? throw new ArgumentNullException(nameof(serviceProvider));
 
+            // Resolve the shared stdout writer so JSON-RPC responses and
+            // notifications/message frames are serialized through one lock.
+            // Falls back to a fresh instance if DI didn't register one (defensive).
+            _stdoutWriter = _serviceProvider.GetService<McpStdoutWriter>() ?? new McpStdoutWriter();
+
             // Allow protocol version to be configured via IConfiguration, using centralized defaults.
             IConfiguration? configuration = _serviceProvider.GetService<IConfiguration>();
             _protocolVersion = McpProtocolDefaults.ResolveProtocolVersion(configuration);
@@ -47,16 +54,14 @@ public McpStdioServer(McpToolRegistry toolRegistry, IServiceProvider serviceProv
         /// <returns>A task representing the asynchronous operation.</returns>
         public async Task RunAsync(CancellationToken cancellationToken)
         {
-            // Use UTF-8 WITHOUT BOM
+            // Use UTF-8 WITHOUT BOM for stdin. Stdout is owned by McpStdoutWriter,
+            // which serializes all writes from McpStdioServer and the MCP logging
+            // pipeline so JSON-RPC frames cannot interleave at the byte level.
             UTF8Encoding utf8NoBom = new(encoderShouldEmitUTF8Identifier: false);
 
             using Stream stdin = Console.OpenStandardInput();
-            using Stream stdout = Console.OpenStandardOutput();
             using StreamReader reader = new(stdin, utf8NoBom);
-            using StreamWriter writer = new(stdout, utf8NoBom) { AutoFlush = true };
 
-            // Redirect Console.Out to use our writer
-            Console.SetOut(writer);
             while (!cancellationToken.IsCancellationRequested)
             {
                 string? line = await reader.ReadLineAsync(cancellationToken);
@@ -298,16 +303,31 @@ private void HandleSetLogLevel(JsonElement? id, JsonElement root)
             // If CLI or Config overrode, this returns false but we still return success to the client
             bool updated = logLevelController.UpdateFromMcp(level);
 
-            // If MCP successfully changed the log level to something other than "none",
-            // ensure Console.Error is pointing to the real stderr (not TextWriter.Null).
-            // This handles the case where MCP stdio mode started with LogLevel.None (quiet startup)
-            // and the client later enables logging via logging/setLevel.
+            // Determine if logging is enabled (level != "none")
+            // Note: Even if CLI/Config overrode the level, we still enable notifications
+            // when the client requests logging. They'll get logs at the overridden level.
             bool isLoggingEnabled = !string.Equals(level, "none", StringComparison.OrdinalIgnoreCase);
+
+            // Only restore stderr when this MCP call actually changed the effective level.
+            // If CLI/Config overrode (updated == false), stderr is already in the correct state:
+            //   - CLI/Config level == "none": stderr was redirected to TextWriter.Null at startup
+            //     and must stay that way; restoring it would re-introduce noisy output even
+            //     though the operator explicitly asked for silence.
+            //   - CLI/Config level != "none": stderr was never redirected, so restoring is a no-op.
             if (updated && isLoggingEnabled)
             {
                 RestoreStderrIfNeeded();
             }
 
+            // Enable or disable MCP log notifications based on the requested level
+            // When CLI/Config overrode, notifications are still enabled - client asked for logs,
+            // they just get them at the CLI/Config level instead of the requested level.
+            IMcpLogNotificationWriter? notificationWriter = _serviceProvider.GetService<IMcpLogNotificationWriter>();
+            if (notificationWriter != null)
+            {
+                notificationWriter.IsEnabled = isLoggingEnabled;
+            }
+
             // Always return success (empty result object) per MCP spec
             WriteResult(id, new { });
         }
@@ -539,39 +559,41 @@ private static string SafeToString(object obj)
 
         /// <summary>
         /// Writes a JSON-RPC result response to the standard output.
+        /// Routed through <see cref="McpStdoutWriter"/> so the write is serialized
+        /// with notifications/message frames from the logging pipeline.
         /// </summary>
         /// <param name="id">The request identifier extracted from the incoming JSON-RPC request. Used to correlate the response with the request.</param>
         /// <param name="resultObject">The result object to include in the response.</param>
-        private static void WriteResult(JsonElement? id, object resultObject)
+        private void WriteResult(JsonElement? id, object resultObject)
         {
             var response = new
             {
-                jsonrpc = "2.0",
+                jsonrpc = McpStdioJsonRpcErrorCodes.JSON_RPC_VERSION,
                 id = id.HasValue ? GetIdValue(id.Value) : null,
                 result = resultObject
             };
 
-            string json = JsonSerializer.Serialize(response);
-            Console.Out.WriteLine(json);
+            _stdoutWriter.WriteLine(JsonSerializer.Serialize(response));
         }
 
         /// <summary>
         /// Writes a JSON-RPC error response to the standard output.
+        /// Routed through <see cref="McpStdoutWriter"/> so the write is serialized
+        /// with notifications/message frames from the logging pipeline.
         /// </summary>
         /// <param name="id">The request identifier extracted from the incoming JSON-RPC request. Used to correlate the response with the request.</param>
         /// <param name="code">The error code.</param>
         /// <param name="message">The error message.</param>
-        private static void WriteError(JsonElement? id, int code, string message)
+        private void WriteError(JsonElement? id, int code, string message)
         {
             var errorObj = new
             {
-                jsonrpc = "2.0",
+                jsonrpc = McpStdioJsonRpcErrorCodes.JSON_RPC_VERSION,
                 id = id.HasValue ? GetIdValue(id.Value) : null,
                 error = new { code, message }
             };
 
-            string json = JsonSerializer.Serialize(errorObj);
-            Console.Out.WriteLine(json);
+            _stdoutWriter.WriteLine(JsonSerializer.Serialize(errorObj));
         }
 
         /// <summary>

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

@@ -0,0 +1,101 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Text;
+
+namespace Azure.DataApiBuilder.Mcp.Core
+{
+    /// <summary>
+    /// Process-wide owner of the MCP stdio server's stdout stream.
+    ///
+    /// In MCP stdio mode, stdout is the JSON-RPC channel and is shared by
+    /// multiple writers — JSON-RPC responses from <see cref="McpStdioServer"/>
+    /// and asynchronous <c>notifications/message</c> frames from the logging
+    /// pipeline. Without coordination, two writers calling <c>WriteLine</c>
+    /// concurrently can interleave at the byte level and corrupt the channel.
+    ///
+    /// This class wraps the underlying <see cref="StreamWriter"/> and serializes
+    /// every write through a single lock so JSON-RPC frames stay intact.
+    /// Registered as a singleton in DI for MCP stdio mode; instantiated lazily
+    /// (the underlying stream is opened on the first write) so non-MCP code
+    /// paths and unit tests can construct the type without side effects.
+    /// </summary>
+    public sealed class McpStdoutWriter : IDisposable
+    {
+        private readonly object _lock = new();
+        private TextWriter? _writer;
+        private bool _disposed;
+
+        /// <summary>
+        /// Production constructor. The underlying stdout stream is opened
+        /// lazily on the first <see cref="WriteLine"/> call.
+        /// </summary>
+        public McpStdoutWriter()
+        {
+        }
+
+        /// <summary>
+        /// Test-only constructor that injects a pre-built writer so unit tests
+        /// can verify lock behavior, disposal semantics, and notification
+        /// framing without touching the real stdout stream.
+        /// </summary>
+        internal McpStdoutWriter(TextWriter writer)
+        {
+            _writer = writer ?? throw new ArgumentNullException(nameof(writer));
+        }
+
+        /// <summary>
+        /// Writes a single line to stdout under a process-wide lock so
+        /// concurrent JSON-RPC responses and notifications cannot interleave.
+        /// No-op after <see cref="Dispose"/>.
+        /// </summary>
+        public void WriteLine(string line)
+        {
+            lock (_lock)
+            {
+                if (_disposed)
+                {
+                    return;
+                }
+
+                EnsureInitialized();
+                _writer!.WriteLine(line);
+            }
+        }
+
+        public void Dispose()
+        {
+            lock (_lock)
+            {
+                if (_disposed)
+                {
+                    return;
+                }
+
+                _disposed = true;
+                _writer?.Dispose();
+                _writer = null;
+            }
+        }
+
+        private void EnsureInitialized()
+        {
+            if (_writer is not null)
+            {
+                return;
+            }
+
+            // Opening the raw stdout stream bypasses any Console.SetOut(...)
+            // redirection. This is intentional: in MCP stdio mode, Program.cs
+            // redirects Console.Out to a sink (TextWriter.Null or stderr) so
+            // stray Console.WriteLine calls from third-party code cannot
+            // corrupt the JSON-RPC channel. Only this class - and only via
+            // WriteLine() - is allowed to write to the real stdout.
+            Stream stdout = Console.OpenStandardOutput();
+            _writer = new StreamWriter(stdout, new UTF8Encoding(encoderShouldEmitUTF8Identifier: false))
+            {
+                AutoFlush = true
+            };
+        }
+    }
+}

src/Azure.DataApiBuilder.Mcp/Model/McpStdioJsonRpcErrorCodes.cs

@@ -7,6 +7,11 @@ namespace Azure.DataApiBuilder.Mcp.Model
     /// </summary>
     internal static class McpStdioJsonRpcErrorCodes
     {
+        /// <summary>
+        /// JSON-RPC protocol version.
+        /// </summary>
+        public const string JSON_RPC_VERSION = "2.0";
+
         /// <summary>
         /// Invalid JSON was received by the server.
         /// An error occurred on the server while parsing the JSON text.

src/Azure.DataApiBuilder.Mcp/Telemetry/McpLogNotificationWriter.cs

@@ -0,0 +1,114 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Text.Json;
+using Azure.DataApiBuilder.Core.Telemetry;
+using Azure.DataApiBuilder.Mcp.Core;
+using Azure.DataApiBuilder.Mcp.Model;
+using Microsoft.Extensions.Logging;
+
+namespace Azure.DataApiBuilder.Mcp.Telemetry;
+
+/// <summary>
+/// Writes log messages as MCP `notifications/message` JSON-RPC notifications.
+/// This allows MCP clients (like MCP Inspector) to receive log output in real-time.
+/// </summary>
+/// <remarks>
+/// MCP spec: https://modelcontextprotocol.io/specification/2025-11-05/server/utilities/logging
+/// The notification format is:
+/// <code>
+/// {
+///   "jsonrpc": "2.0",
+///   "method": "notifications/message",
+///   "params": {
+///     "level": "info",
+///     "logger": "CategoryName",
+///     "data": "The log message"
+///   }
+/// }
+/// </code>
+/// All writes are routed through the shared <see cref="McpStdoutWriter"/> so
+/// notifications cannot interleave with JSON-RPC responses written by
+/// <see cref="McpStdioServer"/>.
+/// </remarks>
+public class McpLogNotificationWriter : IMcpLogNotificationWriter
+{
+    private readonly McpStdoutWriter? _stdoutWriter;
+
+    /// <summary>
+    /// Creates a notification writer that writes through the shared stdout
+    /// writer. The shared writer serializes notifications with JSON-RPC
+    /// responses so concurrent writes do not interleave on the wire.
+    /// </summary>
+    /// <param name="stdoutWriter">
+    /// Shared stdout writer. May be <c>null</c> for unit tests that do not
+    /// exercise the write path; in that case <see cref="WriteNotification"/>
+    /// becomes a no-op.
+    /// </param>
+    public McpLogNotificationWriter(McpStdoutWriter? stdoutWriter = null)
+    {
+        _stdoutWriter = stdoutWriter;
+    }
+
+    /// <summary>
+    /// Gets or sets whether MCP log notifications are enabled. This is the
+    /// single source of truth for whether notifications flow to the client;
+    /// it is consulted by <see cref="McpLogger.IsEnabled(LogLevel)"/> so that
+    /// the gate is enforced once, at log time, before any formatter work runs.
+    /// <see cref="WriteNotification"/> intentionally does not re-check this
+    /// flag — callers must gate via <see cref="McpLogger"/>.
+    /// </summary>
+    public bool IsEnabled { get; set; }
+
+    /// <summary>
+    /// Writes a log message as an MCP notification. The caller is responsible
+    /// for gating on <see cref="IsEnabled"/>; <see cref="McpLogger"/> already
+    /// does this in its <see cref="McpLogger.IsEnabled(LogLevel)"/> override.
+    /// </summary>
+    /// <param name="logLevel">The .NET log level.</param>
+    /// <param name="categoryName">The logger category (typically class name).</param>
+    /// <param name="message">The formatted log message.</param>
+    public void WriteNotification(LogLevel logLevel, string categoryName, string message)
+    {
+        // No IsEnabled check here: the gate lives in McpLogger.IsEnabled so
+        // that we have a single source of truth and a single check site.
+        // The _stdoutWriter null check remains as a defensive guard for unit
+        // tests that construct the writer without a backing stdout.
+        if (_stdoutWriter is null)
+        {
+            return;
+        }
+
+        string mcpLevel = McpLogLevelConverter.ConvertToMcp(logLevel);
+
+        var notification = new
+        {
+            jsonrpc = McpStdioJsonRpcErrorCodes.JSON_RPC_VERSION,
+            method = "notifications/message",
+            @params = new
+            {
+                level = mcpLevel,
+                logger = categoryName,
+                data = message
+            }
+        };
+
+        _stdoutWriter.WriteLine(JsonSerializer.Serialize(notification));
+    }
+}
+
+/// <summary>
+/// Interface for MCP log notification writing.
+/// </summary>
+public interface IMcpLogNotificationWriter
+{
+    /// <summary>
+    /// Gets or sets whether MCP log notifications are enabled.
+    /// </summary>
+    bool IsEnabled { get; set; }
+
+    /// <summary>
+    /// Writes a log message as an MCP notification.
+    /// </summary>
+    void WriteNotification(LogLevel logLevel, string categoryName, string message);
+}