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

Make supposedly unreachable code less reachable (#178)

Browse source 
* Move WriteJsonRpcMessageToBuffer to method

* Fix indentation in McpClient

* Make supposedly unreachable code less reachable

* Guard against multiple RunAsync calls

* Move RunAsync inside of Try in MapMcp
This commit is contained in:
Stephen Halter 2025-04-01 19:29:59 -07:00 • committed by GitHub
commit 674cb1515c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 94 additions and 80 deletions
Download patch file Download diff file Expand all files Collapse all files

4
src/ModelContextProtocol.AspNetCore/McpEndpointRouteBuilderExtensions.cs
View file

@ -47,11 +47,11 @@ public static class McpEndpointRouteBuilderExtensions
{
throw new Exception($"Unreachable given good entropy! Session with ID '{sessionId}' has already been created.");
}
await using var server = McpServerFactory.Create(transport, mcpServerOptions.Value, loggerFactory, endpoints.ServiceProvider);
try
{
var transportTask = transport.RunAsync(cancellationToken: requestAborted);
await using var server = McpServerFactory.Create(transport, mcpServerOptions.Value, loggerFactory, endpoints.ServiceProvider);
try
{
@ -85,7 +85,7 @@ public static class McpEndpointRouteBuilderExtensions
if (!_sessions.TryGetValue(sessionId.ToString(), out var transport))
{
await Results.BadRequest($"Session {sessionId} not found.").ExecuteAsync(context);
await Results.BadRequest($"Session ID not found.").ExecuteAsync(context);
return;
}

32
src/ModelContextProtocol/Client/McpClient.cs
View file

@ -82,29 +82,27 @@ internal sealed class McpClient : McpJsonRpcEndpoint, IMcpClient
{
// Connect transport
_sessionTransport = await _clientTransport.ConnectAsync(cancellationToken).ConfigureAwait(false);
// We don't want the ConnectAsync token to cancel the session after we've successfully connected.
// The base class handles cleaning up the session in DisposeAsync without our help.
StartSession(_sessionTransport, fullSessionCancellationToken: CancellationToken.None);
StartSession(_sessionTransport);
// Perform initialization sequence
using var initializationCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
initializationCts.CancelAfter(_options.InitializationTimeout);
try
{
// Send initialize request
var initializeResponse = await SendRequestAsync<InitializeResult>(
new JsonRpcRequest
{
Method = RequestMethods.Initialize,
Params = new InitializeRequestParams()
try
{
// Send initialize request
var initializeResponse = await SendRequestAsync<InitializeResult>(
new JsonRpcRequest
{
ProtocolVersion = _options.ProtocolVersion,
Capabilities = _options.Capabilities ?? new ClientCapabilities(),
ClientInfo = _options.ClientInfo
}
},
initializationCts.Token).ConfigureAwait(false);
Method = RequestMethods.Initialize,
Params = new InitializeRequestParams()
{
ProtocolVersion = _options.ProtocolVersion,
Capabilities = _options.Capabilities ?? new ClientCapabilities(),
ClientInfo = _options.ClientInfo
}
},
initializationCts.Token).ConfigureAwait(false);
// Store server information
_logger.ServerCapabilitiesReceived(EndpointName,

26
src/ModelContextProtocol/Protocol/Transport/SseResponseStreamTransport.cs
View file

@ -32,19 +32,6 @@ public sealed class SseResponseStreamTransport(Stream sseResponseStream, string
/// <returns>A task representing the send loop that writes JSON-RPC messages to the SSE response stream.</returns>
public Task RunAsync(CancellationToken cancellationToken)
{
void WriteJsonRpcMessageToBuffer(SseItem<IJsonRpcMessage?> item, IBufferWriter<byte> writer)
{
if (item.EventType == "endpoint")
{
writer.Write(Encoding.UTF8.GetBytes(messageEndpoint));
return;
}
JsonSerializer.Serialize(GetUtf8JsonWriter(writer), item.Data, McpJsonUtilities.DefaultOptions.GetTypeInfo<IJsonRpcMessage?>());
}
IsConnected = true;
// The very first SSE event isn't really an IJsonRpcMessage, but there's no API to write a single item of a different type,
// so we fib and special-case the "endpoint" event type in the formatter.
if (!_outgoingSseChannel.Writer.TryWrite(new SseItem<IJsonRpcMessage?>(null, "endpoint")))
@ -52,10 +39,23 @@ public sealed class SseResponseStreamTransport(Stream sseResponseStream, string
throw new InvalidOperationException($"You must call ${nameof(RunAsync)} before calling ${nameof(SendMessageAsync)}.");
}
IsConnected = true;
var sseItems = _outgoingSseChannel.Reader.ReadAllAsync(cancellationToken);
return _sseWriteTask = SseFormatter.WriteAsync(sseItems, sseResponseStream, WriteJsonRpcMessageToBuffer, cancellationToken);
}
private void WriteJsonRpcMessageToBuffer(SseItem<IJsonRpcMessage?> item, IBufferWriter<byte> writer)
{
if (item.EventType == "endpoint")
{
writer.Write(Encoding.UTF8.GetBytes(messageEndpoint));
return;
}
JsonSerializer.Serialize(GetUtf8JsonWriter(writer), item.Data, McpJsonUtilities.DefaultOptions.GetTypeInfo<IJsonRpcMessage?>());
}
/// <inheritdoc/>
public ChannelReader<IJsonRpcMessage> MessageReader => _incomingChannel.Reader;

16
src/ModelContextProtocol/Server/McpServer.cs
View file

@ -14,8 +14,8 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
private readonly EventHandler? _toolsChangedDelegate;
private readonly EventHandler? _promptsChangedDelegate;
private ITransport _sessionTransport;
private string _endpointName;
private int _started;
/// <summary>
/// Creates a new instance of <see cref="McpServer"/>.
@ -32,7 +32,6 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
Throw.IfNull(transport);
Throw.IfNull(options);
_sessionTransport = transport;
ServerOptions = options;
Services = serviceProvider;
_endpointName = $"Server ({options.ServerInfo.Name} {options.ServerInfo.Version})";
@ -74,6 +73,8 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
SetPromptsHandler(options);
SetResourcesHandler(options);
SetSetLoggingLevelHandler(options);
StartSession(transport);
}
public ServerCapabilities? ServerCapabilities { get; set; }
@ -96,11 +97,16 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
/// <inheritdoc />
public async Task RunAsync(CancellationToken cancellationToken = default)
{
if (Interlocked.Exchange(ref _started, 1) != 0)
{
throw new InvalidOperationException($"{nameof(RunAsync)} must only be called once.");
}
try
{
// Start processing messages
StartSession(_sessionTransport, fullSessionCancellationToken: cancellationToken);
await MessageProcessingTask.ConfigureAwait(false);
using var _ = cancellationToken.Register(static s => ((McpServer)s!).CancelSession(), this);
// The McpServer ctor always calls StartSession, so MessageProcessingTask is always set.
await MessageProcessingTask!.ConfigureAwait(false);
}
finally
{

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

@ -22,7 +22,6 @@ internal abstract class McpJsonRpcEndpoint : IAsyncDisposable
private McpSession? _session;
private CancellationTokenSource? _sessionCts;
private int _started;
private readonly SemaphoreSlim _disposeLock = new(1, 1);
private bool _disposed;
@ -61,18 +60,15 @@ internal abstract class McpJsonRpcEndpoint : IAsyncDisposable
protected Task? MessageProcessingTask { get; set; }
[MemberNotNull(nameof(MessageProcessingTask))]
protected void StartSession(ITransport sessionTransport, CancellationToken fullSessionCancellationToken = default)
protected void StartSession(ITransport sessionTransport)
{
if (Interlocked.Exchange(ref _started, 1) != 0)
{
throw new InvalidOperationException("The MCP session has already stared.");
}
_sessionCts = CancellationTokenSource.CreateLinkedTokenSource(fullSessionCancellationToken);
_sessionCts = new CancellationTokenSource();
_session = new McpSession(sessionTransport, EndpointName, _requestHandlers, _notificationHandlers, _logger);
MessageProcessingTask = _session.ProcessMessagesAsync(_sessionCts.Token);
}
protected void CancelSession() => _sessionCts?.Cancel();
public async ValueTask DisposeAsync()
{
using var _ = await _disposeLock.LockAsync().ConfigureAwait(false);

12
tests/ModelContextProtocol.Tests/Server/McpServerFactoryTests.cs
View file

@ -1,8 +1,6 @@
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Protocol.Types;
using ModelContextProtocol.Protocol.Types;
using ModelContextProtocol.Server;
using ModelContextProtocol.Tests.Utils;
using Moq;
namespace ModelContextProtocol.Tests.Server;
@ -25,7 +23,8 @@ public class McpServerFactoryTests : LoggedTest
public async Task Create_Should_Initialize_With_Valid_Parameters()
{
// Arrange & Act
await using IMcpServer server = McpServerFactory.Create(Mock.Of<ITransport>(), _options, LoggerFactory);
await using var transport = new TestServerTransport();
await using IMcpServer server = McpServerFactory.Create(transport, _options, LoggerFactory);
// Assert
Assert.NotNull(server);
@ -39,9 +38,10 @@ public class McpServerFactoryTests : LoggedTest
}
[Fact]
public void Create_Throws_For_Null_Options()
public async Task Create_Throws_For_Null_Options()
{
// Arrange, Act & Assert
Assert.Throws<ArgumentNullException>("serverOptions", () => McpServerFactory.Create(Mock.Of<ITransport>(), null!, LoggerFactory));
await using var transport = new TestServerTransport();
Assert.Throws<ArgumentNullException>("serverOptions", () => McpServerFactory.Create(transport, null!, LoggerFactory));
}
}

72
tests/ModelContextProtocol.Tests/Server/McpServerTests.cs
View file

@ -1,27 +1,21 @@
using Microsoft.Extensions.AI;
using Microsoft.Extensions.DependencyInjection;
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Protocol.Types;
using ModelContextProtocol.Server;
using ModelContextProtocol.Tests.Utils;
using Moq;
using System.Reflection;
namespace ModelContextProtocol.Tests.Server;
public class McpServerTests : LoggedTest
{
private readonly Mock<ITransport> _serverTransport;
private readonly McpServerOptions _options;
private readonly IServiceProvider _serviceProvider;
public McpServerTests(ITestOutputHelper testOutputHelper)
: base(testOutputHelper)
{
_serverTransport = new Mock<ITransport>();
_options = CreateOptions();
_serviceProvider = new ServiceCollection().BuildServiceProvider();
}
private static McpServerOptions CreateOptions(ServerCapabilities? capabilities = null)
@ -39,7 +33,8 @@ public class McpServerTests : LoggedTest
public async Task Constructor_Should_Initialize_With_Valid_Parameters()
{
// Arrange & Act
await using var server = McpServerFactory.Create(_serverTransport.Object, _options, LoggerFactory, _serviceProvider);
await using var transport = new TestServerTransport();
await using var server = McpServerFactory.Create(transport, _options, LoggerFactory);
// Assert
Assert.NotNull(server);
@ -49,21 +44,23 @@ public class McpServerTests : LoggedTest
public void Constructor_Throws_For_Null_Transport()
{
// Arrange, Act & Assert
Assert.Throws<ArgumentNullException>(() => McpServerFactory.Create(null!, _options, LoggerFactory, _serviceProvider));
Assert.Throws<ArgumentNullException>(() => McpServerFactory.Create(null!, _options, LoggerFactory));
}
[Fact]
public void Constructor_Throws_For_Null_Options()
public async Task Constructor_Throws_For_Null_Options()
{
// Arrange, Act & Assert
Assert.Throws<ArgumentNullException>(() => McpServerFactory.Create(_serverTransport.Object, null!, LoggerFactory, _serviceProvider));
await using var transport = new TestServerTransport();
Assert.Throws<ArgumentNullException>(() => McpServerFactory.Create(transport, null!, LoggerFactory));
}
[Fact]
public async Task Constructor_Does_Not_Throw_For_Null_Logger()
{
// Arrange & Act
await using var server = McpServerFactory.Create(_serverTransport.Object, _options, null, _serviceProvider);
await using var transport = new TestServerTransport();
await using var server = McpServerFactory.Create(transport, _options, null);
// Assert
Assert.NotNull(server);
@ -73,7 +70,8 @@ public class McpServerTests : LoggedTest
public async Task Constructor_Does_Not_Throw_For_Null_ServiceProvider()
{
// Arrange & Act
await using var server = McpServerFactory.Create(_serverTransport.Object, _options, LoggerFactory, null);
await using var transport = new TestServerTransport();
await using var server = McpServerFactory.Create(transport, _options, LoggerFactory, null);
// Assert
Assert.NotNull(server);
@ -83,27 +81,23 @@ public class McpServerTests : LoggedTest
public async Task RunAsync_Should_Throw_InvalidOperationException_If_Already_Running()
{
// Arrange
await using var server = McpServerFactory.Create(_serverTransport.Object, _options, LoggerFactory, _serviceProvider);
await using var transport = new TestServerTransport();
await using var server = McpServerFactory.Create(transport, _options, LoggerFactory);
var runTask = server.RunAsync(TestContext.Current.CancellationToken);
// Act & Assert
await Assert.ThrowsAsync<InvalidOperationException>(() => server.RunAsync(TestContext.Current.CancellationToken));
try
{
await runTask;
}
catch (NullReferenceException)
{
// _serverTransport.Object returns a null MessageReader
}
await transport.DisposeAsync();
await runTask;
}
[Fact]
public async Task RequestSamplingAsync_Should_Throw_McpServerException_If_Client_Does_Not_Support_Sampling()
{
// Arrange
await using var server = McpServerFactory.Create(_serverTransport.Object, _options, LoggerFactory, _serviceProvider);
await using var transport = new TestServerTransport();
await using var server = McpServerFactory.Create(transport, _options, LoggerFactory);
SetClientCapabilities(server, new ClientCapabilities());
var action = () => server.RequestSamplingAsync(new CreateMessageRequestParams { Messages = [] }, CancellationToken.None);
@ -117,7 +111,7 @@ public class McpServerTests : LoggedTest
{
// Arrange
await using var transport = new TestServerTransport();
await using var server = McpServerFactory.Create(transport, _options, LoggerFactory, _serviceProvider);
await using var server = McpServerFactory.Create(transport, _options, LoggerFactory);
SetClientCapabilities(server, new ClientCapabilities { Sampling = new SamplingCapability() });
var runTask = server.RunAsync(TestContext.Current.CancellationToken);
@ -138,7 +132,8 @@ public class McpServerTests : LoggedTest
public async Task RequestRootsAsync_Should_Throw_McpServerException_If_Client_Does_Not_Support_Roots()
{
// Arrange
await using var server = McpServerFactory.Create(_serverTransport.Object, _options, LoggerFactory, _serviceProvider);
await using var transport = new TestServerTransport();
await using var server = McpServerFactory.Create(transport, _options, LoggerFactory);
SetClientCapabilities(server, new ClientCapabilities());
// Act & Assert
@ -150,7 +145,7 @@ public class McpServerTests : LoggedTest
{
// Arrange
await using var transport = new TestServerTransport();
await using var server = McpServerFactory.Create(transport, _options, LoggerFactory, _serviceProvider);
await using var server = McpServerFactory.Create(transport, _options, LoggerFactory);
SetClientCapabilities(server, new ClientCapabilities { Roots = new RootsCapability() });
var runTask = server.RunAsync(TestContext.Current.CancellationToken);
@ -507,7 +502,7 @@ public class McpServerTests : LoggedTest
var options = CreateOptions(serverCapabilities);
configureOptions?.Invoke(options);
await using var server = McpServerFactory.Create(transport, options, LoggerFactory, _serviceProvider);
await using var server = McpServerFactory.Create(transport, options, LoggerFactory);
var runTask = server.RunAsync(TestContext.Current.CancellationToken);
@ -542,7 +537,7 @@ public class McpServerTests : LoggedTest
await using var transport = new TestServerTransport();
var options = CreateOptions(serverCapabilities);
Assert.Throws<McpServerException>(() => McpServerFactory.Create(transport, options, LoggerFactory, _serviceProvider));
Assert.Throws<McpServerException>(() => McpServerFactory.Create(transport, options, LoggerFactory));
}
[Fact]
@ -553,7 +548,6 @@ public class McpServerTests : LoggedTest
Assert.Throws<ArgumentException>("server", () => server.AsSamplingChatClient());
}
[Fact]
public async Task AsSamplingChatClient_HandlesRequestResponse()
{
@ -583,6 +577,26 @@ public class McpServerTests : LoggedTest
Assert.Equal(ChatRole.Assistant, response.Messages[0].Role);
}
[Fact]
public async Task Can_SendMessage_Before_RunAsync()
{
await using var transport = new TestServerTransport();
await using var server = McpServerFactory.Create(transport, _options, LoggerFactory);
var logNotification = new JsonRpcNotification()
{
Method = NotificationMethods.LoggingMessageNotification
};
await server.SendMessageAsync(logNotification, TestContext.Current.CancellationToken);
var runTask = server.RunAsync(TestContext.Current.CancellationToken);
await transport.DisposeAsync();
await runTask;
Assert.NotEmpty(transport.SentMessages);
Assert.Same(logNotification, transport.SentMessages[0]);
}
private static void SetClientCapabilities(IMcpServer server, ClientCapabilities capabilities)
{
PropertyInfo? property = server.GetType().GetProperty("ClientCapabilities", BindingFlags.Public | BindingFlags.Instance);
@ -644,7 +658,7 @@ public class McpServerTests : LoggedTest
var notificationReceived = new TaskCompletionSource<JsonRpcNotification>();
var server = McpServerFactory.Create(transport, options, LoggerFactory, _serviceProvider);
var server = McpServerFactory.Create(transport, options, LoggerFactory);
server.AddNotificationHandler(NotificationMethods.ProgressNotification, notification =>
{
notificationReceived.SetResult(notification);