Server-side disconnect

Author: MackinnonBuck

src/ModelContextProtocol.Core/Server/RequestContext.cs

@@ -81,4 +81,21 @@ public McpServer Server
     /// including the method name, parameters, request ID, and associated transport and user information.
     /// </remarks>
     public JsonRpcRequest JsonRpcRequest { get; }
+
+    /// <summary>
+    /// Ends the current response and enables polling for updates from the server.
+    /// </summary>
+    /// <param name="retryInterval">The interval at which the client should poll for updates.</param>
+    /// <param name="cancellationToken">The cancellation token.</param>
+    /// <returns>A <see cref="ValueTask"/> that completes when polling has been enabled.</returns>
+    public async ValueTask EnablePollingAsync(TimeSpan retryInterval, CancellationToken cancellationToken = default)
+    {
+        if (JsonRpcRequest.Context?.RelatedTransport is not StreamableHttpPostTransport transport)
+        {
+            // Polling is only supported for Streamable HTTP POST transports.
+            return;
+        }
+
+        await transport.EnablePollingAsync(retryInterval, cancellationToken);
+    }
 }

src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs

@@ -102,6 +102,25 @@ public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken can
         }
     }
 
+    public async ValueTask EnablePollingAsync(TimeSpan retryInterval, CancellationToken cancellationToken)
+    {
+        var eventStreamWriter = await GetOrCreateEventStreamAsync(cancellationToken).ConfigureAwait(false);
+        if (eventStreamWriter is null)
+        {
+            return;
+        }
+
+        // Set the mode to 'Polling' so that the replay stream ends as soon as all available messages have been sent.
+        // This prevents the client from immediately establishing another long-lived connection.
+        await eventStreamWriter.SetModeAsync(SseEventStreamMode.Polling, cancellationToken).ConfigureAwait(false);
+
+        // Send the priming event with the new retry interval.
+        await _sseWriter.SendPrimingEventAsync(retryInterval, eventStreamWriter, cancellationToken).ConfigureAwait(false);
+
+        // Dispose the writer to close it and force future writes to only apply to the SSE event store.
+        await _sseWriter.DisposeAsync();
+    }
+
     public async ValueTask DisposeAsync()
     {
         await _sseWriter.DisposeAsync().ConfigureAwait(false);