SSE Transport Support (#13)

* First iteration

* Tests

* SSE support
This commit is contained in:
Peder Holdgaard Pedersen 2025-01-12 16:35:39 +01:00 • committed by GitHub
commit e51c78e5e4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 1200 additions and 43 deletions

View file

@ -5,7 +5,7 @@ using System.Text.Json.Serialization;
/// <summary>
/// An error response message in the JSON-RPC protocol.
/// </summary>
internal record JsonRpcError : IJsonRpcMessageWithId
public record JsonRpcError : IJsonRpcMessageWithId
{
/// <summary>
/// JSON-RPC protocol version. Always "2.0".

View file

@ -5,7 +5,7 @@ using System.Text.Json.Serialization;
/// <summary>
/// Detailed error information for JSON-RPC error responses.
/// </summary>
internal record JsonRpcErrorDetail
public record JsonRpcErrorDetail
{
/// <summary>
/// Integer error code.

View file

@ -5,7 +5,7 @@ namespace McpDotNet.Protocol.Messages;
/// <summary>
/// A successful response message in the JSON-RPC protocol.
/// </summary>
internal record JsonRpcResponse : IJsonRpcMessageWithId
public record JsonRpcResponse : IJsonRpcMessageWithId
{
/// <summary>
/// JSON-RPC protocol version. Always "2.0".

View file

@ -1,9 +1,7 @@
// Protocol/Transport/IMcpTransport.cs
using System.Threading.Channels;
using McpDotNet.Protocol.Messages;
namespace McpDotNet.Protocol.Transport;
using System.Threading.Channels;
using global::McpDotNet.Protocol.Messages;
/// <summary>
/// Represents a transport mechanism for MCP communication.
/// </summary>

View file

@ -0,0 +1,303 @@

using System.Diagnostics.Tracing;
using System.Text.Json;
using System.Text;
using McpDotNet.Protocol.Messages;
using McpDotNet.Configuration;
using Microsoft.Extensions.Logging;
using McpDotNet.Logging;
using System.Net.Http.Headers;
using System.Text.Json.Serialization;
using McpDotNet.Utils.Json;
using System.Threading;
namespace McpDotNet.Protocol.Transport;
/// <summary>
///
/// </summary>
public sealed class SseTransport : TransportBase
{
private readonly HttpClient _httpClient;
private readonly SseTransportOptions _options;
private readonly Uri _sseEndpoint;
private Uri? _messageEndpoint;
private readonly CancellationTokenSource _connectionCts;
private Task? _receiveTask;
private readonly ILogger<SseTransport> _logger;
private readonly McpServerConfig _serverConfig;
private readonly JsonSerializerOptions _jsonOptions;
private TaskCompletionSource<bool> _connectionEstablished;
/// <summary>
///
/// </summary>
/// <param name="transportOptions"></param>
/// <param name="serverConfig"></param>
/// <param name="loggerFactory"></param>
public SseTransport(SseTransportOptions transportOptions, McpServerConfig serverConfig, ILoggerFactory loggerFactory)
{
_options = transportOptions;
_serverConfig = serverConfig;
_sseEndpoint = new Uri(serverConfig.Location!);
_httpClient = new HttpClient();
_connectionCts = new CancellationTokenSource();
_logger = loggerFactory.CreateLogger<SseTransport>();
_jsonOptions = new JsonSerializerOptions().ConfigureForMcp(loggerFactory);
_connectionEstablished = new TaskCompletionSource<bool>();
}
/// <inheritdoc/>
public override async Task ConnectAsync(CancellationToken cancellationToken = default)
{
try
{
if (IsConnected)
{
_logger.TransportAlreadyConnected(_serverConfig.Id, _serverConfig.Name);
throw new McpTransportException("Transport is already connected");
}
// Start message receiving loop
_receiveTask = ReceiveMessagesAsync(_connectionCts.Token);
_logger.TransportReadingMessages(_serverConfig.Id, _serverConfig.Name);
await Task.WhenAny(
_connectionEstablished.Task,
Task.Delay(_options.ConnectionTimeout, cancellationToken)
);
if (!IsConnected)
throw new TimeoutException("Failed to receive endpoint event");
}
catch (McpTransportException)
{
// Rethrow transport exceptions
throw;
}
catch (Exception ex)
{
_logger.TransportConnectFailed(_serverConfig.Id, _serverConfig.Name, ex);
await CloseAsync();
throw new McpTransportException("Failed to connect transport", ex);
}
}
public override async Task SendMessageAsync(
IJsonRpcMessage message,
CancellationToken cancellationToken = default)
{
if (_messageEndpoint == null)
throw new InvalidOperationException("Transport not connected");
var content = new StringContent(
JsonSerializer.Serialize(message, _jsonOptions),
Encoding.UTF8,
"application/json"
);
var response = await _httpClient.PostAsync(
_messageEndpoint,
content,
cancellationToken
);
response.EnsureSuccessStatusCode();
var responseContent = await response.Content.ReadAsStringAsync();
// Handle notifications, which don't have responses
if (message is JsonRpcNotification)
{
await HandleNotificationResponseAsync(message, responseContent, cancellationToken);
return;
}
var responseMessage = JsonSerializer.Deserialize<IJsonRpcMessage>(responseContent, _jsonOptions);
if (responseMessage != null)
{
string messageId = "(no id)";
if (responseMessage is IJsonRpcMessageWithId messageWithId)
{
messageId = messageWithId.Id.ToString();
}
_logger.TransportReceivedMessageParsed(_serverConfig.Id, _serverConfig.Name, messageId);
await WriteMessageAsync(responseMessage, cancellationToken);
_logger.TransportMessageWritten(_serverConfig.Id, _serverConfig.Name, messageId);
}
else
{
_logger.TransportMessageParseUnexpectedType(_serverConfig.Id, _serverConfig.Name, responseContent);
}
}
/// <inheritdoc/>
public async Task CloseAsync()
{
_connectionCts.Cancel();
if (_receiveTask != null)
await _receiveTask;
_httpClient.Dispose();
_connectionCts.Dispose();
SetConnected(false);
}
/// <inheritdoc/>
public override async ValueTask DisposeAsync()
{
await CloseAsync();
GC.SuppressFinalize(this);
}
private async Task HandleNotificationResponseAsync(
IJsonRpcMessage message,
string responseContent,
CancellationToken cancellationToken)
{
// Check for error response, but don't require one
try
{
var errorResponse = JsonSerializer.Deserialize<JsonRpcError>(responseContent, _jsonOptions);
if (errorResponse is not null)
{
await WriteMessageAsync(errorResponse, cancellationToken);
return;
}
}
catch (JsonException)
{
// Ignore deserialization errors for notifications
// This is expected as most notifications won't have responses
}
}
internal Uri? MessageEndpoint => _messageEndpoint;
internal SseTransportOptions Options => _options;
private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
{
int reconnectAttempts = 0;
while (!cancellationToken.IsCancellationRequested)
{
try
{
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
);
response.EnsureSuccessStatusCode();
using var stream = await response.Content.ReadAsStreamAsync(cancellationToken);
using var reader = new StreamReader(stream);
// Reset reconnect attempts on successful connection
reconnectAttempts = 0;
string? currentEvent = null;
string? line;
while ((line = await reader.ReadLineAsync()) != null)
{
if (line.StartsWith("event: "))
{
currentEvent = line.Substring(7);
}
else if (line.StartsWith("data: "))
{
var data = line.Substring(6);
if (currentEvent == "endpoint")
{
HandleEndpointEvent(data);
}
else
{
await ProcessSseMessage(data, cancellationToken);
}
}
}
}
catch (Exception ex) when (!cancellationToken.IsCancellationRequested)
{
_logger.TransportConnectionError(_serverConfig.Id, _serverConfig.Name, ex);
reconnectAttempts++;
if (reconnectAttempts >= _options.MaxReconnectAttempts)
{
throw new McpTransportException("Exceeded reconnect limit", ex);
}
await Task.Delay(_options.ReconnectDelay, cancellationToken);
}
}
}
private async Task ProcessSseMessage(string data, CancellationToken cancellationToken)
{
if (!IsConnected)
{
_logger.TransportMessageReceivedBeforeConnected(_serverConfig.Id, _serverConfig.Name, data);
return;
}
try
{
var message = JsonSerializer.Deserialize<IJsonRpcMessage>(data, _jsonOptions);
if (message == null)
{
_logger.TransportMessageParseUnexpectedType(_serverConfig.Id, _serverConfig.Name, data);
return;
}
string messageId = "(no id)";
if (message is IJsonRpcMessageWithId messageWithId)
{
messageId = messageWithId.Id.ToString();
}
_logger.TransportReceivedMessageParsed(_serverConfig.Id, _serverConfig.Name, messageId);
await WriteMessageAsync(message, cancellationToken);
_logger.TransportMessageWritten(_serverConfig.Id, _serverConfig.Name, messageId);
}
catch (JsonException ex)
{
_logger.TransportMessageParseFailed(_serverConfig.Id, _serverConfig.Name, data, ex);
}
}
private void HandleEndpointEvent(string data)
{
try
{
var endpointData = JsonSerializer.Deserialize<EndpointEventData>(data, _jsonOptions);
if (endpointData?.Uri == null)
{
_logger.TransportEndpointEventInvalid(_serverConfig.Id, _serverConfig.Name, data);
return;
}
_messageEndpoint = new Uri(endpointData.Uri);
SetConnected(true);
_connectionEstablished.TrySetResult(true);
}
catch (JsonException ex)
{
_logger.TransportEndpointEventParseFailed(_serverConfig.Id, _serverConfig.Name, data, ex);
throw new McpTransportException("Failed to parse endpoint event", ex);
}
}
private record EndpointEventData
{
[JsonPropertyName("uri")]
public string? Uri { get; init; }
}
}

View file

@ -0,0 +1,27 @@
namespace McpDotNet.Protocol.Transport;
/// <summary>
/// Options for configuring the SSE transport.
/// </summary>
public record SseTransportOptions
{
/// <summary>
/// Timeout for initial connection and endpoint event.
/// </summary>
public TimeSpan ConnectionTimeout { get; init; } = TimeSpan.FromSeconds(30);
/// <summary>
/// Number of reconnection attempts for SSE connection.
/// </summary>
public int MaxReconnectAttempts { get; init; } = 3;
/// <summary>
/// Delay between reconnection attempts.
/// </summary>
public TimeSpan ReconnectDelay { get; init; } = TimeSpan.FromSeconds(5);
/// <summary>
/// Headers to include in HTTP requests.
/// </summary>
public Dictionary<string, string>? AdditionalHeaders { get; init; }
}