added tests

Author: Tyler-R-Kendrick

src/ModelContextProtocol/McpEndpointExtensions.cs

@@ -9,7 +9,7 @@ public static class McpEndpointExtensions
     /// <summary>
     /// Notifies the connected endpoint of an event.
     /// </summary>
-    /// <param name="endpoint">The endpoint issueing the notification.</param>
+    /// <param name="endpoint">The endpoint issuing the notification.</param>
     /// <param name="notification">The notification to send.</param>
     /// <param name="cancellationToken">A token to cancel the operation.</param>
     /// <exception cref="ArgumentNullException"><paramref name="endpoint"/> is <see langword="null"/>.</exception>
@@ -27,7 +27,7 @@ public static Task NotifyAsync(
     /// <summary>
     /// Notifies the connected endpoint of an event.
     /// </summary>
-    /// <param name="endpoint">The endpoint issueing the notification.</param>
+    /// <param name="endpoint">The endpoint issuing the notification.</param>
     /// <param name="method">The method to call.</param>
     /// <param name="parameters">The parameters to send.</param>
     /// <param name="cancellationToken">A token to cancel the operation.</param>
@@ -49,7 +49,7 @@ public static Task NotifyAsync(
     }
 
     /// <summary>Notifies the connected endpoint of progress.</summary>
-    /// <param name="endpoint">The endpoint issueing the notification.</param>
+    /// <param name="endpoint">The endpoint issuing the notification.</param>
     /// <param name="progressToken">The <see cref="ProgressToken"/> identifying the operation.</param>
     /// <param name="progress">The progress update to send.</param>
     /// <param name="cancellationToken">A token to cancel the operation.</param>
@@ -75,7 +75,7 @@ public static Task NotifyProgressAsync(
     /// <summary>
     /// Notifies the connected endpoint that a request has been cancelled.
     /// </summary>
-    /// <param name="endpoint">The endpoint issueing the notification.</param>
+    /// <param name="endpoint">The endpoint issuing the notification.</param>
     /// <param name="requestId">The ID of the request to cancel.</param>
     /// <param name="reason">An optional reason for the cancellation.</param>
     /// <param name="cancellationToken">A token to cancel the operation.</param>

src/ModelContextProtocol/Shared/IMcpSession.cs

@@ -0,0 +1,37 @@
+using ModelContextProtocol.Protocol.Messages;
+
+namespace ModelContextProtocol.Shared;
+
+/// <summary>
+/// Class for managing an MCP JSON-RPC session. This covers both MCP clients and servers.
+/// </summary>
+public interface IMcpSession : IDisposable
+{
+    /// <summary>
+    /// The name of the endpoint for logging and debug purposes.
+    /// </summary>
+    string EndpointName { get; set; }
+
+    /// <summary>
+    /// Starts processing messages from the transport. This method will block until the transport is disconnected.
+    /// This is generally started in a background task or thread from the initialization logic of the derived class.
+    /// </summary>
+    Task ProcessMessagesAsync(CancellationToken cancellationToken);
+
+    /// <summary>
+    /// Sends a generic JSON-RPC request to the server.
+    /// </summary>
+    /// <param name="message">The JSON-RPC request to send.</param>
+    /// <param name="cancellationToken">A token to cancel the operation.</param>
+    /// <returns>A task containing the server's response.</returns>
+    Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Sends a request over the protocol
+    /// </summary>
+    /// <typeparam name="TResult">The MCP Response type.</typeparam>
+    /// <param name="request">The request instance</param>
+    /// <param name="cancellationToken">The token for cancellation.</param>
+    /// <returns>The MCP response.</returns>
+    Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class;
+}

src/ModelContextProtocol/Shared/McpJsonRpcEndpoint.cs

@@ -15,7 +15,7 @@ namespace ModelContextProtocol.Shared;
 /// This is especially true as a client represents a connection to one and only one server, and vice versa.
 /// Any multi-client or multi-server functionality should be implemented at a higher level of abstraction.
 /// </summary>
-internal abstract class McpJsonRpcEndpoint : IMcpEndpoint, IAsyncDisposable
+public abstract class McpJsonRpcEndpoint : IMcpEndpoint, IAsyncDisposable
 {
     private readonly RequestHandlers _requestHandlers = [];
     private readonly NotificationHandlers _notificationHandlers = [];
@@ -27,6 +27,9 @@ internal abstract class McpJsonRpcEndpoint : IMcpEndpoint, IAsyncDisposable
     private readonly SemaphoreSlim _disposeLock = new(1, 1);
     private bool _disposed;
 
+    /// <summary>
+    /// The logger for this endpoint.
+    /// </summary>
     protected readonly ILogger _logger;
 
     /// <summary>
@@ -38,18 +41,53 @@ protected McpJsonRpcEndpoint(ILoggerFactory? loggerFactory = null)
         _logger = loggerFactory?.CreateLogger(GetType()) ?? NullLogger.Instance;
     }
 
+    /// <summary>
+    /// Sets the request handler for a specific method.
+    /// </summary>
+    /// <typeparam name="TRequest">The MCP Request type</typeparam>
+    /// <typeparam name="TResponse">The MCP Response type</typeparam>
+    /// <param name="method">The method name.</param>
+    /// <param name="handler">The handler function.</param>
     protected void SetRequestHandler<TRequest, TResponse>(string method, Func<TRequest?, CancellationToken, Task<TResponse>> handler)
         => _requestHandlers.Set(method, handler);
 
+    /// <summary>
+    /// Sets the request handler for a specific method.
+    /// </summary>
+    /// <param name="method">The method name.</param>
+    /// <param name="handler">The handler function.</param>
     public void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler)
         => _notificationHandlers.Add(method, handler);
 
+    /// <summary>
+    /// Sends a request over the protocol
+    /// </summary>
+    /// <typeparam name="TResult">The MCP Response type.</typeparam>
+    /// <param name="request">The request instance</param>
+    /// <param name="cancellationToken">The token for cancellation.</param>
+    /// <returns>The MCP response.</returns>
     public async Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class
     {
-        using var registration = cancellationToken.Register(() => _ = this.NotifyCancelAsync(request.Id));
+        using var registration = cancellationToken.Register(async () =>
+        {
+            try
+            {
+                await this.NotifyCancelAsync(request.Id).ConfigureAwait(false);
+            }
+            catch (Exception ex)
+            {
+                _logger.LogError(ex, "An error occurred while notifying cancellation for request {RequestId}.", request.Id);
+            }
+        });
         return await GetSessionOrThrow().SendRequestAsync<TResult>(request, cancellationToken);
     }
 
+    /// <summary>
+    /// Sends a notification over the protocol.
+    /// </summary>
+    /// <param name="message">The message to send.</param>
+    /// <param name="cancellationToken">The token for cancellation.</param>
+    /// <returns>A task representing the completion of the operation.</returns>
     public Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
         => GetSessionOrThrow().SendMessageAsync(message, cancellationToken);
 
@@ -63,6 +101,12 @@ public Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancella
     /// </summary>
     protected Task? MessageProcessingTask { get; set; }
 
+    /// <summary>
+    /// Starts the session with the given transport.
+    /// </summary>
+    /// <param name="sessionTransport">The transport to use for the session.</param>
+    /// <param name="fullSessionCancellationToken">A cancellation token for the full session.</param>
+    /// <exception cref="InvalidOperationException">Thrown if the session has already started.</exception>
     [MemberNotNull(nameof(MessageProcessingTask))]
     protected void StartSession(ITransport sessionTransport, CancellationToken fullSessionCancellationToken = default)
     {
@@ -76,6 +120,10 @@ protected void StartSession(ITransport sessionTransport, CancellationToken fullS
         MessageProcessingTask = _session.ProcessMessagesAsync(_sessionCts.Token);
     }
 
+    /// <summary>
+    /// Disposes the endpoint and releases resources.
+    /// </summary>
+    /// <returns>A task representing the completion of the operation.</returns>
     public async ValueTask DisposeAsync()
     {
         using var _ = await _disposeLock.LockAsync().ConfigureAwait(false);
@@ -125,6 +173,11 @@ public virtual async ValueTask DisposeUnsynchronizedAsync()
         _logger.EndpointCleanedUp(EndpointName);
     }
 
-    protected McpSession GetSessionOrThrow()
+    /// <summary>
+    /// Gets the current session.
+    /// </summary>
+    /// <returns>The current session.</returns>
+    /// <exception cref="InvalidOperationException">Thrown if the session is not started.</exception>
+    protected IMcpSession GetSessionOrThrow()
         => _session ?? throw new InvalidOperationException($"This should be unreachable from public API! Call {nameof(StartSession)} before sending messages.");
 }

src/ModelContextProtocol/Shared/McpSession.cs

@@ -14,7 +14,7 @@ namespace ModelContextProtocol.Shared;
 /// <summary>
 /// Class for managing an MCP JSON-RPC session. This covers both MCP clients and servers.
 /// </summary>
-internal sealed class McpSession : IDisposable
+internal sealed class McpSession : IMcpSession
 {
     private readonly ITransport _transport;
     private readonly RequestHandlers _requestHandlers;

tests/ModelContextProtocol.Tests/CancelledNotificationTests.cs

@@ -0,0 +1,84 @@
+using ModelContextProtocol.Protocol.Messages;
+using ModelContextProtocol.Protocol.Types;
+using ModelContextProtocol.Tests.Utils;
+
+namespace ModelContextProtocol.Tests;
+
+/// <summary>
+/// Tests for the cancelled notifications against an IMcpEndpoint.
+/// </summary>
+public class CancelledNotificationTests(
+    McpEndpointTestFixture fixture, ITestOutputHelper testOutputHelper)
+    : LoggedTest(testOutputHelper), IClassFixture<McpEndpointTestFixture>
+{
+    [Fact]
+    public async Task NotifyCancelAsync_SendsCorrectNotification()
+    {
+        // Arrange
+        await using var endpoint = fixture.CreateEndpoint();
+        await using var transport = fixture.CreateTransport();
+        var cancellationToken = TestContext.Current.CancellationToken;
+        endpoint.Start(transport, cancellationToken);
+        
+        var requestId = new RequestId("test-request-id-123");
+        const string reason = "Operation was cancelled by the user";
+
+        // Act
+        await endpoint.NotifyCancelAsync(requestId, reason, cancellationToken);
+
+        // Assert
+        Assert.Single(transport.SentMessages);
+        var notification = Assert.IsType<JsonRpcNotification>(transport.SentMessages[0]);
+        Assert.Equal(NotificationMethods.CancelledNotification, notification.Method);
+        
+        var cancelParams = Assert.IsType<CancelledNotification>(notification.Params);
+        Assert.Equal(requestId, cancelParams.RequestId);
+        Assert.Equal(reason, cancelParams.Reason);
+    }
+
+    [Fact]
+    public async Task SendRequestAsync_Cancellation_SendsNotification()
+    {
+        // Arrange
+        await using var endpoint = fixture.CreateEndpoint();
+        await using var transport = fixture.CreateTransport();
+        endpoint.Start(transport, CancellationToken.None);
+        
+        var requestId = new RequestId("test-request-id-123");
+        JsonRpcRequest request = new()
+        {
+            Id = requestId,
+            Method = "test.method",
+            Params = new { },
+        };
+        using CancellationTokenSource cancellationSource = new();
+        await cancellationSource.CancelAsync();
+        // Act
+        try
+        {
+            await endpoint.SendRequestAsync<EmptyResult>(request, cancellationSource.Token);
+        }
+        catch (OperationCanceledException)
+        {
+            // Expected exception
+        }
+        catch (Exception ex)
+        {
+            Assert.Fail($"Unexpected exception: {ex.Message}");
+        }
+
+        // Assert
+        Assert.NotEmpty(transport.SentMessages);
+        Assert.Equal(2, transport.SentMessages.Count);
+        var notification = Assert.IsType<JsonRpcNotification>(transport.SentMessages[0]);
+        Assert.Equal(NotificationMethods.CancelledNotification, notification.Method);
+        
+        var cancelParams = Assert.IsType<CancelledNotification>(notification.Params);
+        Assert.Equal(requestId, cancelParams.RequestId);
+
+        var requestMessage = Assert.IsType<JsonRpcRequest>(transport.SentMessages[1]);
+        Assert.Equal(request.Id, requestMessage.Id);
+        Assert.Equal(request.Method, requestMessage.Method);
+        Assert.Equal(request.Params, requestMessage.Params);
+    }
+}

tests/ModelContextProtocol.Tests/McpEndpointTestFixture.cs

@@ -0,0 +1,54 @@
+using Microsoft.Extensions.Logging;
+using ModelContextProtocol.Protocol.Messages;
+using ModelContextProtocol.Shared;
+using ModelContextProtocol.Protocol.Transport;
+using System.Threading.Channels;
+
+namespace ModelContextProtocol.Tests;
+
+/// <summary>
+/// Test fixture for McpEndpoint tests that provides shared transport implementations.
+/// </summary>
+public class McpEndpointTestFixture() : IAsyncDisposable
+{
+    private readonly ILoggerFactory _loggerFactory = LoggerFactory.Create(builder => builder.AddConsole());
+
+    /// <summary>
+    /// Creates a test transport.
+    /// </summary>
+    internal TestCancellationTransport CreateTransport() => new();
+
+    /// <summary>
+    /// Creates a test client endpoint.
+    /// </summary>
+    internal TestMcpJsonRpcEndpoint CreateEndpoint() => new();
+
+
+    public ValueTask DisposeAsync() => ValueTask.CompletedTask;
+    
+    internal class TestCancellationTransport : ITransport
+    {
+        public bool IsConnected => true;
+        public List<IJsonRpcMessage> SentMessages { get; } = [];
+        public ChannelReader<IJsonRpcMessage> MessageReader { get; init; }
+            = Channel.CreateUnbounded<IJsonRpcMessage>().Reader;
+        public Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken)
+        {
+            SentMessages.Add(message);
+            return Task.CompletedTask;
+        }
+
+        public ValueTask DisposeAsync() => ValueTask.CompletedTask;
+    }
+
+    internal class TestMcpJsonRpcEndpoint(LoggerFactory? loggerFactory = null)
+        : McpJsonRpcEndpoint(loggerFactory ?? new())
+    {
+        public override string EndpointName => "TestEndpoint";
+
+        public void Start(
+            ITransport transport,
+            CancellationToken token)
+            => StartSession(transport, token);
+    }
+}