diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs index 690098384b..5931530942 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 TextReader? _inputReader; private readonly string _protocolVersion; private const int MAX_LINE_LENGTH = 1024 * 1024; // 1 MB limit for incoming JSON-RPC requests @@ -39,10 +40,11 @@ public class McpStdioServer : IMcpStdioServer DefaultIgnoreCondition = System.Text.Json.Serialization.JsonIgnoreCondition.WhenWritingNull }; - public McpStdioServer(McpToolRegistry toolRegistry, IServiceProvider serviceProvider) + public McpStdioServer(McpToolRegistry toolRegistry, IServiceProvider serviceProvider, TextReader? inputReader = null) { _toolRegistry = toolRegistry ?? throw new ArgumentNullException(nameof(toolRegistry)); _serviceProvider = serviceProvider ?? throw new ArgumentNullException(nameof(serviceProvider)); + _inputReader = inputReader; // Resolve the shared stdout writer so JSON-RPC responses and // notifications/message frames are serialized through one lock. @@ -61,17 +63,20 @@ public McpStdioServer(McpToolRegistry toolRegistry, IServiceProvider serviceProv /// A task representing the asynchronous operation. public async Task RunAsync(CancellationToken cancellationToken) { - // 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 StreamReader reader = new(stdin, utf8NoBom); + // By default read via Console.In so the loop honors the configured + // Console.InputEncoding in stdio mode. + TextReader reader = _inputReader ?? Console.In; while (!cancellationToken.IsCancellationRequested) { string? line = await reader.ReadLineAsync(cancellationToken); + + // EOF (stdin pipe closed) is a normal shutdown signal for stdio mode. + if (line is null) + { + return; + } + if (string.IsNullOrWhiteSpace(line)) { continue; diff --git a/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs b/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs new file mode 100644 index 0000000000..3477e775a5 --- /dev/null +++ b/src/Service.Tests/UnitTests/McpStdioServerRunAsyncTests.cs @@ -0,0 +1,79 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#nullable enable + +using System; +using System.IO; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Azure.DataApiBuilder.Mcp.Core; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.VisualStudio.TestTools.UnitTesting; + +namespace Azure.DataApiBuilder.Service.Tests.UnitTests +{ + [TestClass] + public class McpStdioServerRunAsyncTests + { + [TestMethod] + public async Task RunAsync_EofOnStdin_ExitsGracefullyWithoutOutput() + { + // Empty input immediately yields EOF (ReadLineAsync returns null). + (McpStdioServer server, StringWriter stdoutCapture) = + CreateServerWithCapturedOutput(new StringReader(string.Empty)); + + await server.RunAsync(CancellationToken.None); + + Assert.AreEqual(string.Empty, stdoutCapture.ToString(), + "Server should exit cleanly on EOF without emitting protocol output."); + } + + [TestMethod] + public async Task RunAsync_BlankLineThenShutdown_IgnoresBlankLineAndHandlesShutdown() + { + (McpStdioServer server, StringWriter stdoutCapture) = + CreateServerWithCapturedOutput(new StringReader(Environment.NewLine + + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"shutdown\"}" + + Environment.NewLine)); + + await server.RunAsync(CancellationToken.None); + + string[] lines = stdoutCapture + .ToString() + .Split(Environment.NewLine, StringSplitOptions.RemoveEmptyEntries); + + Assert.AreEqual(1, lines.Length, + "Expected a single response line for shutdown request."); + + using JsonDocument response = JsonDocument.Parse(lines[0]); + JsonElement root = response.RootElement; + + Assert.AreEqual("2.0", root.GetProperty("jsonrpc").GetString(), + "Expected jsonrpc version 2.0 in shutdown response."); + Assert.AreEqual(1, root.GetProperty("id").GetInt32(), + "Expected shutdown response id to match request id."); + Assert.IsTrue(root.GetProperty("result").GetProperty("ok").GetBoolean(), + "Expected shutdown response result.ok to be true."); + } + + private static (McpStdioServer server, StringWriter stdoutCapture) CreateServerWithCapturedOutput(TextReader inputReader) + { + StringWriter stdoutCapture = new(); + McpStdoutWriter stdoutWriter = new(stdoutCapture); + + ServiceCollection services = new(); + services.AddSingleton(stdoutWriter); + services.AddSingleton(); + IServiceProvider serviceProvider = services.BuildServiceProvider(); + + McpStdioServer server = new( + serviceProvider.GetRequiredService(), + serviceProvider, + inputReader); + + return (server, stdoutCapture); + } + } +}