Refactor transports to help enable graceful shutdown (#142)
- 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
GitHub
parent
d47c834fa9
commit
2ab9db0878
2 changed files with 2015 additions and 1847 deletions
|
|
@ -1,11 +1,11 @@
|
|||
using ModelContextProtocol.Configuration;
|
||||
using Microsoft.Extensions.Logging;
|
||||
using ModelContextProtocol.Configuration;
|
||||
using ModelContextProtocol.Logging;
|
||||
using ModelContextProtocol.Protocol.Messages;
|
||||
using ModelContextProtocol.Protocol.Transport;
|
||||
using ModelContextProtocol.Protocol.Types;
|
||||
using ModelContextProtocol.Shared;
|
||||
using ModelContextProtocol.Utils.Json;
|
||||
using Microsoft.Extensions.Logging;
|
||||
using System.Text.Json;
|
||||
|
||||
namespace ModelContextProtocol.Client;
|
||||
|
|
@ -13,23 +13,24 @@ namespace ModelContextProtocol.Client;
|
|||
/// <inheritdoc/>
|
||||
internal sealed class McpClient : McpJsonRpcEndpoint, IMcpClient
|
||||
{
|
||||
private readonly McpClientOptions _options;
|
||||
private readonly IClientTransport _clientTransport;
|
||||
private readonly McpClientOptions _options;
|
||||
|
||||
private volatile bool _isInitializing;
|
||||
private ITransport? _sessionTransport;
|
||||
private CancellationTokenSource? _connectCts;
|
||||
|
||||
/// <summary>
|
||||
/// Initializes a new instance of the <see cref="McpClient"/> class.
|
||||
/// </summary>
|
||||
/// <param name="transport">The transport to use for communication with the server.</param>
|
||||
/// <param name="clientTransport">The transport to use for communication with the server.</param>
|
||||
/// <param name="options">Options for the client, defining protocol version and capabilities.</param>
|
||||
/// <param name="serverConfig">The server configuration.</param>
|
||||
/// <param name="loggerFactory">The logger factory.</param>
|
||||
public McpClient(IClientTransport transport, McpClientOptions options, McpServerConfig serverConfig, ILoggerFactory? loggerFactory)
|
||||
: base(transport, loggerFactory)
|
||||
public McpClient(IClientTransport clientTransport, McpClientOptions options, McpServerConfig serverConfig, ILoggerFactory? loggerFactory)
|
||||
: base(loggerFactory)
|
||||
{
|
||||
_clientTransport = clientTransport;
|
||||
_options = options;
|
||||
_clientTransport = transport;
|
||||
|
||||
EndpointName = $"Client ({serverConfig.Id}: {serverConfig.Name})";
|
||||
|
||||
|
|
@ -70,95 +71,95 @@ internal sealed class McpClient : McpJsonRpcEndpoint, IMcpClient
|
|||
/// <inheritdoc/>
|
||||
public override string EndpointName { get; }
|
||||
|
||||
/// <inheritdoc/>
|
||||
public async Task ConnectAsync(CancellationToken cancellationToken = default)
|
||||
{
|
||||
if (IsInitialized)
|
||||
{
|
||||
_logger.ClientAlreadyInitialized(EndpointName);
|
||||
return;
|
||||
}
|
||||
_connectCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
||||
cancellationToken = _connectCts.Token;
|
||||
|
||||
if (_isInitializing)
|
||||
{
|
||||
_logger.ClientAlreadyInitializing(EndpointName);
|
||||
throw new InvalidOperationException("Client is already initializing");
|
||||
}
|
||||
|
||||
_isInitializing = true;
|
||||
try
|
||||
{
|
||||
CancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
||||
|
||||
// Connect transport
|
||||
await _clientTransport.ConnectAsync(CancellationTokenSource.Token).ConfigureAwait(false);
|
||||
|
||||
// Start processing messages
|
||||
MessageProcessingTask = ProcessMessagesAsync(CancellationTokenSource.Token);
|
||||
_sessionTransport = await _clientTransport.ConnectAsync(cancellationToken).ConfigureAwait(false);
|
||||
InitializeSession(_sessionTransport);
|
||||
// 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(fullSessionCancellationToken: CancellationToken.None);
|
||||
|
||||
// Perform initialization sequence
|
||||
await InitializeAsync(CancellationTokenSource.Token).ConfigureAwait(false);
|
||||
using var initializationCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
||||
initializationCts.CancelAfter(_options.InitializationTimeout);
|
||||
|
||||
IsInitialized = true;
|
||||
try
|
||||
{
|
||||
// Send initialize request
|
||||
var initializeResponse = await SendRequestAsync<InitializeResult>(
|
||||
new JsonRpcRequest
|
||||
{
|
||||
Method = "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,
|
||||
capabilities: JsonSerializer.Serialize(initializeResponse.Capabilities, McpJsonUtilities.JsonContext.Default.ServerCapabilities),
|
||||
serverInfo: JsonSerializer.Serialize(initializeResponse.ServerInfo, McpJsonUtilities.JsonContext.Default.Implementation));
|
||||
|
||||
ServerCapabilities = initializeResponse.Capabilities;
|
||||
ServerInfo = initializeResponse.ServerInfo;
|
||||
ServerInstructions = initializeResponse.Instructions;
|
||||
|
||||
// Validate protocol version
|
||||
if (initializeResponse.ProtocolVersion != _options.ProtocolVersion)
|
||||
{
|
||||
_logger.ServerProtocolVersionMismatch(EndpointName, _options.ProtocolVersion, initializeResponse.ProtocolVersion);
|
||||
throw new McpClientException($"Server protocol version mismatch. Expected {_options.ProtocolVersion}, got {initializeResponse.ProtocolVersion}");
|
||||
}
|
||||
|
||||
// Send initialized notification
|
||||
await SendMessageAsync(
|
||||
new JsonRpcNotification { Method = "notifications/initialized" },
|
||||
initializationCts.Token).ConfigureAwait(false);
|
||||
}
|
||||
catch (OperationCanceledException) when (initializationCts.IsCancellationRequested)
|
||||
{
|
||||
_logger.ClientInitializationTimeout(EndpointName);
|
||||
throw new McpClientException("Initialization timed out");
|
||||
}
|
||||
}
|
||||
catch (Exception e)
|
||||
{
|
||||
_logger.ClientInitializationError(EndpointName, e);
|
||||
await CleanupAsync().ConfigureAwait(false);
|
||||
await DisposeAsync().ConfigureAwait(false);
|
||||
throw;
|
||||
}
|
||||
finally
|
||||
{
|
||||
_isInitializing = false;
|
||||
}
|
||||
}
|
||||
|
||||
private async Task InitializeAsync(CancellationToken cancellationToken)
|
||||
/// <inheritdoc/>
|
||||
public override async ValueTask DisposeUnsynchronizedAsync()
|
||||
{
|
||||
using var initializationCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
||||
initializationCts.CancelAfter(_options.InitializationTimeout);
|
||||
if (_connectCts is not null)
|
||||
{
|
||||
await _connectCts.CancelAsync().ConfigureAwait(false);
|
||||
}
|
||||
|
||||
try
|
||||
{
|
||||
// Send initialize request
|
||||
var initializeResponse = await SendRequestAsync<InitializeResult>(
|
||||
new JsonRpcRequest
|
||||
{
|
||||
Method = "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,
|
||||
capabilities: JsonSerializer.Serialize(initializeResponse.Capabilities, McpJsonUtilities.JsonContext.Default.ServerCapabilities),
|
||||
serverInfo: JsonSerializer.Serialize(initializeResponse.ServerInfo, McpJsonUtilities.JsonContext.Default.Implementation));
|
||||
|
||||
ServerCapabilities = initializeResponse.Capabilities;
|
||||
ServerInfo = initializeResponse.ServerInfo;
|
||||
ServerInstructions = initializeResponse.Instructions;
|
||||
|
||||
// Validate protocol version
|
||||
if (initializeResponse.ProtocolVersion != _options.ProtocolVersion)
|
||||
await base.DisposeUnsynchronizedAsync().ConfigureAwait(false);
|
||||
}
|
||||
finally
|
||||
{
|
||||
if (_sessionTransport is not null)
|
||||
{
|
||||
_logger.ServerProtocolVersionMismatch(EndpointName, _options.ProtocolVersion, initializeResponse.ProtocolVersion);
|
||||
throw new McpClientException($"Server protocol version mismatch. Expected {_options.ProtocolVersion}, got {initializeResponse.ProtocolVersion}");
|
||||
await _sessionTransport.DisposeAsync().ConfigureAwait(false);
|
||||
}
|
||||
|
||||
// Send initialized notification
|
||||
await SendMessageAsync(
|
||||
new JsonRpcNotification { Method = "notifications/initialized" },
|
||||
initializationCts.Token).ConfigureAwait(false);
|
||||
}
|
||||
catch (OperationCanceledException) when (initializationCts.IsCancellationRequested)
|
||||
{
|
||||
_logger.ClientInitializationTimeout(EndpointName);
|
||||
throw new McpClientException("Initialization timed out");
|
||||
_connectCts?.Dispose();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -82,7 +82,14 @@ public static class McpClientFactory
|
|||
}
|
||||
catch
|
||||
{
|
||||
await transport.DisposeAsync().ConfigureAwait(false);
|
||||
if (transport is IAsyncDisposable asyncDisposableTransport)
|
||||
{
|
||||
await asyncDisposableTransport.DisposeAsync().ConfigureAwait(false);
|
||||
}
|
||||
else if (transport is IDisposable disposableTransport)
|
||||
{
|
||||
disposableTransport.Dispose();
|
||||
}
|
||||
throw;
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue