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
1 changed files with 2015 additions and 1847 deletions
|
|
@ -1,36 +1,20 @@
|
|||
using System.Net.Http.Headers;
|
||||
using System.Net.ServerSentEvents;
|
||||
using System.Text;
|
||||
using System.Text.Json;
|
||||
using Microsoft.Extensions.Logging;
|
||||
using ModelContextProtocol.Configuration;
|
||||
using ModelContextProtocol.Logging;
|
||||
using ModelContextProtocol.Protocol.Messages;
|
||||
using ModelContextProtocol.Utils;
|
||||
using ModelContextProtocol.Utils.Json;
|
||||
using Microsoft.Extensions.Logging;
|
||||
using Microsoft.Extensions.Logging.Abstractions;
|
||||
|
||||
namespace ModelContextProtocol.Protocol.Transport;
|
||||
|
||||
/// <summary>
|
||||
/// The ServerSideEvents client transport implementation
|
||||
/// </summary>
|
||||
public sealed class SseClientTransport : TransportBase, IClientTransport
|
||||
public sealed class SseClientTransport : IClientTransport, IAsyncDisposable
|
||||
{
|
||||
private readonly HttpClient _httpClient;
|
||||
private readonly SseClientTransportOptions _options;
|
||||
private readonly Uri _sseEndpoint;
|
||||
private Uri? _messageEndpoint;
|
||||
private CancellationTokenSource? _connectionCts;
|
||||
private Task? _receiveTask;
|
||||
private readonly ILogger _logger;
|
||||
private readonly McpServerConfig _serverConfig;
|
||||
private readonly JsonSerializerOptions _jsonOptions;
|
||||
private readonly TaskCompletionSource<bool> _connectionEstablished;
|
||||
private readonly HttpClient _httpClient;
|
||||
private readonly ILoggerFactory? _loggerFactory;
|
||||
private readonly bool _ownsHttpClient;
|
||||
|
||||
private string EndpointName => $"Client (SSE) for ({_serverConfig.Id}: {_serverConfig.Name})";
|
||||
|
||||
/// <summary>
|
||||
/// SSE transport for client endpoints. Unlike stdio it does not launch a process, but connects to an existing server.
|
||||
/// The HTTP server can be local or remote, and must support the SSE protocol.
|
||||
|
|
@ -53,7 +37,6 @@ public sealed class SseClientTransport : TransportBase, IClientTransport
|
|||
/// <param name="loggerFactory">Logger factory for creating loggers.</param>
|
||||
/// <param name="ownsHttpClient">True to dispose HTTP client on close connection.</param>
|
||||
public SseClientTransport(SseClientTransportOptions transportOptions, McpServerConfig serverConfig, HttpClient httpClient, ILoggerFactory? loggerFactory, bool ownsHttpClient = false)
|
||||
: base(loggerFactory)
|
||||
{
|
||||
Throw.IfNull(transportOptions);
|
||||
Throw.IfNull(serverConfig);
|
||||
|
|
@ -61,281 +44,36 @@ public sealed class SseClientTransport : TransportBase, IClientTransport
|
|||
|
||||
_options = transportOptions;
|
||||
_serverConfig = serverConfig;
|
||||
_sseEndpoint = new Uri(serverConfig.Location!);
|
||||
_httpClient = httpClient;
|
||||
_connectionCts = new CancellationTokenSource();
|
||||
_logger = (ILogger?)loggerFactory?.CreateLogger<SseClientTransport>() ?? NullLogger.Instance;
|
||||
_jsonOptions = McpJsonUtilities.DefaultOptions;
|
||||
_connectionEstablished = new TaskCompletionSource<bool>();
|
||||
_loggerFactory = loggerFactory;
|
||||
_ownsHttpClient = ownsHttpClient;
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public async Task ConnectAsync(CancellationToken cancellationToken = default)
|
||||
/// <inheritdoc />
|
||||
public async Task<ITransport> ConnectAsync(CancellationToken cancellationToken = default)
|
||||
{
|
||||
var sessionTransport = new SseClientSessionTransport(_options, _serverConfig, _httpClient, _loggerFactory);
|
||||
|
||||
try
|
||||
{
|
||||
if (IsConnected)
|
||||
{
|
||||
_logger.TransportAlreadyConnected(EndpointName);
|
||||
throw new McpTransportException("Transport is already connected");
|
||||
}
|
||||
|
||||
// Start message receiving loop
|
||||
_receiveTask = ReceiveMessagesAsync(_connectionCts!.Token);
|
||||
|
||||
_logger.TransportReadingMessages(EndpointName);
|
||||
|
||||
await _connectionEstablished.Task.WaitAsync(_options.ConnectionTimeout, cancellationToken).ConfigureAwait(false);
|
||||
await sessionTransport.ConnectAsync(cancellationToken).ConfigureAwait(false);
|
||||
return sessionTransport;
|
||||
}
|
||||
catch (McpTransportException)
|
||||
catch
|
||||
{
|
||||
// Rethrow transport exceptions
|
||||
await sessionTransport.DisposeAsync().ConfigureAwait(false);
|
||||
throw;
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
_logger.TransportConnectFailed(EndpointName, ex);
|
||||
await CloseAsync().ConfigureAwait(false);
|
||||
throw new McpTransportException("Failed to connect transport", ex);
|
||||
}
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public override async Task SendMessageAsync(
|
||||
IJsonRpcMessage message,
|
||||
CancellationToken cancellationToken = default)
|
||||
/// <inheritdoc />
|
||||
public ValueTask DisposeAsync()
|
||||
{
|
||||
if (_messageEndpoint == null)
|
||||
throw new InvalidOperationException("Transport not connected");
|
||||
|
||||
using var content = new StringContent(
|
||||
JsonSerializer.Serialize(message, _jsonOptions.GetTypeInfo<IJsonRpcMessage>()),
|
||||
Encoding.UTF8,
|
||||
"application/json"
|
||||
);
|
||||
|
||||
string messageId = "(no id)";
|
||||
|
||||
if (message is IJsonRpcMessageWithId messageWithId)
|
||||
{
|
||||
messageId = messageWithId.Id.ToString();
|
||||
}
|
||||
|
||||
var response = await _httpClient.PostAsync(
|
||||
_messageEndpoint,
|
||||
content,
|
||||
cancellationToken
|
||||
).ConfigureAwait(false);
|
||||
|
||||
response.EnsureSuccessStatusCode();
|
||||
|
||||
var responseContent = await response.Content.ReadAsStringAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// Check if the message was an initialize request
|
||||
if (message is JsonRpcRequest request && request.Method == "initialize")
|
||||
{
|
||||
// If the response is not a JSON-RPC response, it is an SSE message
|
||||
if (responseContent.Equals("accepted", StringComparison.OrdinalIgnoreCase))
|
||||
{
|
||||
_logger.SSETransportPostAccepted(EndpointName, messageId);
|
||||
// The response will arrive as an SSE message
|
||||
}
|
||||
else
|
||||
{
|
||||
JsonRpcResponse initializeResponse = JsonSerializer.Deserialize(responseContent, _jsonOptions.GetTypeInfo<JsonRpcResponse>()) ??
|
||||
throw new McpTransportException("Failed to initialize client");
|
||||
|
||||
_logger.TransportReceivedMessageParsed(EndpointName, messageId);
|
||||
await WriteMessageAsync(initializeResponse, cancellationToken).ConfigureAwait(false);
|
||||
_logger.TransportMessageWritten(EndpointName, messageId);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// Otherwise, check if the response was accepted (the response will come as an SSE message)
|
||||
if (responseContent.Equals("accepted", StringComparison.OrdinalIgnoreCase))
|
||||
{
|
||||
_logger.SSETransportPostAccepted(EndpointName, messageId);
|
||||
}
|
||||
else
|
||||
{
|
||||
_logger.SSETransportPostNotAccepted(EndpointName, messageId, responseContent);
|
||||
throw new McpTransportException("Failed to send message");
|
||||
}
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public async Task CloseAsync()
|
||||
{
|
||||
if (_connectionCts != null)
|
||||
{
|
||||
await _connectionCts.CancelAsync().ConfigureAwait(false);
|
||||
_connectionCts.Dispose();
|
||||
_connectionCts = null;
|
||||
}
|
||||
if (_receiveTask != null)
|
||||
await _receiveTask.ConfigureAwait(false);
|
||||
|
||||
SetConnected(false);
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public override async ValueTask DisposeAsync()
|
||||
{
|
||||
try
|
||||
{
|
||||
await CloseAsync().ConfigureAwait(false);
|
||||
}
|
||||
catch (Exception)
|
||||
{
|
||||
// Ignore exceptions on close
|
||||
}
|
||||
|
||||
if (_ownsHttpClient)
|
||||
_httpClient?.Dispose();
|
||||
|
||||
_connectionCts?.Dispose();
|
||||
|
||||
GC.SuppressFinalize(this);
|
||||
}
|
||||
|
||||
internal Uri? MessageEndpoint => _messageEndpoint;
|
||||
|
||||
internal SseClientTransportOptions Options => _options;
|
||||
|
||||
private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
int reconnectAttempts = 0;
|
||||
|
||||
while (!cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
try
|
||||
{
|
||||
using var request = new HttpRequestMessage(HttpMethod.Get, _sseEndpoint);
|
||||
request.Headers.Accept.Add(new MediaTypeWithQualityHeaderValue("text/event-stream"));
|
||||
|
||||
using var response = await _httpClient.SendAsync(
|
||||
request,
|
||||
HttpCompletionOption.ResponseHeadersRead,
|
||||
cancellationToken
|
||||
).ConfigureAwait(false);
|
||||
|
||||
response.EnsureSuccessStatusCode();
|
||||
|
||||
using var stream = await response.Content.ReadAsStreamAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// Reset reconnect attempts on successful connection
|
||||
reconnectAttempts = 0;
|
||||
|
||||
await foreach (SseItem<string> sseEvent in SseParser.Create(stream).EnumerateAsync(cancellationToken).ConfigureAwait(false))
|
||||
{
|
||||
switch (sseEvent.EventType)
|
||||
{
|
||||
case "endpoint":
|
||||
HandleEndpointEvent(sseEvent.Data);
|
||||
break;
|
||||
|
||||
case "message":
|
||||
await ProcessSseMessage(sseEvent.Data, cancellationToken).ConfigureAwait(false);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
_logger.TransportReadMessagesCancelled(EndpointName);
|
||||
// Normal shutdown
|
||||
}
|
||||
catch (IOException) when (cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
_logger.TransportReadMessagesCancelled(EndpointName);
|
||||
// Normal shutdown
|
||||
}
|
||||
catch (Exception ex) when (!cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
_logger.TransportConnectionError(EndpointName, ex);
|
||||
|
||||
reconnectAttempts++;
|
||||
if (reconnectAttempts >= _options.MaxReconnectAttempts)
|
||||
{
|
||||
throw new McpTransportException("Exceeded reconnect limit", ex);
|
||||
}
|
||||
|
||||
await Task.Delay(_options.ReconnectDelay, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async Task ProcessSseMessage(string data, CancellationToken cancellationToken)
|
||||
{
|
||||
if (!IsConnected)
|
||||
{
|
||||
_logger.TransportMessageReceivedBeforeConnected(EndpointName, data);
|
||||
return;
|
||||
_httpClient.Dispose();
|
||||
}
|
||||
|
||||
try
|
||||
{
|
||||
var message = JsonSerializer.Deserialize(data, _jsonOptions.GetTypeInfo<IJsonRpcMessage>());
|
||||
if (message == null)
|
||||
{
|
||||
_logger.TransportMessageParseUnexpectedType(EndpointName, data);
|
||||
return;
|
||||
}
|
||||
|
||||
string messageId = "(no id)";
|
||||
if (message is IJsonRpcMessageWithId messageWithId)
|
||||
{
|
||||
messageId = messageWithId.Id.ToString();
|
||||
}
|
||||
|
||||
_logger.TransportReceivedMessageParsed(EndpointName, messageId);
|
||||
await WriteMessageAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
_logger.TransportMessageWritten(EndpointName, messageId);
|
||||
}
|
||||
catch (JsonException ex)
|
||||
{
|
||||
_logger.TransportMessageParseFailed(EndpointName, data, ex);
|
||||
}
|
||||
}
|
||||
|
||||
private void HandleEndpointEvent(string data)
|
||||
{
|
||||
try
|
||||
{
|
||||
if (string.IsNullOrEmpty(data))
|
||||
{
|
||||
_logger.TransportEndpointEventInvalid(EndpointName, data);
|
||||
return;
|
||||
}
|
||||
|
||||
// Check if data is absolute URI
|
||||
if (data.StartsWith("http://", StringComparison.OrdinalIgnoreCase) || data.StartsWith("https://", StringComparison.OrdinalIgnoreCase))
|
||||
{
|
||||
// Since the endpoint is an absolute URI, we can use it directly
|
||||
_messageEndpoint = new Uri(data);
|
||||
}
|
||||
else
|
||||
{
|
||||
// If the endpoint is a relative URI, we need to combine it with the relative path of the SSE endpoint
|
||||
var hostUrl = _sseEndpoint.AbsoluteUri;
|
||||
if (hostUrl.EndsWith("/sse", StringComparison.Ordinal))
|
||||
hostUrl = hostUrl[..^4];
|
||||
|
||||
var endpointUri = $"{hostUrl.TrimEnd('/')}/{data.TrimStart('/')}";
|
||||
|
||||
_messageEndpoint = new Uri(endpointUri);
|
||||
}
|
||||
|
||||
// Set connected state
|
||||
SetConnected(true);
|
||||
_connectionEstablished.TrySetResult(true);
|
||||
}
|
||||
catch (JsonException ex)
|
||||
{
|
||||
_logger.TransportEndpointEventParseFailed(EndpointName, data, ex);
|
||||
throw new McpTransportException("Failed to parse endpoint event", ex);
|
||||
}
|
||||
return default;
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue