Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 39 additions & 17 deletions src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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
Expand All @@ -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);
Expand All @@ -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);
Expand Down Expand Up @@ -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 { });
}
Expand Down Expand Up @@ -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>
Expand Down
101 changes: 101 additions & 0 deletions src/Azure.DataApiBuilder.Mcp/Core/McpStdoutWriter.cs
Original file line number Diff line number Diff line change
@@ -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)
Comment thread
RubenCerna2079 marked this conversation as resolved.
{
_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
};
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
114 changes: 114 additions & 0 deletions src/Azure.DataApiBuilder.Mcp/Telemetry/McpLogNotificationWriter.cs
Original file line number Diff line number Diff line change
@@ -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);
}
Loading