Add more tests (#51)

* Add more tests

* Fix some code smells

* Remove more code smells

* Combine "ServerIntegrationTests" and "ClientIntegrationTests" because both are client integration tests just using different servers

* Clean up resources after test

* Improve kill logic for spawned child processes

* Improved testability for Stdio

* Fix SseIntegration tests

src/mcpdotnet/Client/McpClientFactory.cs

@@ -1,4 +1,5 @@
-using System.Runtime.InteropServices;
+using System.Globalization;
+using System.Runtime.InteropServices;
 using McpDotNet.Configuration;
 using McpDotNet.Logging;
 using McpDotNet.Protocol.Transport;
@@ -11,7 +12,7 @@ namespace McpDotNet.Client;
 /// All server configurations must be passed in the constructor. 
 /// Capabilities (as defined in client options) are shared across all clients, as the client host can always decide not to use certain capabilities.
 /// </summary>
-public class McpClientFactory
+public class McpClientFactory : IDisposable
 {
     private const string ARGUMENTS_OPTIONS_KEY = "arguments";
     private const string COMMAND_OPTIONS_KEY = "command";
@@ -22,6 +23,12 @@ public class McpClientFactory
     private readonly Func<IClientTransport, McpServerConfig, McpClientOptions, IMcpClient> _clientFactoryMethod;
     private readonly ILoggerFactory _loggerFactory;
     private readonly ILogger<McpClientFactory> _logger;
+    private bool _isDisposed;
+
+    /// <summary>
+    /// Gets or sets a value indicating whether clients should be disposed when the factory is disposed.
+    /// </summary>
+    public bool DisposeClientsOnDispose { get; set; } = true;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="McpClientFactory"/> class.
@@ -109,7 +116,8 @@ public class McpClientFactory
                 WorkingDirectory = config.TransportOptions?.GetValueOrDefault("workingDirectory"),
                 EnvironmentVariables = config.TransportOptions?
                     .Where(kv => kv.Key.StartsWith("env:", StringComparison.Ordinal))
-                    .ToDictionary(kv => kv.Key.Substring(4), kv => kv.Value)
+                    .ToDictionary(kv => kv.Key.Substring(4), kv => kv.Value),
+                ShutdownTimeout = TimeSpan.TryParse(config.TransportOptions?.GetValueOrDefault("shutdownTimeout"), CultureInfo.InvariantCulture, out var timespan) ? timespan : StdioClientTransportOptions.DefaultShutdownTimeout
             }, config, _loggerFactory);
         }
 
@@ -181,4 +189,37 @@ public class McpClientFactory
             ? $"/c {command} {args}"
             : $"/c {command}";
     }
+
+    /// <summary>
+    /// Disposes all clients created by the factory.
+    /// </summary>
+    /// <param name="disposing"></param>
+    protected virtual void Dispose(bool disposing)
+    {
+        if (!_isDisposed)
+        {
+            if (disposing && DisposeClientsOnDispose)
+                DisposeClients();
+
+            _isDisposed = true;
+        }
+    }
+
+    private void DisposeClients()
+    {
+        foreach (var client in _clients.Values)
+        {
+            client?.DisposeAsync().AsTask().Wait();
+        }
+
+        _clients.Clear();
+    }
+
+    /// <inheritdoc />
+    public void Dispose()
+    {
+        // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
+        Dispose(disposing: true);
+        GC.SuppressFinalize(this);
+    }
 }

src/mcpdotnet/Configuration/McpServerBuilderExtensions.Tools.cs

@@ -116,7 +116,7 @@ public static partial class McpServerBuilderExtensions
         if (toolTypes.Count == 0)
             throw new ArgumentException("No types with marked methods found in the assembly.", nameof(assembly));
 
-        return WithTools(builder, toolTypes.ToArray());
+        return WithTools(builder, [.. toolTypes]);
     }
 
     private static Tool CreateTool(MethodInfo method, McpToolAttribute mcpAttribute)
@@ -187,7 +187,7 @@ public static partial class McpServerBuilderExtensions
                 BindingFlags.DoNotWrapExceptions |
 #endif
                 BindingFlags.Default;
