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

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
4 changed files with 94 additions and 80 deletions
Download patch file Download diff file Expand all files Collapse all files

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);