modelcontextprotocol/csharp-sdk
Watch
1
0
Fork
You've already forked csharp-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Refactor transports to help enable graceful shutdown (#142)

Browse source 
- This also paves the way for better multi-session support
- We should definitely rethink names for the transport API
- For now, I kept the names similar as possible, so we can
  focus on the API shape

Co-authored-by: Stephen Toub <stoub@microsoft.com>
This commit is contained in:
Stephen Halter 2025-03-30 17:10:14 -07:00 • committed by GitHub
commit 2ab9db0878
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 2015 additions and 1847 deletions
Download patch file Download diff file Expand all files Collapse all files

402
src/ModelContextProtocol/Shared/McpJsonRpcEndpoint.cs
View file

@ -1,15 +1,10 @@
using ModelContextProtocol.Client;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using ModelContextProtocol.Logging;
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Utils;
using ModelContextProtocol.Utils.Json;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using System.Collections.Concurrent;
using System.Text.Json;
using System.Diagnostics.CodeAnalysis;
namespace ModelContextProtocol.Shared;
@ -22,356 +17,93 @@ namespace ModelContextProtocol.Shared;
/// </summary>
internal abstract class McpJsonRpcEndpoint : IAsyncDisposable
{
private readonly ITransport _transport;
private readonly ConcurrentDictionary<RequestId, TaskCompletionSource<IJsonRpcMessage>> _pendingRequests;
private readonly ConcurrentDictionary<string, List<Func<JsonRpcNotification, Task>>> _notificationHandlers;
private readonly Dictionary<string, Func<JsonRpcRequest, CancellationToken, Task<object?>>> _requestHandlers = [];
private int _nextRequestId;
private readonly JsonSerializerOptions _jsonOptions;
private bool _isDisposed;
private readonly RequestHandlers _requestHandlers = [];
private readonly NotificationHandlers _notificationHandlers = [];
private McpSession? _session;
private CancellationTokenSource? _sessionCts;
private int _started;
private readonly SemaphoreSlim _disposeLock = new(1, 1);
private bool _disposed;
protected readonly ILogger _logger;
/// <summary>
/// Initializes a new instance of the <see cref="McpJsonRpcEndpoint"/> class.
/// </summary>
/// <param name="transport">An MCP transport implementation.</param>
/// <param name="loggerFactory">The logger factory.</param>
protected McpJsonRpcEndpoint(ITransport transport, ILoggerFactory? loggerFactory = null)
protected McpJsonRpcEndpoint(ILoggerFactory? loggerFactory = null)
{
Throw.IfNull(transport);
_transport = transport;
_pendingRequests = new();
_notificationHandlers = new();
_nextRequestId = 1;
_jsonOptions = McpJsonUtilities.DefaultOptions;
_logger = loggerFactory?.CreateLogger(GetType()) ?? NullLogger.Instance;
}
/// <summary>
/// Gets whether the endpoint is initialized and ready to process messages.
/// </summary>
public bool IsInitialized { get; set; }
protected void SetRequestHandler<TRequest, TResponse>(string method, Func<TRequest?, CancellationToken, Task<TResponse>> handler)
=> _requestHandlers.Set(method, handler);
public void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler)
=> _notificationHandlers.Add(method, handler);
public Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class
=> GetSessionOrThrow().SendRequestAsync<TResult>(request, cancellationToken);
public Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
=> GetSessionOrThrow().SendMessageAsync(message, cancellationToken);
/// <summary>
/// Gets the name of the endpoint for logging and debug purposes.
/// </summary>
public abstract string EndpointName { get; }
/// <summary>
/// Gets the transport implementation for the endpoint. Should generally not be needed outside of tests.
/// Sub-classes should store IClientTransport or IServerTransport injected during construction instead of casting this field.
/// </summary>
internal ITransport Transport => _transport;
/// <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>
internal async Task ProcessMessagesAsync(CancellationToken cancellationToken)
{
try
{
await foreach (var message in _transport.MessageReader.ReadAllAsync(cancellationToken).ConfigureAwait(false))
{
_logger.TransportMessageRead(EndpointName, message.GetType().Name);
// Fire and forget the message handling task to avoid blocking the transport
// If awaiting the task, the transport will not be able to read more messages,
// which could lead to a deadlock if the handler sends a message back
_ = ProcessMessageAsync();
async Task ProcessMessageAsync()
{
#if NET
await Task.CompletedTask.ConfigureAwait(ConfigureAwaitOptions.ForceYielding);
#else
await default(ForceYielding);
#endif
try
{
await HandleMessageAsync(message, cancellationToken).ConfigureAwait(false);
}
catch (Exception ex)
{
var payload = JsonSerializer.Serialize(message, _jsonOptions.GetTypeInfo<IJsonRpcMessage>());
_logger.MessageHandlerError(EndpointName, message.GetType().Name, payload, ex);
}
}
}
}
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
{
// Normal shutdown
_logger.EndpointMessageProcessingCancelled(EndpointName);
}
catch (NullReferenceException)
{
// Ignore reader disposal and mocked transport
}
}
private async Task HandleMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken)
{
switch (message)
{
case JsonRpcRequest request:
await HandleRequest(request, cancellationToken).ConfigureAwait(false);
break;
case IJsonRpcMessageWithId messageWithId:
HandleMessageWithId(message, messageWithId);
break;
case JsonRpcNotification notification:
await HandleNotification(notification).ConfigureAwait(false);
break;
default:
_logger.EndpointHandlerUnexpectedMessageType(EndpointName, message.GetType().Name);
break;
}
}
private async Task HandleNotification(JsonRpcNotification notification)
{
if (_notificationHandlers.TryGetValue(notification.Method, out var handlers))
{
foreach (var notificationHandler in handlers)
{
try
{
await notificationHandler(notification).ConfigureAwait(false);
}
catch (Exception ex)
{
// Log handler error but continue processing
_logger.NotificationHandlerError(EndpointName, notification.Method, ex);
}
}
}
}
private void HandleMessageWithId(IJsonRpcMessage message, IJsonRpcMessageWithId messageWithId)
{
if (!messageWithId.Id.IsValid)
{
_logger.RequestHasInvalidId(EndpointName);
}
else if (_pendingRequests.TryRemove(messageWithId.Id, out var tcs))
{
_logger.ResponseMatchedPendingRequest(EndpointName, messageWithId.Id.ToString());
tcs.TrySetResult(message);
}
else
{
_logger.NoRequestFoundForMessageWithId(EndpointName, messageWithId.Id.ToString());
}
}
private async Task HandleRequest(JsonRpcRequest request, CancellationToken cancellationToken)
{
if (_requestHandlers.TryGetValue(request.Method, out var handler))
{
try
{
_logger.RequestHandlerCalled(EndpointName, request.Method);
var result = await handler(request, cancellationToken).ConfigureAwait(false);
_logger.RequestHandlerCompleted(EndpointName, request.Method);
await _transport.SendMessageAsync(new JsonRpcResponse
{
Id = request.Id,
JsonRpc = "2.0",
Result = result
}, cancellationToken).ConfigureAwait(false);
}
catch (Exception ex)
{
_logger.RequestHandlerError(EndpointName, request.Method, ex);
// Send error response
await _transport.SendMessageAsync(new JsonRpcError
{
Id = request.Id,
JsonRpc = "2.0",
Error = new JsonRpcErrorDetail
{
Code = -32000, // Implementation defined error
Message = ex.Message
}
}, cancellationToken).ConfigureAwait(false);
}
}
else
{
_logger.NoHandlerFoundForRequest(EndpointName, request.Method);
}
}
/// <summary>
/// Sends a generic JSON-RPC request to the server.
/// It is strongly recommended use the capability-specific methods instead of this one.
/// Use this method for custom requests or those not yet covered explicitly by the endpoint implementation.
/// </summary>
/// <typeparam name="TResult">The expected response type.</typeparam>
/// <param name="request">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>
public async Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken) where TResult : class
{
if (!_transport.IsConnected)
{
_logger.EndpointNotConnected(EndpointName);
throw new McpClientException("Transport is not connected");
}
// Set request ID
request.Id = RequestId.FromNumber(Interlocked.Increment(ref _nextRequestId));
var tcs = new TaskCompletionSource<IJsonRpcMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
_pendingRequests[request.Id] = tcs;
try
{
// Expensive logging, use the logging framework to check if the logger is enabled
if (_logger.IsEnabled(LogLevel.Debug))
{
_logger.SendingRequestPayload(EndpointName, JsonSerializer.Serialize(request, _jsonOptions.GetTypeInfo<JsonRpcRequest>()));
}
// Less expensive information logging
_logger.SendingRequest(EndpointName, request.Method);
await _transport.SendMessageAsync(request, cancellationToken).ConfigureAwait(false);
_logger.RequestSentAwaitingResponse(EndpointName, request.Method, request.Id.ToString());
var response = await tcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
if (response is JsonRpcError error)
{
_logger.RequestFailed(EndpointName, request.Method, error.Error.Message, error.Error.Code);
throw new McpClientException($"Request failed (server side): {error.Error.Message}", error.Error.Code);
}
if (response is JsonRpcResponse success)
{
// Convert the Result object to JSON and back to get our strongly-typed result
var resultJson = JsonSerializer.Serialize(success.Result, _jsonOptions.GetTypeInfo<object?>());
var resultObject = JsonSerializer.Deserialize(resultJson, _jsonOptions.GetTypeInfo<TResult>());
// Not expensive logging because we're already converting to JSON in order to get the result object
_logger.RequestResponseReceivedPayload(EndpointName, resultJson);
_logger.RequestResponseReceived(EndpointName, request.Method);
if (resultObject != null)
{
return resultObject;
}
// Result object was null, this is unexpected
_logger.RequestResponseTypeConversionError(EndpointName, request.Method, typeof(TResult));
throw new McpClientException($"Unexpected response type {JsonSerializer.Serialize(success.Result, _jsonOptions.GetTypeInfo<TResult>())}, expected {typeof(TResult)}");
}
// Unexpected response type
_logger.RequestInvalidResponseType(EndpointName, request.Method);
throw new McpClientException("Invalid response type");
}
finally
{
_pendingRequests.TryRemove(request.Id, out _);
}
}
public Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
{
Throw.IfNull(message);
if (!_transport.IsConnected)
{
_logger.ClientNotConnected(EndpointName);
throw new McpClientException("Transport is not connected");
}
if (_logger.IsEnabled(LogLevel.Debug))
{
_logger.SendingMessage(EndpointName, JsonSerializer.Serialize(message, _jsonOptions.GetTypeInfo<IJsonRpcMessage>()));
}
return _transport.SendMessageAsync(message, cancellationToken);
}
/// <summary>
/// Registers a handler for incoming notifications of a specific method.
///
/// <see cref="NotificationMethods">Constants for common notification methods</see>
/// </summary>
/// <param name="method">The notification method to handle.</param>
/// <param name="handler">The async handler function to process notifications.</param>
public void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler)
{
var handlers = _notificationHandlers.GetOrAdd(method, _ => []);
lock (handlers)
{
handlers.Add(handler);
}
}
/// <inheritdoc/>
public async ValueTask DisposeAsync()
{
await CleanupAsync().ConfigureAwait(false);
GC.SuppressFinalize(this);
}
/// <summary>
/// Registers a handler for incoming requests of a specific method.
/// </summary>
/// <typeparam name="TRequest">Type of request payload</typeparam>
/// <typeparam name="TResponse">Type of response payload (not full RPC response</typeparam>
/// <param name="method">Method identifier to register for</param>
/// <param name="handler">Handler to be called when a request with specified method identifier is received</param>
protected void SetRequestHandler<TRequest, TResponse>(string method, Func<TRequest?, CancellationToken, Task<TResponse>> handler)
{
Throw.IfNull(method);
Throw.IfNull(handler);
_requestHandlers[method] = async (request, cancellationToken) =>
{
// Convert the params JsonElement to our type using the same options
var jsonString = JsonSerializer.Serialize(request.Params, _jsonOptions.GetTypeInfo<object?>());
var typedRequest = JsonSerializer.Deserialize(jsonString, _jsonOptions.GetTypeInfo<TRequest>());
return await handler(typedRequest, cancellationToken).ConfigureAwait(false);
};
}
/// <summary>
/// Task that processes incoming messages from the transport.
/// </summary>
protected Task? MessageProcessingTask { get; set; }
/// <summary>
/// CancellationTokenSource used to cancel the message processing task.
/// </summary>
protected CancellationTokenSource? CancellationTokenSource { get; set; }
protected void InitializeSession(ITransport sessionTransport)
{
_session = new McpSession(sessionTransport, EndpointName, _requestHandlers, _notificationHandlers, _logger);
}
[MemberNotNull(nameof(MessageProcessingTask))]
protected void StartSession(CancellationToken fullSessionCancellationToken = default)
{
if (Interlocked.Exchange(ref _started, 1) != 0)
{
throw new InvalidOperationException("The MCP session has already stared.");
}
_sessionCts = CancellationTokenSource.CreateLinkedTokenSource(fullSessionCancellationToken);
MessageProcessingTask = GetSessionOrThrow().ProcessMessagesAsync(_sessionCts.Token);
}
public async ValueTask DisposeAsync()
{
using var _ = await _disposeLock.LockAsync().ConfigureAwait(false);
if (_disposed)
{
return;
}
_disposed = true;
await DisposeUnsynchronizedAsync().ConfigureAwait(false);
}
/// <summary>
/// Cleans up the endpoint and releases resources.
/// </summary>
/// <returns></returns>
protected virtual async Task CleanupAsync()
public virtual async ValueTask DisposeUnsynchronizedAsync()
{
if (_isDisposed)
return;
_isDisposed = true;
_logger.CleaningUpEndpoint(EndpointName);
if (CancellationTokenSource != null)
await CancellationTokenSource.CancelAsync().ConfigureAwait(false);
if (_sessionCts is not null)
{
await _sessionCts.CancelAsync().ConfigureAwait(false);
}
if (MessageProcessingTask != null)
if (MessageProcessingTask is not null)
{
try
{
@ -383,18 +115,12 @@ internal abstract class McpJsonRpcEndpoint : IAsyncDisposable
}
}
// Complete all pending requests with cancellation
foreach (var entry in _pendingRequests)
{
entry.Value.TrySetCanceled();
}
_pendingRequests.Clear();
await _transport.DisposeAsync().ConfigureAwait(false);
CancellationTokenSource?.Dispose();
IsInitialized = false;
_session?.Dispose();
_sessionCts?.Dispose();
_logger.EndpointCleanedUp(EndpointName);
}
protected McpSession GetSessionOrThrow()
=> _session ?? throw new InvalidOperationException($"This should be unreachable from public API! Call {nameof(InitializeSession)} before sending messages.");
}

299
src/ModelContextProtocol/Shared/McpSession.cs Normal file
View file

@ -0,0 +1,299 @@
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using ModelContextProtocol.Client;
using ModelContextProtocol.Logging;
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Utils;
using ModelContextProtocol.Utils.Json;
using System.Collections.Concurrent;
using System.Text.Json;
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
{
private readonly ITransport _transport;
private readonly RequestHandlers _requestHandlers;
private readonly NotificationHandlers _notificationHandlers;
private readonly ConcurrentDictionary<RequestId, TaskCompletionSource<IJsonRpcMessage>> _pendingRequests = [];
private readonly JsonSerializerOptions _jsonOptions;
private readonly ILogger _logger;
private int _nextRequestId;
/// <summary>
/// Initializes a new instance of the <see cref="McpSession"/> class.
/// </summary>
/// <param name="transport">An MCP transport implementation.</param>
/// <param name="endpointName">The name of the endpoint for logging and debug purposes.</param>
/// <param name="requestHandlers">A collection of request handlers.</param>
/// <param name="notificationHandlers">A collection of notification handlers.</param>
/// <param name="logger">The logger.</param>
public McpSession(
ITransport transport,
string endpointName,
RequestHandlers requestHandlers,
NotificationHandlers notificationHandlers,
ILogger logger)
{
Throw.IfNull(transport);
_transport = transport;
EndpointName = endpointName;
_requestHandlers = requestHandlers;
_notificationHandlers = notificationHandlers;
_jsonOptions = McpJsonUtilities.DefaultOptions;
_logger = logger ?? NullLogger.Instance;
}
/// <summary>
/// Gets and sets the name of the endpoint for logging and debug purposes.
/// </summary>
public 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>
public async Task ProcessMessagesAsync(CancellationToken cancellationToken)
{
try
{
await foreach (var message in _transport.MessageReader.ReadAllAsync(cancellationToken).ConfigureAwait(false))
{
_logger.TransportMessageRead(EndpointName, message.GetType().Name);
// Fire and forget the message handling task to avoid blocking the transport
// If awaiting the task, the transport will not be able to read more messages,
// which could lead to a deadlock if the handler sends a message back
_ = ProcessMessageAsync();
async Task ProcessMessageAsync()
{
#if NET
await Task.CompletedTask.ConfigureAwait(ConfigureAwaitOptions.ForceYielding);
#else
await default(ForceYielding);
#endif
try
{
await HandleMessageAsync(message, cancellationToken).ConfigureAwait(false);
}
catch (Exception ex)
{
var payload = JsonSerializer.Serialize(message, _jsonOptions.GetTypeInfo<IJsonRpcMessage>());
_logger.MessageHandlerError(EndpointName, message.GetType().Name, payload, ex);
}
}
}
}
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
{
// Normal shutdown
_logger.EndpointMessageProcessingCancelled(EndpointName);
}
}
private async Task HandleMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken)
{
switch (message)
{
case JsonRpcRequest request:
await HandleRequest(request, cancellationToken).ConfigureAwait(false);
break;
case IJsonRpcMessageWithId messageWithId:
HandleMessageWithId(message, messageWithId);
break;
case JsonRpcNotification notification:
await HandleNotification(notification).ConfigureAwait(false);
break;
default:
_logger.EndpointHandlerUnexpectedMessageType(EndpointName, message.GetType().Name);
break;
}
}
private async Task HandleNotification(JsonRpcNotification notification)
{
if (_notificationHandlers.TryGetValue(notification.Method, out var handlers))
{
foreach (var notificationHandler in handlers)
{
try
{
await notificationHandler(notification).ConfigureAwait(false);
}
catch (Exception ex)
{
// Log handler error but continue processing
_logger.NotificationHandlerError(EndpointName, notification.Method, ex);
}
}
}
}
private void HandleMessageWithId(IJsonRpcMessage message, IJsonRpcMessageWithId messageWithId)
{
if (!messageWithId.Id.IsValid)
{
_logger.RequestHasInvalidId(EndpointName);
}
else if (_pendingRequests.TryRemove(messageWithId.Id, out var tcs))
{
_logger.ResponseMatchedPendingRequest(EndpointName, messageWithId.Id.ToString());
tcs.TrySetResult(message);
}
else
{
_logger.NoRequestFoundForMessageWithId(EndpointName, messageWithId.Id.ToString());
}
}
private async Task HandleRequest(JsonRpcRequest request, CancellationToken cancellationToken)
{
if (_requestHandlers.TryGetValue(request.Method, out var handler))
{
try
{
_logger.RequestHandlerCalled(EndpointName, request.Method);
var result = await handler(request, cancellationToken).ConfigureAwait(false);
_logger.RequestHandlerCompleted(EndpointName, request.Method);
await _transport.SendMessageAsync(new JsonRpcResponse
{
Id = request.Id,
JsonRpc = "2.0",
Result = result
}, cancellationToken).ConfigureAwait(false);
}
catch (Exception ex)
{
_logger.RequestHandlerError(EndpointName, request.Method, ex);
// Send error response
await _transport.SendMessageAsync(new JsonRpcError
{
Id = request.Id,
JsonRpc = "2.0",
Error = new JsonRpcErrorDetail
{
Code = -32000, // Implementation defined error
Message = ex.Message
}
}, cancellationToken).ConfigureAwait(false);
}
}
else
{
_logger.NoHandlerFoundForRequest(EndpointName, request.Method);
}
}
/// <summary>
/// Sends a generic JSON-RPC request to the server.
/// It is strongly recommended use the capability-specific methods instead of this one.
/// Use this method for custom requests or those not yet covered explicitly by the endpoint implementation.
/// </summary>
/// <typeparam name="TResult">The expected response type.</typeparam>
/// <param name="request">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>
public async Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken) where TResult : class
{
if (!_transport.IsConnected)
{
_logger.EndpointNotConnected(EndpointName);
throw new McpClientException("Transport is not connected");
}
// Set request ID
request.Id = RequestId.FromNumber(Interlocked.Increment(ref _nextRequestId));
var tcs = new TaskCompletionSource<IJsonRpcMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
_pendingRequests[request.Id] = tcs;
try
{
// Expensive logging, use the logging framework to check if the logger is enabled
if (_logger.IsEnabled(LogLevel.Debug))
{
_logger.SendingRequestPayload(EndpointName, JsonSerializer.Serialize(request, _jsonOptions.GetTypeInfo<JsonRpcRequest>()));
}
// Less expensive information logging
_logger.SendingRequest(EndpointName, request.Method);
await _transport.SendMessageAsync(request, cancellationToken).ConfigureAwait(false);
_logger.RequestSentAwaitingResponse(EndpointName, request.Method, request.Id.ToString());
var response = await tcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
if (response is JsonRpcError error)
{
_logger.RequestFailed(EndpointName, request.Method, error.Error.Message, error.Error.Code);
throw new McpClientException($"Request failed (server side): {error.Error.Message}", error.Error.Code);
}
if (response is JsonRpcResponse success)
{
// Convert the Result object to JSON and back to get our strongly-typed result
var resultJson = JsonSerializer.Serialize(success.Result, _jsonOptions.GetTypeInfo<object?>());
var resultObject = JsonSerializer.Deserialize(resultJson, _jsonOptions.GetTypeInfo<TResult>());
// Not expensive logging because we're already converting to JSON in order to get the result object
_logger.RequestResponseReceivedPayload(EndpointName, resultJson);
_logger.RequestResponseReceived(EndpointName, request.Method);
if (resultObject != null)
{
return resultObject;
}
// Result object was null, this is unexpected
_logger.RequestResponseTypeConversionError(EndpointName, request.Method, typeof(TResult));
throw new McpClientException($"Unexpected response type {JsonSerializer.Serialize(success.Result, _jsonOptions.GetTypeInfo<TResult>())}, expected {typeof(TResult)}");
}
// Unexpected response type
_logger.RequestInvalidResponseType(EndpointName, request.Method);
throw new McpClientException("Invalid response type");
}
finally
{
_pendingRequests.TryRemove(request.Id, out _);
}
}
public Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
{
Throw.IfNull(message);
if (!_transport.IsConnected)
{
_logger.ClientNotConnected(EndpointName);
throw new McpClientException("Transport is not connected");
}
if (_logger.IsEnabled(LogLevel.Debug))
{
_logger.SendingMessage(EndpointName, JsonSerializer.Serialize(message, _jsonOptions.GetTypeInfo<IJsonRpcMessage>()));
}
return _transport.SendMessageAsync(message, cancellationToken);
}
public void Dispose()
{
// Complete all pending requests with cancellation
foreach (var entry in _pendingRequests)
{
entry.Value.TrySetCanceled();
}
_pendingRequests.Clear();
}
}

16
src/ModelContextProtocol/Shared/NotificationHandlers.cs Normal file
View file

@ -0,0 +1,16 @@
using ModelContextProtocol.Protocol.Messages;
using System.Collections.Concurrent;
namespace ModelContextProtocol.Shared;
internal sealed class NotificationHandlers : ConcurrentDictionary<string, List<Func<JsonRpcNotification, Task>>>
{
public void Add(string method, Func<JsonRpcNotification, Task> handler)
{
var handlers = GetOrAdd(method, _ => []);
lock (handlers)
{
handlers.Add(handler);
}
}
}

31
src/ModelContextProtocol/Shared/RequestHandlers.cs Normal file
View file

@ -0,0 +1,31 @@
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Utils;
using ModelContextProtocol.Utils.Json;
using System.Text.Json;
namespace ModelContextProtocol.Shared;
internal sealed class RequestHandlers : Dictionary<string, Func<JsonRpcRequest, CancellationToken, Task<object?>>>
{
/// <summary>
/// Registers a handler for incoming requests of a specific method.
/// </summary>
/// <typeparam name="TRequest">Type of request payload</typeparam>
/// <typeparam name="TResponse">Type of response payload (not full RPC response</typeparam>
/// <param name="method">Method identifier to register for</param>
/// <param name="handler">Handler to be called when a request with specified method identifier is received</param>
public void Set<TRequest, TResponse>(string method, Func<TRequest?, CancellationToken, Task<TResponse>> handler)
{
Throw.IfNull(method);
Throw.IfNull(handler);
this[method] = async (request, cancellationToken) =>
{
// Convert the params JsonElement to our type using the same options
var jsonString = JsonSerializer.Serialize(request.Params, McpJsonUtilities.DefaultOptions.GetTypeInfo<object?>());
var typedRequest = JsonSerializer.Deserialize(jsonString, McpJsonUtilities.DefaultOptions.GetTypeInfo<TRequest>());
return await handler(typedRequest, cancellationToken).ConfigureAwait(false);
};
}
}