-            result = method.Invoke(objectInstance, InvokeFlags, binder: null, parameters.ToArray(), culture: null);
+            result = method.Invoke(objectInstance, InvokeFlags, binder: null, [.. parameters], culture: null);
         }
         catch (TargetInvocationException e) when (e.InnerException is not null)
         {
@@ -206,7 +206,7 @@ public static partial class McpServerBuilderExtensions
             return new CallToolResponse { Content = [new Content() { Text = resultString, Type = "text" }] };
 
         if (result is string[] resultStringArray)
-            return new CallToolResponse { Content = resultStringArray.Select(s => new Content() { Text = s, Type = "text" }).ToList() };
+            return new CallToolResponse { Content = [.. resultStringArray.Select(s => new Content() { Text = s, Type = "text" })] };
 
         if (result is null)
             return new CallToolResponse { Content = [new Content() { Text = "null" }] };

src/mcpdotnet/Configuration/McpServerServiceCollectionExtension.cs

@@ -1,6 +1,7 @@
 using System.Reflection;
 using McpDotNet.Configuration;
 using McpDotNet.Hosting;
+using McpDotNet.Protocol.Transport;
 using McpDotNet.Protocol.Types;
 using McpDotNet.Server;
 using Microsoft.Extensions.DependencyInjection;
@@ -38,6 +39,7 @@ public static class McpServerServiceCollectionExtension
         services.AddSingleton<IMcpServerFactory, McpServerFactory>();
         services.AddHostedService<McpServerHostedService>();
         services.AddOptions();
+        services.AddSingleton(InputOutputStreams.Default);
 
         services.AddSingleton<IMcpServer>(sp => sp.GetRequiredService<IMcpServerFactory>().CreateServer());
 

src/mcpdotnet/Protocol/Transport/InputOutputStreams.cs

@@ -0,0 +1,26 @@
+namespace McpDotNet.Protocol.Transport;
+
+/// <summary>
+/// Container for the input and output streams.
+/// </summary>
+public class InputOutputStreams
+{
+    /// <summary>
+    /// Gets the default input and output streams.
+    /// </summary>
+    public static readonly InputOutputStreams Default = new()
+    {
+        Output = Console.Out,
+        Input = Console.In
+    };
+
+    /// <summary>
+    /// Gets or sets the output stream.
+    /// </summary>
+    public required TextWriter Output { get; init; }
+
+    /// <summary>
+    /// Gets or sets the input stream.
+    /// </summary>
+    public required TextReader Input { get; init; }
+}

src/mcpdotnet/Protocol/Transport/SseClientTransport.cs

@@ -119,7 +119,7 @@ public sealed class SseClientTransport : TransportBase, IClientTransport
         if (_messageEndpoint == null)
             throw new InvalidOperationException("Transport not connected");
 
-        var content = new StringContent(
+        using var content = new StringContent(
             JsonSerializer.Serialize(message, _jsonOptions),
             Encoding.UTF8,
             "application/json"
@@ -217,7 +217,7 @@ public sealed class SseClientTransport : TransportBase, IClientTransport
         {
             try
             {
-                var request = new HttpRequestMessage(HttpMethod.Get, _sseEndpoint);
+                using var request = new HttpRequestMessage(HttpMethod.Get, _sseEndpoint);
                 request.Headers.Accept.Add(new MediaTypeWithQualityHeaderValue("text/event-stream"));
 
                 using var response = await _httpClient.SendAsync(
@@ -328,7 +328,7 @@ public sealed class SseClientTransport : TransportBase, IClientTransport
                 if (hostUrl.EndsWith("/sse", StringComparison.Ordinal))
                     hostUrl = hostUrl[..^4];
 
-                var endpointUri = hostUrl + data;
+                var endpointUri = $"{hostUrl.TrimEnd('/')}/{data.TrimStart('/')}";
 
                 _messageEndpoint = new Uri(endpointUri);
             }

src/mcpdotnet/Protocol/Transport/StdioClientTransport.cs

@@ -5,6 +5,7 @@ using System.Text.Json;
 using McpDotNet.Configuration;
 using McpDotNet.Logging;
 using McpDotNet.Protocol.Messages;
+using McpDotNet.Utils;
 using McpDotNet.Utils.Json;
 using Microsoft.Extensions.Logging;
 
@@ -245,12 +246,10 @@ public sealed class StdioClientTransport : TransportBase, IClientTransport
 
                 // Wait for the process to exit
                 _logger.TransportWaitingForShutdown(EndpointName);
-                if (!_process.WaitForExit((int)_options.ShutdownTimeout.TotalMilliseconds))
-                {
-                    // If it doesn't exit gracefully, terminate it
-                    _logger.TransportKillingProcess(EndpointName);
-                    _process.Kill(true);
-                }
+
+                // Kill the while process tree because the process may spawn child processes
+                // and Node.js does not kill its children when it exits properly
+                _process.KillTree(_options.ShutdownTimeout);
             }
             catch (Exception ex)
             {
@@ -262,9 +261,11 @@ public sealed class StdioClientTransport : TransportBase, IClientTransport
         }
 
         if (_shutdownCts != null)
+        {
             await _shutdownCts.CancelAsync().ConfigureAwait(false);
-        _shutdownCts?.Dispose();
-        _shutdownCts = null;
+            _shutdownCts.Dispose();
+            _shutdownCts = null;
+        }
 
         if (_readTask != null)
         {

src/mcpdotnet/Protocol/Transport/StdioClientTransportOptions.cs

@@ -6,6 +6,11 @@ namespace McpDotNet.Protocol.Transport;
 /// </summary>
 public record StdioClientTransportOptions
 {
+    /// <summary>
+    /// The default timeout to wait for the server to shut down gracefully.
+    /// </summary>
+    public static readonly TimeSpan DefaultShutdownTimeout = TimeSpan.FromSeconds(5);
+
     /// <summary>
     /// The command to execute to start the server process.
     /// </summary>
@@ -29,5 +34,5 @@ public record StdioClientTransportOptions
     /// <summary>
     /// The timeout to wait for the server to shut down gracefully.
     /// </summary>
-    public TimeSpan ShutdownTimeout { get; init; } = TimeSpan.FromSeconds(5);
+    public TimeSpan ShutdownTimeout { get; init; } = DefaultShutdownTimeout;
 }

src/mcpdotnet/Protocol/Transport/StdioServerTransport.cs

@@ -18,6 +18,8 @@ public sealed class StdioServerTransport : TransportBase, IServerTransport
     private readonly JsonSerializerOptions _jsonOptions;
     private Task? _readTask;
     private CancellationTokenSource? _shutdownCts;
+    private readonly TextReader _input;
+    private readonly TextWriter _output;
 
     private string EndpointName => $"Server (stdio) ({_serverName})";
 
@@ -26,8 +28,14 @@ public sealed class StdioServerTransport : TransportBase, IServerTransport
     /// </summary>
     /// <param name="serverOptions">The server options.</param>
     /// <param name="loggerFactory">A logger factory for creating loggers.</param>
-    public StdioServerTransport(McpServerOptions serverOptions, ILoggerFactory? loggerFactory)
-        : this(serverOptions is not null ? serverOptions.ServerInfo.Name : throw new ArgumentNullException(nameof(serverOptions)), loggerFactory)
+    /// <param name="inputOutputStreams">Container for the input and output streams</param>
+    public StdioServerTransport(McpServerOptions serverOptions, ILoggerFactory? loggerFactory, InputOutputStreams inputOutputStreams)
+        : this(
+              serverOptions is not null ? serverOptions.ServerInfo.Name : throw new ArgumentNullException(nameof(serverOptions)),
+              loggerFactory,
+              inputOutputStreams is not null ? inputOutputStreams.Output : throw new ArgumentNullException(nameof(inputOutputStreams)),
+              inputOutputStreams.Input
+              )
     {
     }
 
@@ -37,11 +45,26 @@ public sealed class StdioServerTransport : TransportBase, IServerTransport
     /// <param name="serverName">The name of the server.</param>
     /// <param name="loggerFactory">A logger factory for creating loggers.</param>
     public StdioServerTransport(string serverName, ILoggerFactory? loggerFactory)
+        : this(serverName, loggerFactory, Console.Out, Console.In)
+    {
+
+    }
+
+    /// <summary>
+    /// Initializes a new instance of the StdioServerTransport class.
+    /// </summary>
+    /// <param name="serverName">The name of the server.</param>
+    /// <param name="loggerFactory">A logger factory for creating loggers.</param>
+    /// <param name="output">The output stream writer</param>
+    /// <param name="input">The input stream reader</param>
+    public StdioServerTransport(string serverName, ILoggerFactory? loggerFactory, TextWriter output, TextReader input)
         : base(loggerFactory)
     {
         _serverName = serverName;
         _logger = loggerFactory is not null ? loggerFactory.CreateLogger<StdioClientTransport>() : NullLogger.Instance;
         _jsonOptions = JsonSerializerOptionsExtensions.DefaultOptions;
+        _input = input ?? throw new ArgumentNullException(nameof(input));
+        _output = output ?? throw new ArgumentNullException(nameof(output));
     }
 
     /// <inheritdoc/>
@@ -56,7 +79,6 @@ public sealed class StdioServerTransport : TransportBase, IServerTransport
         return Task.CompletedTask;
     }
 
-
     /// <inheritdoc/>
     public override async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
     {
@@ -77,8 +99,8 @@ public sealed class StdioServerTransport : TransportBase, IServerTransport
             var json = JsonSerializer.Serialize(message, _jsonOptions);
             _logger.TransportSendingMessage(EndpointName, id, json);
 
-            await Console.Out.WriteLineAsync(json.AsMemory(), cancellationToken).ConfigureAwait(false);
-            await Console.Out.FlushAsync(cancellationToken).ConfigureAwait(false);
+            await _output.WriteLineAsync(json.AsMemory(), cancellationToken).ConfigureAwait(false);
+            await _output.FlushAsync(cancellationToken).ConfigureAwait(false);
 
             _logger.TransportSentMessage(EndpointName, id);
         }
@@ -106,7 +128,7 @@ public sealed class StdioServerTransport : TransportBase, IServerTransport
             {
                 _logger.TransportWaitingForMessage(EndpointName);
 
-                var reader = Console.In;
+                var reader = _input;
                 var line = await reader.ReadLineAsync(cancellationToken).ConfigureAwait(false);
                 if (line == null)
                 {

src/mcpdotnet/Shared/McpJsonRpcEndpoint.cs

@@ -25,6 +25,7 @@ internal abstract class McpJsonRpcEndpoint : IAsyncDisposable
     private int _nextRequestId;
     private readonly JsonSerializerOptions _jsonOptions;
     private readonly ILogger<McpClient> _logger;
+    private bool _isDisposed;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="McpJsonRpcEndpoint"/> class.
@@ -398,6 +399,11 @@ internal abstract class McpJsonRpcEndpoint : IAsyncDisposable
     /// <returns></returns>
     protected async Task CleanupAsync()
     {
+        if (_isDisposed)
+            return;
+
+        _isDisposed = true;
+
         _logger.CleaningUpEndpoint(EndpointName);
 
         if (CancellationTokenSource != null)

src/mcpdotnet/Utils/ProcessHelper.cs

@@ -0,0 +1,108 @@
+using System.Diagnostics;
+using System.Runtime.InteropServices;
+
+namespace McpDotNet.Utils;
+
+/// <summary>
+/// Helper class for working with processes.
+/// </summary>
+internal static class ProcessHelper
+{
+    private static readonly bool _isWindows = RuntimeInformation.IsOSPlatform(OSPlatform.Windows);
+    private static readonly TimeSpan _defaultTimeout = TimeSpan.FromSeconds(30);
+
+    public static void KillTree(this Process process) => process.KillTree(_defaultTimeout);
+
+    public static void KillTree(this Process process, TimeSpan timeout)
+    {
+        var pid = process.Id;
+        if (_isWindows)
+        {
+            RunProcessAndWaitForExit(
+                "taskkill",
+                $"/T /F /PID {pid}",
+                timeout,
+                out var _);
+        }
+        else
+        {
+            var children = new HashSet<int>();
+            GetAllChildIdsUnix(pid, children, timeout);
+            foreach (var childId in children)
+            {
+                KillProcessUnix(childId, timeout);
+            }
+            KillProcessUnix(pid, timeout);
+        }
+
+        // wait until the process finishes exiting/getting killed. 
+        // We don't want to wait forever here because the task is already supposed to be dieing, we just want to give it long enough
+        // to try and flush what it can and stop. If it cannot do that in a reasonable time frame then we will just ignore it.
+        process.WaitForExit((int)timeout.TotalMilliseconds);
+    }
+
+    private static void GetAllChildIdsUnix(int parentId, ISet<int> children, TimeSpan timeout)
+    {
+        var exitcode = RunProcessAndWaitForExit(
+            "pgrep",
+            $"-P {parentId}",
+            timeout,
+            out var stdout);
+
+        if (exitcode == 0 && !string.IsNullOrEmpty(stdout))
+        {
+            using var reader = new StringReader(stdout);
+            while (true)
+            {
+                var text = reader.ReadLine();
+                if (text == null)
+                    return;
+
+                if (int.TryParse(text, out var id))
+                {
+                    children.Add(id);
+                    // Recursively get the children
+                    GetAllChildIdsUnix(id, children, timeout);
+                }
+            }
+        }
+    }
+
+    private static void KillProcessUnix(int processId, TimeSpan timeout)
+    {
+        RunProcessAndWaitForExit(
+            "kill",
+            $"-TERM {processId}",
+            timeout,
+            out var _);
+    }
+
+    private static int RunProcessAndWaitForExit(string fileName, string arguments, TimeSpan timeout, out string? stdout)
+    {
+        var startInfo = new ProcessStartInfo
+        {
+            FileName = fileName,
+            Arguments = arguments,
+            RedirectStandardOutput = true,
+            RedirectStandardError = true,
+            UseShellExecute = false,
+        };
+
+        stdout = null;
+
+        var process = Process.Start(startInfo);
+        if (process == null)
+            return -1;
+
+        if (process.WaitForExit((int)timeout.TotalMilliseconds))
+        {
+            stdout = process.StandardOutput.ReadToEnd();
+        }
+        else
+        {
+            process.Kill();
+        }
+
+        return process.ExitCode;
+    }
+}