Move the SseResponseStreamTransport out of sample (#47)
This commit is contained in:
parent
e1a5828036
commit
1b1c7fa8ca
6 changed files with 103 additions and 125 deletions
|
|
@ -13,7 +13,7 @@
|
|||
<PackageVersion Include="Microsoft.Extensions.AI.Abstractions" Version="$(MicrosoftExtensionsAIVersion)" />
|
||||
<PackageVersion Include="Microsoft.Extensions.Hosting.Abstractions" Version="$(MicrosoftExtensionsVersion)" />
|
||||
<PackageVersion Include="Microsoft.Extensions.Logging.Abstractions" Version="$(MicrosoftExtensionsVersion)" />
|
||||
<PackageVersion Include="System.Net.ServerSentEvents" Version="$(SystemVersion)" />
|
||||
<PackageVersion Include="System.Net.ServerSentEvents" Version="$(System10Version)" />
|
||||
<PackageVersion Include="System.Text.Json" Version="$(SystemVersion)" />
|
||||
<PackageVersion Include="System.Threading.Channels" Version="$(SystemVersion)" />
|
||||
|
||||
|
|
|
|||
|
|
@ -10,8 +10,4 @@
|
|||
<ProjectReference Include="..\..\src\ModelContextProtocol\ModelContextProtocol.csproj" />
|
||||
</ItemGroup>
|
||||
|
||||
<ItemGroup>
|
||||
<PackageReference Include="System.Net.ServerSentEvents" VersionOverride="10.0.0-preview.1.25080.5" />
|
||||
</ItemGroup>
|
||||
|
||||
</Project>
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
using ModelContextProtocol.Server;
|
||||
using ModelContextProtocol.Utils.Json;
|
||||
using Microsoft.Extensions.Options;
|
||||
using ModelContextProtocol.Protocol.Transport;
|
||||
|
||||
namespace AspNetCoreSseServer;
|
||||
|
||||
|
|
@ -10,7 +11,7 @@ public static class McpEndpointRouteBuilderExtensions
|
|||
public static IEndpointConventionBuilder MapMcpSse(this IEndpointRouteBuilder endpoints)
|
||||
{
|
||||
IMcpServer? server = null;
|
||||
SseServerStreamTransport? transport = null;
|
||||
SseResponseStreamTransport? transport = null;
|
||||
var loggerFactory = endpoints.ServiceProvider.GetRequiredService<ILoggerFactory>();
|
||||
var mcpServerOptions = endpoints.ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>();
|
||||
|
||||
|
|
@ -18,7 +19,7 @@ public static class McpEndpointRouteBuilderExtensions
|
|||
|
||||
routeGroup.MapGet("/sse", async (HttpResponse response, CancellationToken requestAborted) =>
|
||||
{
|
||||
await using var localTransport = transport = new SseServerStreamTransport(response.Body);
|
||||
await using var localTransport = transport = new SseResponseStreamTransport(response.Body);
|
||||
await using var localServer = server = McpServerFactory.Create(transport, mcpServerOptions.Value, loggerFactory, endpoints.ServiceProvider);
|
||||
|
||||
await localServer.StartAsync(requestAborted);
|
||||
|
|
@ -37,7 +38,7 @@ public static class McpEndpointRouteBuilderExtensions
|
|||
}
|
||||
});
|
||||
|
||||
routeGroup.MapPost("/message", async (HttpContext context) =>
|
||||
routeGroup.MapPost("/message", async context =>
|
||||
{
|
||||
if (transport is null)
|
||||
{
|
||||
|
|
|
|||
|
|
@ -1,14 +1,11 @@
|
|||
using ModelContextProtocol.Server;
|
||||
using System.Diagnostics.CodeAnalysis;
|
||||
using System.Net;
|
||||
using System.Text;
|
||||
using System.Net;
|
||||
using ModelContextProtocol.Server;
|
||||
|
||||
namespace ModelContextProtocol.Protocol.Transport;
|
||||
|
||||
/// <summary>
|
||||
/// HTTP server provider using HttpListener.
|
||||
/// </summary>
|
||||
[ExcludeFromCodeCoverage]
|
||||
internal class HttpListenerServerProvider : IDisposable
|
||||
{
|
||||
private static readonly byte[] s_accepted = "Accepted"u8.ToArray();
|
||||
|
|
@ -19,8 +16,6 @@ internal class HttpListenerServerProvider : IDisposable
|
|||
private readonly int _port;
|
||||
private HttpListener? _listener;
|
||||
private CancellationTokenSource? _cts;
|
||||
private Func<string, CancellationToken, bool>? _messageHandler;
|
||||
private StreamWriter? _streamWriter;
|
||||
private bool _isRunning;
|
||||
|
||||
/// <summary>
|
||||
|
|
@ -32,31 +27,8 @@ internal class HttpListenerServerProvider : IDisposable
|
|||
_port = port;
|
||||
}
|
||||
|
||||
public Task<string> GetSseEndpointUri()
|
||||
{
|
||||
return Task.FromResult($"http://localhost:{_port}{SseEndpoint}");
|
||||
}
|
||||
|
||||
public Task InitializeMessageHandler(Func<string, CancellationToken, bool> messageHandler)
|
||||
{
|
||||
_messageHandler = messageHandler;
|
||||
return Task.CompletedTask;
|
||||
}
|
||||
|
||||
public async Task SendEvent(string data, string eventId)
|
||||
{
|
||||
if (_streamWriter == null)
|
||||
{
|
||||
throw new McpServerException("Stream writer not initialized");
|
||||
}
|
||||
if (eventId != null)
|
||||
{
|
||||
await _streamWriter.WriteLineAsync($"id: {eventId}").ConfigureAwait(false);
|
||||
}
|
||||
await _streamWriter.WriteLineAsync($"data: {data}").ConfigureAwait(false);
|
||||
await _streamWriter.WriteLineAsync().ConfigureAwait(false); // Empty line to finish the event
|
||||
await _streamWriter.FlushAsync().ConfigureAwait(false);
|
||||
}
|
||||
public required Func<Stream, CancellationToken, Task> OnSseConnectionAsync { get; set; }
|
||||
public required Func<Stream, CancellationToken, Task<bool>> OnMessageAsync { get; set; }
|
||||
|
||||
/// <inheritdoc/>
|
||||
public Task StartAsync(CancellationToken cancellationToken = default)
|
||||
|
|
@ -64,7 +36,7 @@ internal class HttpListenerServerProvider : IDisposable
|
|||
if (_isRunning)
|
||||
return Task.CompletedTask;
|
||||
|
||||
_cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
||||
_cts = new CancellationTokenSource();
|
||||
_listener = new HttpListener();
|
||||
_listener.Prefixes.Add($"http://localhost:{_port}/");
|
||||
_listener.Start();
|
||||
|
|
@ -84,8 +56,6 @@ internal class HttpListenerServerProvider : IDisposable
|
|||
_cts?.Cancel();
|
||||
_listener?.Stop();
|
||||
|
||||
_streamWriter?.Close();
|
||||
|
||||
_isRunning = false;
|
||||
return Task.CompletedTask;
|
||||
}
|
||||
|
|
@ -170,28 +140,10 @@ internal class HttpListenerServerProvider : IDisposable
|
|||
response.Headers.Add("Cache-Control", "no-cache");
|
||||
response.Headers.Add("Connection", "keep-alive");
|
||||
|
||||
// Get the output stream and create a StreamWriter
|
||||
var outputStream = response.OutputStream;
|
||||
_streamWriter = new StreamWriter(outputStream, Encoding.UTF8) { AutoFlush = true };
|
||||
|
||||
// Keep the connection open until cancelled
|
||||
try
|
||||
{
|
||||
// Immediately send the "endpoint" event with the POST URL
|
||||
await _streamWriter.WriteLineAsync("event: endpoint").ConfigureAwait(false);
|
||||
await _streamWriter.WriteLineAsync($"data: {MessageEndpoint}").ConfigureAwait(false);
|
||||
await _streamWriter.WriteLineAsync().ConfigureAwait(false); // blank line to end an SSE message
|
||||
await _streamWriter.FlushAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// Keep the connection open by "pinging" or just waiting
|
||||
// until the client disconnects or the server is canceled.
|
||||
while (!cancellationToken.IsCancellationRequested && response.OutputStream.CanWrite)
|
||||
{
|
||||
// Do a periodic no-op to keep connection alive:
|
||||
await _streamWriter.WriteLineAsync(": keep-alive").ConfigureAwait(false);
|
||||
await _streamWriter.FlushAsync(cancellationToken).ConfigureAwait(false);
|
||||
await Task.Delay(TimeSpan.FromSeconds(5), cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
await OnSseConnectionAsync(response.OutputStream, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
catch (TaskCanceledException)
|
||||
{
|
||||
|
|
@ -206,7 +158,6 @@ internal class HttpListenerServerProvider : IDisposable
|
|||
// Remove client on disconnect
|
||||
try
|
||||
{
|
||||
_streamWriter.Close();
|
||||
response.Close();
|
||||
}
|
||||
catch { /* Ignore errors during cleanup */ }
|
||||
|
|
@ -218,15 +169,8 @@ internal class HttpListenerServerProvider : IDisposable
|
|||
var request = context.Request;
|
||||
var response = context.Response;
|
||||
|
||||
// Read the request body
|
||||
string requestBody;
|
||||
using (var reader = new StreamReader(request.InputStream, request.ContentEncoding))
|
||||
{
|
||||
requestBody = await reader.ReadToEndAsync(cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
// Process the message asynchronously
|
||||
if (_messageHandler != null && _messageHandler(requestBody, cancellationToken))
|
||||
if (await OnMessageAsync(request.InputStream, cancellationToken))
|
||||
{
|
||||
// Return 202 Accepted
|
||||
response.StatusCode = 202;
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
using System.Text.Json;
|
||||
using System.Net;
|
||||
using System.Text.Json;
|
||||
using Microsoft.Extensions.Logging;
|
||||
using ModelContextProtocol.Protocol.Messages;
|
||||
using ModelContextProtocol.Utils.Json;
|
||||
|
|
@ -9,16 +10,15 @@ using ModelContextProtocol.Utils;
|
|||
namespace ModelContextProtocol.Protocol.Transport;
|
||||
|
||||
/// <summary>
|
||||
/// Implements the MCP transport protocol over standard input/output streams.
|
||||
/// Implements the MCP transport protocol using <see cref="HttpListener"/>.
|
||||
/// </summary>
|
||||
public sealed class HttpListenerSseServerTransport : TransportBase, IServerTransport
|
||||
{
|
||||
private readonly string _serverName;
|
||||
private readonly HttpListenerServerProvider _httpServerProvider;
|
||||
private readonly ILogger<HttpListenerSseServerTransport> _logger;
|
||||
private readonly JsonSerializerOptions _jsonOptions;
|
||||
private CancellationTokenSource? _shutdownCts;
|
||||
|
||||
private SseResponseStreamTransport? _sseResponseStreamTransport;
|
||||
|
||||
private string EndpointName => $"Server (SSE) ({_serverName})";
|
||||
|
||||
/// <summary>
|
||||
|
|
@ -43,28 +43,23 @@ public sealed class HttpListenerSseServerTransport : TransportBase, IServerTrans
|
|||
{
|
||||
_serverName = serverName;
|
||||
_logger = loggerFactory.CreateLogger<HttpListenerSseServerTransport>();
|
||||
_jsonOptions = McpJsonUtilities.DefaultOptions;
|
||||
_httpServerProvider = new HttpListenerServerProvider(port);
|
||||
_httpServerProvider = new HttpListenerServerProvider(port)
|
||||
{
|
||||
OnSseConnectionAsync = OnSseConnectionAsync,
|
||||
OnMessageAsync = OnMessageAsync,
|
||||
};
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public Task StartListeningAsync(CancellationToken cancellationToken = default)
|
||||
{
|
||||
_shutdownCts = new CancellationTokenSource();
|
||||
|
||||
_httpServerProvider.InitializeMessageHandler(HttpMessageHandler);
|
||||
_httpServerProvider.StartAsync(cancellationToken);
|
||||
|
||||
SetConnected(true);
|
||||
|
||||
return Task.CompletedTask;
|
||||
return _httpServerProvider.StartAsync(cancellationToken);
|
||||
}
|
||||
|
||||
|
||||
/// <inheritdoc/>
|
||||
public override async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
|
||||
{
|
||||
if (!IsConnected)
|
||||
if (!IsConnected || _sseResponseStreamTransport is null)
|
||||
{
|
||||
_logger.TransportNotConnected(EndpointName);
|
||||
throw new McpTransportException("Transport is not connected");
|
||||
|
|
@ -78,10 +73,13 @@ public sealed class HttpListenerSseServerTransport : TransportBase, IServerTrans
|
|||
|
||||
try
|
||||
{
|
||||
var json = JsonSerializer.Serialize(message, _jsonOptions.GetTypeInfo<IJsonRpcMessage>());
|
||||
_logger.TransportSendingMessage(EndpointName, id, json);
|
||||
if (_logger.IsEnabled(LogLevel.Debug))
|
||||
{
|
||||
var json = JsonSerializer.Serialize(message, McpJsonUtilities.DefaultOptions.GetTypeInfo<IJsonRpcMessage>());
|
||||
_logger.TransportSendingMessage(EndpointName, id, json);
|
||||
}
|
||||
|
||||
await _httpServerProvider.SendEvent(json, "message").ConfigureAwait(false);
|
||||
await _sseResponseStreamTransport.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
_logger.TransportSentMessage(EndpointName, id);
|
||||
}
|
||||
|
|
@ -99,49 +97,61 @@ public sealed class HttpListenerSseServerTransport : TransportBase, IServerTrans
|
|||
GC.SuppressFinalize(this);
|
||||
}
|
||||
|
||||
private async Task CleanupAsync(CancellationToken cancellationToken)
|
||||
private Task CleanupAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
_logger.TransportCleaningUp(EndpointName);
|
||||
|
||||
if (_shutdownCts != null)
|
||||
{
|
||||
await _shutdownCts.CancelAsync().ConfigureAwait(false);
|
||||
_shutdownCts.Dispose();
|
||||
_shutdownCts = null;
|
||||
}
|
||||
|
||||
_httpServerProvider.Dispose();
|
||||
|
||||
SetConnected(false);
|
||||
|
||||
_logger.TransportCleanedUp(EndpointName);
|
||||
return Task.CompletedTask;
|
||||
}
|
||||
|
||||
private async Task OnSseConnectionAsync(Stream responseStream, CancellationToken cancellationToken)
|
||||
{
|
||||
await using var sseResponseStreamTransport = new SseResponseStreamTransport(responseStream);
|
||||
_sseResponseStreamTransport = sseResponseStreamTransport;
|
||||
SetConnected(true);
|
||||
await sseResponseStreamTransport.RunAsync(cancellationToken);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Handles HTTP messages received by the HTTP server provider.
|
||||
/// </summary>
|
||||
/// <returns>true if the message was accepted (return 202), false otherwise (return 400)</returns>
|
||||
private bool HttpMessageHandler(string request, CancellationToken cancellationToken)
|
||||
private async Task<bool> OnMessageAsync(Stream requestStream, CancellationToken cancellationToken)
|
||||
{
|
||||
_logger.TransportReceivedMessage(EndpointName, request);
|
||||
string request;
|
||||
IJsonRpcMessage? message = null;
|
||||
|
||||
if (_logger.IsEnabled(LogLevel.Information))
|
||||
{
|
||||
using var reader = new StreamReader(requestStream);
|
||||
request = await reader.ReadToEndAsync(cancellationToken).ConfigureAwait(false);
|
||||
message = JsonSerializer.Deserialize(request, McpJsonUtilities.DefaultOptions.GetTypeInfo<IJsonRpcMessage>());
|
||||
|
||||
_logger.TransportReceivedMessage(EndpointName, request);
|
||||
}
|
||||
else
|
||||
{
|
||||
request = "(Enable information-level logs to see the request)";
|
||||
}
|
||||
|
||||
try
|
||||
{
|
||||
var message = JsonSerializer.Deserialize(request, _jsonOptions.GetTypeInfo<IJsonRpcMessage>());
|
||||
message ??= await JsonSerializer.DeserializeAsync(requestStream, McpJsonUtilities.DefaultOptions.GetTypeInfo<IJsonRpcMessage>());
|
||||
if (message != null)
|
||||
{
|
||||
// Fire-and-forget the message to the message channel
|
||||
Task.Run(async () =>
|
||||
string messageId = "(no id)";
|
||||
if (message is IJsonRpcMessageWithId messageWithId)
|
||||
{
|
||||
string messageId = "(no id)";
|
||||
if (message is IJsonRpcMessageWithId messageWithId)
|
||||
{
|
||||
messageId = messageWithId.Id.ToString();
|
||||
}
|
||||
messageId = messageWithId.Id.ToString();
|
||||
}
|
||||
|
||||
_logger.TransportReceivedMessageParsed(EndpointName, messageId);
|
||||
await WriteMessageAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
_logger.TransportMessageWritten(EndpointName, messageId);
|
||||
}, cancellationToken);
|
||||
_logger.TransportReceivedMessageParsed(EndpointName, messageId);
|
||||
await WriteMessageAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
_logger.TransportMessageWritten(EndpointName, messageId);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,111 @@
|
|||
using System.Buffers;
|
||||
using System.Net.ServerSentEvents;
|
||||
using System.Text.Json;
|
||||
using System.Threading.Channels;
|
||||
using ModelContextProtocol.Protocol.Messages;
|
||||
using ModelContextProtocol.Utils.Json;
|
||||
|
||||
namespace ModelContextProtocol.Protocol.Transport;
|
||||
|
||||
/// <summary>
|
||||
/// Implements the MCP SSE server transport protocol using the SSE response <see cref="Stream"/>.
|
||||
/// </summary>
|
||||
/// <param name="sseResponseStream">The stream to write the SSE response body to.</param>
|
||||
public sealed class SseResponseStreamTransport(Stream sseResponseStream) : ITransport
|
||||
{
|
||||
private readonly Channel<IJsonRpcMessage> _incomingChannel = CreateSingleItemChannel<IJsonRpcMessage>();
|
||||
private readonly Channel<SseItem<IJsonRpcMessage?>> _outgoingSseChannel = CreateSingleItemChannel<SseItem<IJsonRpcMessage?>>();
|
||||
|
||||
private Task? _sseWriteTask;
|
||||
private Utf8JsonWriter? _jsonWriter;
|
||||
|
||||
/// <inheritdoc/>
|
||||
public bool IsConnected => _sseWriteTask?.IsCompleted == false;
|
||||
|
||||
/// <summary>
|
||||
/// Starts the transport and writes the JSON-RPC messages sent via <see cref="SendMessageAsync(IJsonRpcMessage, CancellationToken)"/>
|
||||
/// to the SSE response stream until cancelled or disposed.
|
||||
/// </summary>
|
||||
/// <param name="cancellationToken">A token to cancel writing to the SSE response stream.</param>
|
||||
/// <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("/message"u8);
|
||||
return;
|
||||
}
|
||||
|
||||
JsonSerializer.Serialize(GetUtf8JsonWriter(writer), item.Data, McpJsonUtilities.DefaultOptions.GetTypeInfo<IJsonRpcMessage?>());
|
||||
}
|
||||
|
||||
// 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.
|
||||
_outgoingSseChannel.Writer.TryWrite(new SseItem<IJsonRpcMessage?>(null, "endpoint"));
|
||||
|
||||
var sseItems = _outgoingSseChannel.Reader.ReadAllAsync(cancellationToken);
|
||||
return _sseWriteTask = SseFormatter.WriteAsync(sseItems, sseResponseStream, WriteJsonRpcMessageToBuffer, cancellationToken);
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public ChannelReader<IJsonRpcMessage> MessageReader => _incomingChannel.Reader;
|
||||
|
||||
/// <inheritdoc/>
|
||||
public ValueTask DisposeAsync()
|
||||
{
|
||||
_incomingChannel.Writer.TryComplete();
|
||||
_outgoingSseChannel.Writer.TryComplete();
|
||||
return new ValueTask(_sseWriteTask ?? Task.CompletedTask);
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
|
||||
{
|
||||
if (_sseWriteTask is null)
|
||||
{
|
||||
throw new InvalidOperationException($"Transport is not connected. Make sure to call {nameof(RunAsync)} first.");
|
||||
}
|
||||
|
||||
await _outgoingSseChannel.Writer.WriteAsync(new SseItem<IJsonRpcMessage?>(message), cancellationToken).AsTask();
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Handles incoming JSON-RPC messages received on the /message endpoint.
|
||||
/// </summary>
|
||||
/// <param name="message">The JSON-RPC message received.</param>
|
||||
/// <param name="cancellationToken">A token to cancel the operation.</param>
|
||||
/// <returns>A task representing the potentially asynchronous operation to buffer or process the JSON-RPC message.</returns>
|
||||
/// <exception cref="InvalidOperationException">Thrown when there is an attempt to process a message before calling <see cref="RunAsync(CancellationToken)"/>.</exception>
|
||||
public async Task OnMessageReceivedAsync(IJsonRpcMessage message, CancellationToken cancellationToken)
|
||||
{
|
||||
if (_sseWriteTask is null)
|
||||
{
|
||||
throw new InvalidOperationException($"Transport is not connected. Make sure to call {nameof(RunAsync)} first.");
|
||||
}
|
||||
|
||||
await _incomingChannel.Writer.WriteAsync(message, cancellationToken).AsTask();
|
||||
}
|
||||
|
||||
private static Channel<T> CreateSingleItemChannel<T>() =>
|
||||
Channel.CreateBounded<T>(new BoundedChannelOptions(1)
|
||||
{
|
||||
SingleReader = true,
|
||||
SingleWriter = false,
|
||||
});
|
||||
|
||||
private Utf8JsonWriter GetUtf8JsonWriter(IBufferWriter<byte> writer)
|
||||
{
|
||||
if (_jsonWriter is null)
|
||||
{
|
||||
_jsonWriter = new Utf8JsonWriter(writer);
|
||||
}
|
||||
else
|
||||
{
|
||||
_jsonWriter.Reset(writer);
|
||||
}
|
||||
|
||||
return _jsonWriter;
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue