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
3 changed files with 2015 additions and 1847 deletions
|
|
@ -8,11 +8,6 @@ namespace ModelContextProtocol.Server;
|
|||
/// </summary>
|
||||
public interface IMcpServer : IAsyncDisposable
|
||||
{
|
||||
/// <summary>
|
||||
/// Gets a value indicating whether the server has been initialized.
|
||||
/// </summary>
|
||||
bool IsInitialized { get; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets the capabilities supported by the client.
|
||||
/// </summary>
|
||||
|
|
@ -48,23 +43,23 @@ public interface IMcpServer : IAsyncDisposable
|
|||
void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler);
|
||||
|
||||
/// <summary>
|
||||
/// Starts the server and begins listening for client requests.
|
||||
/// Runs the server, listening for and handling client requests.
|
||||
/// </summary>
|
||||
Task StartAsync(CancellationToken cancellationToken = default);
|
||||
Task RunAsync(CancellationToken cancellationToken = default);
|
||||
|
||||
/// <summary>
|
||||
/// Sends a generic JSON-RPC request to the client.
|
||||
/// NB! This is a temporary method that is available to send not yet implemented feature messages.
|
||||
/// NB! This is a temporary method that is available to send not yet implemented feature messages.
|
||||
/// Once all MCP features are implemented this will be made private, as it is purely a convenience for those who wish to implement features ahead of the library.
|
||||
/// </summary>
|
||||
/// <typeparam name="T">The expected response type.</typeparam>
|
||||
/// <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 client's response.</returns>
|
||||
Task<T> SendRequestAsync<T>(JsonRpcRequest request, CancellationToken cancellationToken) where T : class;
|
||||
Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class;
|
||||
|
||||
/// <summary>
|
||||
/// Sends a message to the server.
|
||||
/// Sends a message to the client.
|
||||
/// </summary>
|
||||
/// <param name="message">The message.</param>
|
||||
/// <param name="cancellationToken">A token to cancel the operation.</param>
|
||||
|
|
|
|||
|
|
@ -13,29 +13,63 @@ namespace ModelContextProtocol.Server;
|
|||
internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
|
||||
{
|
||||
private readonly IServerTransport? _serverTransport;
|
||||
private readonly string _serverDescription;
|
||||
private readonly EventHandler? _toolsChangedDelegate;
|
||||
|
||||
private volatile bool _isInitializing;
|
||||
private ITransport? _sessionTransport;
|
||||
private string _endpointName;
|
||||
|
||||
/// <summary>
|
||||
/// Creates a new instance of <see cref="McpServer"/>.
|
||||
/// </summary>
|
||||
/// <param name="transport">Transport to use for the server</param>
|
||||
/// <param name="options">Configuration options for this server, including capabilities.
|
||||
/// Make sure to accurately reflect exactly what capabilities the server supports and does not support.</param>
|
||||
/// <param name="serverTransport">Transport to use for the server that is ready to accept new sessions asynchronously.</param>
|
||||
/// <param name="options">Configuration options for this server, including capabilities.
|
||||
/// Make sure to accurately reflect exactly what capabilities the server supports and does not support.</param>
|
||||
/// <param name="loggerFactory">Logger factory to use for logging</param>
|
||||
/// <param name="serviceProvider">Optional service provider to use for dependency injection</param>
|
||||
/// <exception cref="McpServerException"></exception>
|
||||
public McpServer(IServerTransport serverTransport, McpServerOptions options, ILoggerFactory? loggerFactory, IServiceProvider? serviceProvider)
|
||||
: this(options, loggerFactory, serviceProvider)
|
||||
{
|
||||
Throw.IfNull(serverTransport);
|
||||
|
||||
_serverTransport = serverTransport;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Creates a new instance of <see cref="McpServer"/>.
|
||||
/// </summary>
|
||||
/// <param name="transport">Transport to use for the server representing an already-established session.</param>
|
||||
/// <param name="options">Configuration options for this server, including capabilities.
|
||||
/// Make sure to accurately reflect exactly what capabilities the server supports and does not support.</param>
|
||||
/// <param name="loggerFactory">Logger factory to use for logging</param>
|
||||
/// <param name="serviceProvider">Optional service provider to use for dependency injection</param>
|
||||
/// <exception cref="McpServerException"></exception>
|
||||
public McpServer(ITransport transport, McpServerOptions options, ILoggerFactory? loggerFactory, IServiceProvider? serviceProvider)
|
||||
: base(transport, loggerFactory)
|
||||
: this(options, loggerFactory, serviceProvider)
|
||||
{
|
||||
Throw.IfNull(transport);
|
||||
|
||||
_sessionTransport = transport;
|
||||
InitializeSession(transport);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Creates a new instance of <see cref="McpServer"/>.
|
||||
/// </summary>
|
||||
/// <param name="options">Configuration options for this server, including capabilities.
|
||||
/// Make sure to accurately reflect exactly what capabilities the server supports and does not support.</param>
|
||||
/// <param name="loggerFactory">Logger factory to use for logging</param>
|
||||
/// <param name="serviceProvider">Optional service provider to use for dependency injection</param>
|
||||
/// <exception cref="McpServerException"></exception>
|
||||
private McpServer(McpServerOptions options, ILoggerFactory? loggerFactory, IServiceProvider? serviceProvider)
|
||||
: base(loggerFactory)
|
||||
{
|
||||
Throw.IfNull(options);
|
||||
|
||||
_serverTransport = transport as IServerTransport;
|
||||
ServerOptions = options;
|
||||
Services = serviceProvider;
|
||||
_serverDescription = $"{options.ServerInfo.Name} {options.ServerInfo.Version}";
|
||||
_endpointName = $"Server ({options.ServerInfo.Name} {options.ServerInfo.Version})";
|
||||
|
||||
_toolsChangedDelegate = delegate
|
||||
{
|
||||
_ = SendMessageAsync(new JsonRpcNotification()
|
||||
|
|
@ -51,7 +85,6 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
|
|||
tools.Changed += _toolsChangedDelegate;
|
||||
}
|
||||
|
||||
IsInitialized = true;
|
||||
return Task.CompletedTask;
|
||||
});
|
||||
|
||||
|
|
@ -67,6 +100,7 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
|
|||
|
||||
public ServerCapabilities? ServerCapabilities { get; set; }
|
||||
|
||||
/// <inheritdoc />
|
||||
public ClientCapabilities? ClientCapabilities { get; set; }
|
||||
|
||||
/// <inheritdoc />
|
||||
|
|
@ -79,55 +113,65 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
|
|||
public IServiceProvider? Services { get; }
|
||||
|
||||
/// <inheritdoc />
|
||||
public override string EndpointName =>
|
||||
$"Server ({_serverDescription}), Client ({ClientInfo?.Name} {ClientInfo?.Version})";
|
||||
public override string EndpointName => _endpointName;
|
||||
|
||||
/// <inheritdoc />
|
||||
public async Task StartAsync(CancellationToken cancellationToken = default)
|
||||
public async Task AcceptSessionAsync(CancellationToken cancellationToken = default)
|
||||
{
|
||||
if (_isInitializing)
|
||||
{
|
||||
_logger.ServerAlreadyInitializing(EndpointName);
|
||||
throw new InvalidOperationException("Server is already initializing");
|
||||
}
|
||||
_isInitializing = true;
|
||||
|
||||
if (IsInitialized)
|
||||
{
|
||||
_logger.ServerAlreadyInitializing(EndpointName);
|
||||
return;
|
||||
}
|
||||
// Below is effectively an assertion. The McpServerFactory should only use this with the IServerTransport constructor.
|
||||
Throw.IfNull(_serverTransport);
|
||||
|
||||
try
|
||||
{
|
||||
CancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
||||
_sessionTransport = await _serverTransport.AcceptAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
if (_serverTransport is not null)
|
||||
if (_sessionTransport is null)
|
||||
{
|
||||
// Start listening for messages
|
||||
await _serverTransport.StartListeningAsync(CancellationTokenSource.Token).ConfigureAwait(false);
|
||||
throw new McpServerException("The server transport closed before a client started a new session.");
|
||||
}
|
||||
|
||||
// Start processing messages
|
||||
MessageProcessingTask = ProcessMessagesAsync(CancellationTokenSource.Token);
|
||||
|
||||
// Unlike McpClient, we're not done initializing until we've received a message from the client, so we don't set IsInitialized here
|
||||
InitializeSession(_sessionTransport);
|
||||
}
|
||||
catch (Exception e)
|
||||
{
|
||||
_logger.ServerInitializationError(EndpointName, e);
|
||||
await CleanupAsync().ConfigureAwait(false);
|
||||
throw;
|
||||
}
|
||||
}
|
||||
|
||||
protected override Task CleanupAsync()
|
||||
/// <inheritdoc />
|
||||
public async Task RunAsync(CancellationToken cancellationToken = default)
|
||||
{
|
||||
try
|
||||
{
|
||||
// Start processing messages
|
||||
StartSession(fullSessionCancellationToken: cancellationToken);
|
||||
await MessageProcessingTask.ConfigureAwait(false);
|
||||
}
|
||||
finally
|
||||
{
|
||||
await DisposeAsync().ConfigureAwait(false);
|
||||
}
|
||||
}
|
||||
|
||||
public override async ValueTask DisposeUnsynchronizedAsync()
|
||||
{
|
||||
if (ServerOptions.Capabilities?.Tools?.ToolCollection is { } tools)
|
||||
{
|
||||
tools.Changed -= _toolsChangedDelegate;
|
||||
}
|
||||
return base.CleanupAsync();
|
||||
|
||||
try
|
||||
{
|
||||
await base.DisposeUnsynchronizedAsync().ConfigureAwait(false);
|
||||
}
|
||||
finally
|
||||
{
|
||||
if (_serverTransport is not null && _sessionTransport is not null)
|
||||
{
|
||||
// We created the _sessionTransport from the _serverTransport, so we own it.
|
||||
await _sessionTransport.DisposeAsync().ConfigureAwait(false);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private void SetPingHandler()
|
||||
|
|
@ -143,6 +187,11 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
|
|||
{
|
||||
ClientCapabilities = request?.Capabilities ?? new();
|
||||
ClientInfo = request?.ClientInfo;
|
||||
|
||||
// Use the ClientInfo to update the session EndpointName for logging.
|
||||
_endpointName = $"{_endpointName}, Client ({ClientInfo?.Name} {ClientInfo?.Version})";
|
||||
GetSessionOrThrow().EndpointName = EndpointName;
|
||||
|
||||
return Task.FromResult(new InitializeResult()
|
||||
{
|
||||
ProtocolVersion = options.ProtocolVersion,
|
||||
|
|
|
|||
|
|
@ -10,27 +10,65 @@ namespace ModelContextProtocol.Server;
|
|||
public static class McpServerFactory
|
||||
{
|
||||
/// <summary>
|
||||
/// Initializes a new instance of the <see cref="McpServerFactory"/> class.
|
||||
/// Initializes a new instance of the <see cref="IMcpServer"/> class.
|
||||
/// </summary>
|
||||
/// <param name="serverTransport">Transport to use for the server</param>
|
||||
/// <param name="transport">Transport to use for the server representing an already-established MCP session.</param>
|
||||
/// <param name="serverOptions">
|
||||
/// Configuration options for this server, including capabilities.
|
||||
/// Make sure to accurately reflect exactly what capabilities the server supports and does not support.
|
||||
/// </param>
|
||||
/// <param name="serviceProvider">Optional service provider to create new instances.</param>
|
||||
/// <param name="loggerFactory">Logger factory to use for logging</param>
|
||||
/// <returns>An <see cref="McpServer"/> that's started and ready to receive connections.</returns>
|
||||
/// <exception cref="ArgumentNullException"><paramref name="serverTransport"/> is <see langword="null"/>.</exception>
|
||||
/// <param name="serviceProvider">Optional service provider to create new instances.</param>
|
||||
/// <returns>An <see cref="IMcpServer"/>.</returns>
|
||||
/// <exception cref="ArgumentNullException"><paramref name="transport"/> is <see langword="null"/>.</exception>
|
||||
/// <exception cref="ArgumentNullException"><paramref name="serverOptions"/> is <see langword="null"/>.</exception>
|
||||
public static IMcpServer Create(
|
||||
ITransport serverTransport,
|
||||
ITransport transport,
|
||||
McpServerOptions serverOptions,
|
||||
ILoggerFactory? loggerFactory = null,
|
||||
IServiceProvider? serviceProvider = null)
|
||||
{
|
||||
Throw.IfNull(transport);
|
||||
Throw.IfNull(serverOptions);
|
||||
|
||||
return new McpServer(transport, serverOptions, loggerFactory, serviceProvider);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Waits for the client to establish a new MCP session, then initializes a new instance of the <see cref="IMcpServer"/> class.
|
||||
/// </summary>
|
||||
/// <param name="serverTransport">Transport to use for the server that is ready to accept new MCP sessions asynchronously.</param>
|
||||
/// <param name="serverOptions">
|
||||
/// Configuration options for this server, including capabilities.
|
||||
/// Make sure to accurately reflect exactly what capabilities the server supports and does not support.
|
||||
/// </param>
|
||||
/// <param name="loggerFactory">Logger factory to use for logging</param>
|
||||
/// <param name="serviceProvider">Optional service provider to create new instances.</param>
|
||||
/// <param name="cancellationToken">Cancel waiting for a client to establish a new MCP session.</param>
|
||||
/// <returns>An <see cref="IMcpServer"/>.</returns>
|
||||
/// <exception cref="ArgumentNullException"><paramref name="serverTransport"/> is <see langword="null"/>.</exception>
|
||||
/// <exception cref="ArgumentNullException"><paramref name="serverOptions"/> is <see langword="null"/>.</exception>
|
||||
public static async Task<IMcpServer> AcceptAsync(
|
||||
IServerTransport serverTransport,
|
||||
McpServerOptions serverOptions,
|
||||
ILoggerFactory? loggerFactory = null,
|
||||
IServiceProvider? serviceProvider = null,
|
||||
CancellationToken cancellationToken = default)
|
||||
{
|
||||
Throw.IfNull(serverTransport);
|
||||
Throw.IfNull(serverOptions);
|
||||
|
||||
return new McpServer(serverTransport, serverOptions, loggerFactory, serviceProvider);
|
||||
var mcpServer = new McpServer(serverTransport, serverOptions, loggerFactory, serviceProvider);
|
||||
|
||||
try
|
||||
{
|
||||
await mcpServer.AcceptSessionAsync(cancellationToken).ConfigureAwait(false);
|
||||
return mcpServer;
|
||||
}
|
||||
catch
|
||||
{
|
||||
await mcpServer.DisposeAsync().ConfigureAwait(false);
|
||||
throw;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue