Update HttpListenerSseServerTransport to use SseResponseStreamTransport

halter73 · halter73 · commit 265107f51d56 · 2025-03-21T13:13:34.000-07:00
diff --git a/src/ModelContextProtocol/Protocol/Transport/HttpListenerServerProvider.cs b/src/ModelContextProtocol/Protocol/Transport/HttpListenerServerProvider.cs
@@ -1,6 +1,5 @@
-﻿using ModelContextProtocol.Server;
-using System.Net;
-using System.Text;
+﻿using System.Net;
+using ModelContextProtocol.Server;
 
 namespace ModelContextProtocol.Protocol.Transport;
 
@@ -17,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>
@@ -30,34 +27,16 @@ public HttpListenerServerProvider(int port)
         _port = port;
     }
 
-    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<string, CancellationToken, Task<bool>> OnMessageAsync { get; set; }
 
     /// <inheritdoc/>
     public Task StartAsync(CancellationToken cancellationToken = default)
     {
         if (_isRunning)
             return Task.CompletedTask;
 
-        _cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        _cts = new CancellationTokenSource();
         _listener = new HttpListener();
         _listener.Prefixes.Add($"http://localhost:{_port}/");
         _listener.Start();
@@ -77,8 +56,6 @@ public Task StopAsync(CancellationToken cancellationToken = default)
         _cts?.Cancel();
         _listener?.Stop();
 
-        _streamWriter?.Close();
-
         _isRunning = false;
         return Task.CompletedTask;
     }
@@ -163,28 +140,10 @@ private async Task HandleSseConnectionAsync(HttpListenerContext context, Cancell
         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)
         {
@@ -199,7 +158,6 @@ private async Task HandleSseConnectionAsync(HttpListenerContext context, Cancell
             // Remove client on disconnect
             try
             {
-                _streamWriter.Close();
                 response.Close();
             }
             catch { /* Ignore errors during cleanup */ }
@@ -219,7 +177,7 @@ private async Task HandleMessageAsync(HttpListenerContext context, CancellationT
         }
 
         // Process the message asynchronously
-        if (_messageHandler != null && _messageHandler(requestBody, cancellationToken))
+        if (await OnMessageAsync(requestBody, cancellationToken))
         {
             // Return 202 Accepted
             response.StatusCode = 202;
diff --git a/src/ModelContextProtocol/Protocol/Transport/HttpListenerSseServerTransport.cs b/src/ModelContextProtocol/Protocol/Transport/HttpListenerSseServerTransport.cs
@@ -17,9 +17,8 @@ public sealed class HttpListenerSseServerTransport : TransportBase, IServerTrans
     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>
@@ -44,28 +43,23 @@ public HttpListenerSseServerTransport(string serverName, int port, ILoggerFactor
     {
         _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");
@@ -79,10 +73,10 @@ public override async Task SendMessageAsync(IJsonRpcMessage message, Cancellatio
 
         try
         {
-            var json = JsonSerializer.Serialize(message, _jsonOptions.GetTypeInfo<IJsonRpcMessage>());
+            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);
         }
@@ -100,49 +94,47 @@ public override async ValueTask DisposeAsync()
         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(string request, CancellationToken cancellationToken)
     {
         _logger.TransportReceivedMessage(EndpointName, request);
 
         try
         {
-            var message = JsonSerializer.Deserialize(request, _jsonOptions.GetTypeInfo<IJsonRpcMessage>());
+            var message = JsonSerializer.Deserialize(request, 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();
-                    }
-
-                    _logger.TransportReceivedMessageParsed(EndpointName, messageId);
-                    await WriteMessageAsync(message, cancellationToken).ConfigureAwait(false);
-                    _logger.TransportMessageWritten(EndpointName, messageId);
-                }, cancellationToken);
+                    messageId = messageWithId.Id.ToString();
+                }
+
+                _logger.TransportReceivedMessageParsed(EndpointName, messageId);
+                await WriteMessageAsync(message, cancellationToken).ConfigureAwait(false);
+                _logger.TransportMessageWritten(EndpointName, messageId);
 
                 return true;
             }
