Streamable HTTP resumability + redelivery + SSE polling via server-side disconnect (#1077)
Co-authored-by: Stephen Halter <halter73@gmail.com>
This commit is contained in:
parent
4cb2fd4655
commit
b8311756b4
18 changed files with 2675 additions and 246 deletions
2
src/ModelContextProtocol.Core/.editorconfig
Normal file
2
src/ModelContextProtocol.Core/.editorconfig
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
[*.cs]
|
||||
dotnet_diagnostic.CA2007.severity = error # CA2007: Do not directly await a Task without ConfigureAwait
|
||||
|
|
@ -108,4 +108,31 @@ public sealed class HttpClientTransportOptions
|
|||
/// Gets sor sets the authorization provider to use for authentication.
|
||||
/// </summary>
|
||||
public ClientOAuthOptions? OAuth { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the maximum number of consecutive reconnection attempts when an SSE stream is disconnected.
|
||||
/// </summary>
|
||||
/// <value>
|
||||
/// The maximum number of reconnection attempts. The default is 2.
|
||||
/// </value>
|
||||
/// <remarks>
|
||||
/// When an SSE stream is disconnected (e.g., due to a network issue), the client will attempt to
|
||||
/// reconnect using the Last-Event-ID header to resume from where it left off. This property controls
|
||||
/// how many reconnection attempts are made before giving up.
|
||||
/// </remarks>
|
||||
public int MaxReconnectionAttempts { get; set; } = 2;
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the default interval at which the client attempts reconnection after an SSE stream is disconnected.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// <p>
|
||||
/// The default value is 1 second.
|
||||
/// </p>
|
||||
/// <p>
|
||||
/// If the server sends a message specifying a different reconnection interval, that new value will be used for all
|
||||
/// subsequent reconnection attempts for that stream.
|
||||
/// </p>
|
||||
/// </remarks>
|
||||
public TimeSpan DefaultReconnectionInterval { get; set; } = TimeSpan.FromSeconds(1);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ internal sealed class StdioClientSessionTransport(
|
|||
{
|
||||
try
|
||||
{
|
||||
await base.SendMessageAsync(message, cancellationToken);
|
||||
await base.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
catch (IOException)
|
||||
{
|
||||
|
|
@ -58,7 +58,7 @@ internal sealed class StdioClientSessionTransport(
|
|||
}
|
||||
|
||||
// And handle cleanup in the base type.
|
||||
await base.CleanupAsync(error, cancellationToken);
|
||||
await base.CleanupAsync(error, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
private async ValueTask<Exception?> GetUnexpectedExitExceptionAsync(CancellationToken cancellationToken)
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ using System.Net.ServerSentEvents;
|
|||
using System.Text.Json;
|
||||
using ModelContextProtocol.Protocol;
|
||||
using System.Threading.Channels;
|
||||
using System.Net;
|
||||
|
||||
namespace ModelContextProtocol.Client;
|
||||
|
||||
|
|
@ -105,8 +106,18 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
|
|||
}
|
||||
else if (response.Content.Headers.ContentType?.MediaType == "text/event-stream")
|
||||
{
|
||||
using var responseBodyStream = await response.Content.ReadAsStreamAsync(cancellationToken);
|
||||
rpcResponseOrError = await ProcessSseResponseAsync(responseBodyStream, rpcRequest, cancellationToken).ConfigureAwait(false);
|
||||
var sseState = new SseStreamState();
|
||||
using var responseBodyStream = await response.Content.ReadAsStreamAsync(cancellationToken).ConfigureAwait(false);
|
||||
var sseResponse = await ProcessSseResponseAsync(responseBodyStream, rpcRequest, sseState, cancellationToken).ConfigureAwait(false);
|
||||
rpcResponseOrError = sseResponse.Response;
|
||||
|
||||
// Resumability: If POST SSE stream ended without a response but we have a Last-Event-ID (from priming),
|
||||
// attempt to resume by sending a GET request with Last-Event-ID header. The server will replay
|
||||
// events from the event store, allowing us to receive the pending response.
|
||||
if (rpcResponseOrError is null && rpcRequest is not null && sseState.LastEventId is not null)
|
||||
{
|
||||
rpcResponseOrError = await SendGetSseRequestWithRetriesAsync(rpcRequest, sseState, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
}
|
||||
|
||||
if (rpcRequest is null)
|
||||
|
|
@ -155,7 +166,7 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
|
|||
// Send DELETE request to terminate the session. Only send if we have a session ID, per MCP spec.
|
||||
if (_options.OwnsSession && !string.IsNullOrEmpty(SessionId))
|
||||
{
|
||||
await SendDeleteRequest();
|
||||
await SendDeleteRequest().ConfigureAwait(false);
|
||||
}
|
||||
|
||||
if (_getReceiveTask != null)
|
||||
|
|
@ -188,56 +199,147 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
|
|||
|
||||
private async Task ReceiveUnsolicitedMessagesAsync()
|
||||
{
|
||||
// Send a GET request to handle any unsolicited messages not sent over a POST response.
|
||||
using var request = new HttpRequestMessage(HttpMethod.Get, _options.Endpoint);
|
||||
request.Headers.Accept.Add(s_textEventStreamMediaType);
|
||||
CopyAdditionalHeaders(request.Headers, _options.AdditionalHeaders, SessionId, _negotiatedProtocolVersion);
|
||||
var state = new SseStreamState();
|
||||
|
||||
// Server support for the GET request is optional. If it fails, we don't care. It just means we won't receive unsolicited messages.
|
||||
HttpResponseMessage response;
|
||||
try
|
||||
// Continuously receive unsolicited messages until canceled
|
||||
while (!_connectionCts.Token.IsCancellationRequested)
|
||||
{
|
||||
response = await _httpClient.SendAsync(request, message: null, _connectionCts.Token).ConfigureAwait(false);
|
||||
}
|
||||
catch (HttpRequestException)
|
||||
{
|
||||
return;
|
||||
}
|
||||
await SendGetSseRequestWithRetriesAsync(
|
||||
relatedRpcRequest: null,
|
||||
state,
|
||||
_connectionCts.Token).ConfigureAwait(false);
|
||||
|
||||
using (response)
|
||||
{
|
||||
if (!response.IsSuccessStatusCode)
|
||||
// If we exhausted retries without receiving any events, stop trying
|
||||
if (state.LastEventId is null)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
using var responseStream = await response.Content.ReadAsStreamAsync(_connectionCts.Token).ConfigureAwait(false);
|
||||
await ProcessSseResponseAsync(responseStream, relatedRpcRequest: null, _connectionCts.Token).ConfigureAwait(false);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task<JsonRpcMessageWithId?> ProcessSseResponseAsync(Stream responseStream, JsonRpcRequest? relatedRpcRequest, CancellationToken cancellationToken)
|
||||
/// <summary>
|
||||
/// Sends a GET request for SSE with retry logic and resumability support.
|
||||
/// </summary>
|
||||
private async Task<JsonRpcMessageWithId?> SendGetSseRequestWithRetriesAsync(
|
||||
JsonRpcRequest? relatedRpcRequest,
|
||||
SseStreamState state,
|
||||
CancellationToken cancellationToken)
|
||||
{
|
||||
await foreach (SseItem<string> sseEvent in SseParser.Create(responseStream).EnumerateAsync(cancellationToken).ConfigureAwait(false))
|
||||
int attempt = 0;
|
||||
|
||||
// Delay before first attempt if we're reconnecting (have a Last-Event-ID)
|
||||
bool shouldDelay = state.LastEventId is not null;
|
||||
|
||||
while (attempt < _options.MaxReconnectionAttempts)
|
||||
{
|
||||
if (sseEvent.EventType != "message")
|
||||
cancellationToken.ThrowIfCancellationRequested();
|
||||
|
||||
if (shouldDelay)
|
||||
{
|
||||
var delay = state.RetryInterval ?? _options.DefaultReconnectionInterval;
|
||||
await Task.Delay(delay, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
shouldDelay = true;
|
||||
|
||||
using var request = new HttpRequestMessage(HttpMethod.Get, _options.Endpoint);
|
||||
request.Headers.Accept.Add(s_textEventStreamMediaType);
|
||||
CopyAdditionalHeaders(request.Headers, _options.AdditionalHeaders, SessionId, _negotiatedProtocolVersion, state.LastEventId);
|
||||
|
||||
HttpResponseMessage response;
|
||||
try
|
||||
{
|
||||
response = await _httpClient.SendAsync(request, message: null, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
catch (HttpRequestException)
|
||||
{
|
||||
attempt++;
|
||||
continue;
|
||||
}
|
||||
|
||||
var rpcResponseOrError = await ProcessMessageAsync(sseEvent.Data, relatedRpcRequest, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// The server SHOULD end the HTTP response body here anyway, but we won't leave it to chance. This transport makes
|
||||
// a GET request for any notifications that might need to be sent after the completion of each POST.
|
||||
if (rpcResponseOrError is not null)
|
||||
using (response)
|
||||
{
|
||||
return rpcResponseOrError;
|
||||
if (response.StatusCode >= HttpStatusCode.InternalServerError)
|
||||
{
|
||||
// Server error; retry.
|
||||
attempt++;
|
||||
continue;
|
||||
}
|
||||
|
||||
if (!response.IsSuccessStatusCode)
|
||||
{
|
||||
// If the server could be reached but returned a non-success status code,
|
||||
// retrying likely won't change that.
|
||||
return null;
|
||||
}
|
||||
|
||||
using var responseStream = await response.Content.ReadAsStreamAsync(cancellationToken).ConfigureAwait(false);
|
||||
var sseResponse = await ProcessSseResponseAsync(responseStream, relatedRpcRequest, state, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
if (sseResponse.Response is { } rpcResponseOrError)
|
||||
{
|
||||
return rpcResponseOrError;
|
||||
}
|
||||
|
||||
// If we reach here, then the stream closed without the response.
|
||||
|
||||
if (sseResponse.IsNetworkError || state.LastEventId is null)
|
||||
{
|
||||
// No event ID means server may not support resumability; don't retry indefinitely.
|
||||
attempt++;
|
||||
}
|
||||
else
|
||||
{
|
||||
// We have an event ID, so we continue polling to receive more events.
|
||||
// The server should eventually send a response or return an error.
|
||||
attempt = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
private async Task<SseResponse> ProcessSseResponseAsync(
|
||||
Stream responseStream,
|
||||
JsonRpcRequest? relatedRpcRequest,
|
||||
SseStreamState state,
|
||||
CancellationToken cancellationToken)
|
||||
{
|
||||
try
|
||||
{
|
||||
await foreach (SseItem<string> sseEvent in SseParser.Create(responseStream).EnumerateAsync(cancellationToken).ConfigureAwait(false))
|
||||
{
|
||||
// Track event ID and retry interval for resumability
|
||||
if (!string.IsNullOrEmpty(sseEvent.EventId))
|
||||
{
|
||||
state.LastEventId = sseEvent.EventId;
|
||||
}
|
||||
if (sseEvent.ReconnectionInterval.HasValue)
|
||||
{
|
||||
state.RetryInterval = sseEvent.ReconnectionInterval.Value;
|
||||
}
|
||||
|
||||
// Skip events with empty data
|
||||
if (string.IsNullOrEmpty(sseEvent.Data))
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
var rpcResponseOrError = await ProcessMessageAsync(sseEvent.Data, relatedRpcRequest, cancellationToken).ConfigureAwait(false);
|
||||
if (rpcResponseOrError is not null)
|
||||
{
|
||||
return new() { Response = rpcResponseOrError };
|
||||
}
|
||||
}
|
||||
}
|
||||
catch (Exception ex) when (ex is IOException or HttpRequestException)
|
||||
{
|
||||
return new() { IsNetworkError = true };
|
||||
}
|
||||
|
||||
return default;
|
||||
}
|
||||
|
||||
private async Task<JsonRpcMessageWithId?> ProcessMessageAsync(string data, JsonRpcRequest? relatedRpcRequest, CancellationToken cancellationToken)
|
||||
{
|
||||
LogTransportReceivedMessageSensitive(Name, data);
|
||||
|
|
@ -292,7 +394,8 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
|
|||
HttpRequestHeaders headers,
|
||||
IDictionary<string, string>? additionalHeaders,
|
||||
string? sessionId,
|
||||
string? protocolVersion)
|
||||
string? protocolVersion,
|
||||
string? lastEventId = null)
|
||||
{
|
||||
if (sessionId is not null)
|
||||
{
|
||||
|
|
@ -304,6 +407,11 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
|
|||
headers.Add("MCP-Protocol-Version", protocolVersion);
|
||||
}
|
||||
|
||||
if (lastEventId is not null)
|
||||
{
|
||||
headers.Add("Last-Event-ID", lastEventId);
|
||||
}
|
||||
|
||||
if (additionalHeaders is null)
|
||||
{
|
||||
return;
|
||||
|
|
@ -317,4 +425,22 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Tracks state across SSE stream connections.
|
||||
/// </summary>
|
||||
private sealed class SseStreamState
|
||||
{
|
||||
public string? LastEventId { get; set; }
|
||||
public TimeSpan? RetryInterval { get; set; }
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Represents the result of processing an SSE response.
|
||||
/// </summary>
|
||||
private readonly struct SseResponse
|
||||
{
|
||||
public JsonRpcMessageWithId? Response { get; init; }
|
||||
public bool IsNetworkError { get; init; }
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -29,16 +29,38 @@ internal sealed partial class McpSessionHandler : IAsyncDisposable
|
|||
"mcp.server.operation.duration", "MCP request or notification duration as observed on the receiver from the time it was received until the result or ack is sent.");
|
||||
|
||||
/// <summary>The latest version of the protocol supported by this implementation.</summary>
|
||||
internal const string LatestProtocolVersion = "2025-06-18";
|
||||
internal const string LatestProtocolVersion = "2025-11-25";
|
||||
|
||||
/// <summary>All protocol versions supported by this implementation.</summary>
|
||||
internal static readonly string[] SupportedProtocolVersions =
|
||||
[
|
||||
"2024-11-05",
|
||||
"2025-03-26",
|
||||
"2025-06-18",
|
||||
LatestProtocolVersion,
|
||||
];
|
||||
|
||||
/// <summary>
|
||||
/// Checks if the given protocol version supports priming events.
|
||||
/// </summary>
|
||||
/// <param name="protocolVersion">The protocol version to check.</param>
|
||||
/// <returns>True if the protocol version supports resumability.</returns>
|
||||
/// <remarks>
|
||||
/// Priming events are only supported in protocol version >= 2025-11-25.
|
||||
/// Older clients may crash when receiving SSE events with empty data.
|
||||
/// </remarks>
|
||||
internal static bool SupportsPrimingEvent(string? protocolVersion)
|
||||
{
|
||||
const string MinResumabilityProtocolVersion = "2025-11-25";
|
||||
|
||||
if (protocolVersion is null)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
return string.Compare(protocolVersion, MinResumabilityProtocolVersion, StringComparison.Ordinal) >= 0;
|
||||
}
|
||||
|
||||
private readonly bool _isServer;
|
||||
private readonly string _transportKind;
|
||||
private readonly ITransport _transport;
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@
|
|||
<Compile Include="..\Common\Throw.cs" Link="Throw.cs" />
|
||||
<Compile Include="..\Common\Obsoletions.cs" Link="Obsoletions.cs" />
|
||||
<Compile Include="..\Common\Experimentals.cs" Link="Experimentals.cs" />
|
||||
<Compile Include="..\Common\ServerSentEvents\**\*.cs" Link="ServerSentEvents\%(RecursiveDir)%(FileName)%(Extension)" />
|
||||
</ItemGroup>
|
||||
|
||||
<!-- Dependencies only needed by netstandard2.0 -->
|
||||
|
|
|
|||
|
|
@ -0,0 +1,36 @@
|
|||
using ModelContextProtocol.Protocol;
|
||||
using System.Net.ServerSentEvents;
|
||||
|
||||
namespace ModelContextProtocol.Server;
|
||||
|
||||
/// <summary>
|
||||
/// Provides read access to an SSE event stream, allowing events to be consumed asynchronously.
|
||||
/// </summary>
|
||||
public interface ISseEventStreamReader
|
||||
{
|
||||
/// <summary>
|
||||
/// Gets the session ID associated with the stream being read.
|
||||
/// </summary>
|
||||
string SessionId { get; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets the ID of the stream.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// This value is guaranteed to be unique on a per-session basis.
|
||||
/// </remarks>
|
||||
string StreamId { get; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets the messages from the stream as an <see cref="IAsyncEnumerable{T}"/>.
|
||||
/// </summary>
|
||||
/// <param name="cancellationToken">A token to cancel the operation.</param>
|
||||
/// <returns>An <see cref="IAsyncEnumerable{T}"/> of <see cref="SseItem{T}"/> containing JSON-RPC messages.</returns>
|
||||
/// <remarks>
|
||||
/// If the stream's mode is set to <see cref="SseEventStreamMode.Polling"/>, the returned
|
||||
/// messages will only include the currently-available events starting at the last event ID specified
|
||||
/// when the reader was created. Otherwise, the returned messages will continue until the associated
|
||||
/// <see cref="ISseEventStreamWriter"/> is disposed.
|
||||
/// </remarks>
|
||||
IAsyncEnumerable<SseItem<JsonRpcMessage?>> ReadEventsAsync(CancellationToken cancellationToken = default);
|
||||
}
|
||||
23
src/ModelContextProtocol.Core/Server/ISseEventStreamStore.cs
Normal file
23
src/ModelContextProtocol.Core/Server/ISseEventStreamStore.cs
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
namespace ModelContextProtocol.Server;
|
||||
|
||||
/// <summary>
|
||||
/// Provides storage and retrieval of SSE event streams, enabling resumability and redelivery of events.
|
||||
/// </summary>
|
||||
public interface ISseEventStreamStore
|
||||
{
|
||||
/// <summary>
|
||||
/// Creates a new SSE event stream with the specified options.
|
||||
/// </summary>
|
||||
/// <param name="options">The configuration options for the new stream.</param>
|
||||
/// <param name="cancellationToken">A token to cancel the operation.</param>
|
||||
/// <returns>A writer for the newly created event stream.</returns>
|
||||
ValueTask<ISseEventStreamWriter> CreateStreamAsync(SseEventStreamOptions options, CancellationToken cancellationToken = default);
|
||||
|
||||
/// <summary>
|
||||
/// Gets a reader for an existing event stream based on the last event ID.
|
||||
/// </summary>
|
||||
/// <param name="lastEventId">The ID of the last event received by the client, used to resume from that point.</param>
|
||||
/// <param name="cancellationToken">A token to cancel the operation.</param>
|
||||
/// <returns>A reader for the event stream, or <c>null</c> if no matching stream is found.</returns>
|
||||
ValueTask<ISseEventStreamReader?> GetStreamReaderAsync(string lastEventId, CancellationToken cancellationToken = default);
|
||||
}
|
||||
|
|
@ -0,0 +1,30 @@
|
|||
using ModelContextProtocol.Protocol;
|
||||
using System.Net.ServerSentEvents;
|
||||
|
||||
namespace ModelContextProtocol.Server;
|
||||
|
||||
/// <summary>
|
||||
/// Provides write access to an SSE event stream, allowing events to be written and tracked with unique IDs.
|
||||
/// </summary>
|
||||
public interface ISseEventStreamWriter : IAsyncDisposable
|
||||
{
|
||||
/// <summary>
|
||||
/// Sets the mode of the event stream.
|
||||
/// </summary>
|
||||
/// <param name="mode">The new mode to set for the event stream.</param>
|
||||
/// <param name="cancellationToken">A token to cancel the operation.</param>
|
||||
/// <returns>A task that represents the asynchronous operation.</returns>
|
||||
ValueTask SetModeAsync(SseEventStreamMode mode, CancellationToken cancellationToken = default);
|
||||
|
||||
/// <summary>
|
||||
/// Writes an event to the stream.
|
||||
/// </summary>
|
||||
/// <param name="sseItem">The original <see cref="SseItem{T}"/>.</param>
|
||||
/// <param name="cancellationToken">A token to cancel the operation.</param>
|
||||
/// <returns>A new <see cref="SseItem{T}"/> with a populated event ID.</returns>
|
||||
/// <remarks>
|
||||
/// If the provided <paramref name="sseItem"/> already has an event ID, this method skips writing the event.
|
||||
/// Otherwise, an event ID unique to all sessions and streams is generated and assigned to the event.
|
||||
/// </remarks>
|
||||
ValueTask<SseItem<JsonRpcMessage?>> WriteEventAsync(SseItem<JsonRpcMessage?> sseItem, CancellationToken cancellationToken = default);
|
||||
}
|
||||
|
|
@ -561,7 +561,7 @@ internal sealed partial class McpServerImpl : McpServer
|
|||
|
||||
try
|
||||
{
|
||||
return await handler(request, cancellationToken);
|
||||
return await handler(request, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
catch (Exception e) when (e is not OperationCanceledException and not McpProtocolException)
|
||||
{
|
||||
|
|
|
|||
|
|
@ -0,0 +1,70 @@
|
|||
using ModelContextProtocol.Protocol;
|
||||
using System.Buffers;
|
||||
using System.Net.ServerSentEvents;
|
||||
using System.Text.Json;
|
||||
|
||||
namespace ModelContextProtocol.Server;
|
||||
|
||||
/// <summary>
|
||||
/// Provides MCP extension methods for <see cref="SseEventWriter"/>.
|
||||
/// </summary>
|
||||
internal static class McpSseEventWriterExtensions
|
||||
{
|
||||
[ThreadStatic]
|
||||
private static Utf8JsonWriter? _jsonWriter;
|
||||
|
||||
/// <summary>
|
||||
/// Writes an SSE item containing a <see cref="JsonRpcMessage"/>.
|
||||
/// </summary>
|
||||
/// <param name="writer">The <see cref="SseEventWriter"/>.</param>
|
||||
/// <param name="item">The SSE item containing the <see cref="JsonRpcMessage"/>.</param>
|
||||
/// <param name="cancellationToken">The token to monitor for cancellation requests.</param>
|
||||
/// <returns>A task representing the asynchronous write operation.</returns>
|
||||
public static ValueTask WriteAsync(this SseEventWriter writer, SseItem<JsonRpcMessage?> item, CancellationToken cancellationToken = default)
|
||||
=> writer.WriteAsync(item, FormatJsonRpcMessage, cancellationToken);
|
||||
|
||||
/// <summary>
|
||||
/// Writes an SSE item containing a <see cref="string"/>.
|
||||
/// </summary>
|
||||
/// <param name="writer">The <see cref="SseEventWriter"/>.</param>
|
||||
/// <param name="item">The SSE item containing the string.</param>
|
||||
/// <param name="cancellationToken">The token to monitor for cancellation requests.</param>
|
||||
/// <returns>A task representing the asynchronous write operation.</returns>
|
||||
public static ValueTask WriteAsync(this SseEventWriter writer, SseItem<string> item, CancellationToken cancellationToken = default)
|
||||
=> writer.WriteAsync(item, FormatString, cancellationToken);
|
||||
|
||||
/// <summary>
|
||||
/// Formats a <see cref="JsonRpcMessage"/> message by writing it as JSON to the buffer writer.
|
||||
/// </summary>
|
||||
private static void FormatJsonRpcMessage(SseItem<JsonRpcMessage?> item, IBufferWriter<byte> writer)
|
||||
{
|
||||
if (item.Data is null)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
if (_jsonWriter is null)
|
||||
{
|
||||
_jsonWriter = new Utf8JsonWriter(writer);
|
||||
}
|
||||
else
|
||||
{
|
||||
_jsonWriter.Reset(writer);
|
||||
}
|
||||
|
||||
JsonSerializer.Serialize(_jsonWriter, item.Data, McpJsonUtilities.JsonContext.Default.JsonRpcMessage!);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Formats a string by writing it as UTF-8 to the buffer writer.
|
||||
/// </summary>
|
||||
private static void FormatString(SseItem<string> item, IBufferWriter<byte> writer)
|
||||
{
|
||||
if (item.Data is null)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
writer.WriteUtf8String(item.Data);
|
||||
}
|
||||
}
|
||||
|
|
@ -82,4 +82,21 @@ public sealed class RequestContext<TParams>
|
|||
/// 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>
|
||||
/// <exception cref="InvalidOperationException">Thrown when the transport does not support polling.</exception>
|
||||
public async ValueTask EnablePollingAsync(TimeSpan retryInterval, CancellationToken cancellationToken = default)
|
||||
{
|
||||
if (JsonRpcRequest.Context?.RelatedTransport is not StreamableHttpPostTransport transport)
|
||||
{
|
||||
throw new InvalidOperationException("Polling is only supported for Streamable HTTP transports.");
|
||||
}
|
||||
|
||||
await transport.EnablePollingAsync(retryInterval, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
20
src/ModelContextProtocol.Core/Server/SseEventStreamMode.cs
Normal file
20
src/ModelContextProtocol.Core/Server/SseEventStreamMode.cs
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
namespace ModelContextProtocol.Server;
|
||||
|
||||
/// <summary>
|
||||
/// Represents the mode of an SSE event stream.
|
||||
/// </summary>
|
||||
public enum SseEventStreamMode
|
||||
{
|
||||
/// <summary>
|
||||
/// Causes the event stream returned by <see cref="ISseEventStreamReader.ReadEventsAsync(System.Threading.CancellationToken)"/> to only end when
|
||||
/// the associated <see cref="ISseEventStreamWriter"/> gets disposed.
|
||||
/// </summary>
|
||||
Streaming = 0,
|
||||
|
||||
/// <summary>
|
||||
/// Causes the event stream returned by <see cref="ISseEventStreamReader.ReadEventsAsync(System.Threading.CancellationToken)"/> to end
|
||||
/// after the most recent event has been consumed. This forces clients to keep making new requests in order to receive
|
||||
/// the latest messages.
|
||||
/// </summary>
|
||||
Polling = 1,
|
||||
}
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
namespace ModelContextProtocol.Server;
|
||||
|
||||
/// <summary>
|
||||
/// Configuration options for creating an SSE event stream.
|
||||
/// </summary>
|
||||
public sealed class SseEventStreamOptions
|
||||
{
|
||||
/// <summary>
|
||||
/// Gets or sets the session ID associated with the event stream.
|
||||
/// </summary>
|
||||
public required string SessionId { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the stream ID that uniquely identifies this stream within a session.
|
||||
/// </summary>
|
||||
public required string StreamId { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the mode of the event stream. Defaults to <see cref="SseEventStreamMode.Streaming"/>.
|
||||
/// </summary>
|
||||
public SseEventStreamMode Mode { get; set; }
|
||||
}
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
using ModelContextProtocol.Protocol;
|
||||
using System.Net.ServerSentEvents;
|
||||
using System.Security.Claims;
|
||||
using System.Threading.Channels;
|
||||
|
||||
|
|
@ -27,14 +28,17 @@ namespace ModelContextProtocol.Server;
|
|||
/// <param name="sessionId">The identifier corresponding to the current MCP session.</param>
|
||||
public sealed class SseResponseStreamTransport(Stream sseResponseStream, string? messageEndpoint = "/message", string? sessionId = null) : ITransport
|
||||
{
|
||||
private readonly SseWriter _sseWriter = new(messageEndpoint);
|
||||
private readonly Channel<JsonRpcMessage> _incomingChannel = Channel.CreateBounded<JsonRpcMessage>(new BoundedChannelOptions(1)
|
||||
{
|
||||
SingleReader = true,
|
||||
SingleWriter = false,
|
||||
});
|
||||
private readonly SemaphoreSlim _lock = new(1, 1);
|
||||
private readonly TaskCompletionSource<bool> _completedTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
private readonly SseEventWriter _sseWriter = new(sseResponseStream);
|
||||
|
||||
private bool _isConnected;
|
||||
private bool _disposed;
|
||||
|
||||
/// <summary>
|
||||
/// Starts the transport and writes the JSON-RPC messages sent via <see cref="SendMessageAsync"/>
|
||||
|
|
@ -45,7 +49,14 @@ public sealed class SseResponseStreamTransport(Stream sseResponseStream, string?
|
|||
public async Task RunAsync(CancellationToken cancellationToken = default)
|
||||
{
|
||||
_isConnected = true;
|
||||
await _sseWriter.WriteAllAsync(sseResponseStream, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// Write the endpoint event first
|
||||
if (messageEndpoint is not null)
|
||||
{
|
||||
await _sseWriter.WriteAsync(SseItem.Endpoint(messageEndpoint), cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
await _completedTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
|
|
@ -57,17 +68,34 @@ public sealed class SseResponseStreamTransport(Stream sseResponseStream, string?
|
|||
/// <inheritdoc/>
|
||||
public async ValueTask DisposeAsync()
|
||||
{
|
||||
using var _ = await _lock.LockAsync().ConfigureAwait(false);
|
||||
|
||||
if (_disposed)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
_disposed = true;
|
||||
_isConnected = false;
|
||||
_incomingChannel.Writer.TryComplete();
|
||||
await _sseWriter.DisposeAsync().ConfigureAwait(false);
|
||||
_completedTcs.TrySetResult(true);
|
||||
_sseWriter.Dispose();
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
|
||||
{
|
||||
Throw.IfNull(message);
|
||||
// If the underlying writer has been disposed, just drop the message.
|
||||
await _sseWriter.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
using var _ = await _lock.LockAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// If disposed, just drop the message.
|
||||
if (_disposed)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
await _sseWriter.WriteAsync(SseItem.Message(message), cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
|
|
|
|||
|
|
@ -1,120 +0,0 @@
|
|||
using ModelContextProtocol.Protocol;
|
||||
using System.Buffers;
|
||||
using System.Net.ServerSentEvents;
|
||||
using System.Text;
|
||||
using System.Text.Json;
|
||||
using System.Threading.Channels;
|
||||
|
||||
namespace ModelContextProtocol.Server;
|
||||
|
||||
internal sealed class SseWriter(string? messageEndpoint = null, BoundedChannelOptions? channelOptions = null) : IAsyncDisposable
|
||||
{
|
||||
private readonly Channel<SseItem<JsonRpcMessage?>> _messages = Channel.CreateBounded<SseItem<JsonRpcMessage?>>(channelOptions ?? new BoundedChannelOptions(1)
|
||||
{
|
||||
SingleReader = true,
|
||||
SingleWriter = false,
|
||||
});
|
||||
|
||||
private Utf8JsonWriter? _jsonWriter;
|
||||
private Task? _writeTask;
|
||||
private CancellationToken? _writeCancellationToken;
|
||||
|
||||
private readonly SemaphoreSlim _disposeLock = new(1, 1);
|
||||
private bool _disposed;
|
||||
|
||||
public Func<IAsyncEnumerable<SseItem<JsonRpcMessage?>>, CancellationToken, IAsyncEnumerable<SseItem<JsonRpcMessage?>>>? MessageFilter { get; set; }
|
||||
|
||||
public Task WriteAllAsync(Stream sseResponseStream, CancellationToken cancellationToken)
|
||||
{
|
||||
Throw.IfNull(sseResponseStream);
|
||||
|
||||
// When messageEndpoint is set, 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.
|
||||
if (messageEndpoint is not null && !_messages.Writer.TryWrite(new SseItem<JsonRpcMessage?>(null, "endpoint")))
|
||||
{
|
||||
throw new InvalidOperationException("You must call RunAsync before calling SendMessageAsync.");
|
||||
}
|
||||
|
||||
_writeCancellationToken = cancellationToken;
|
||||
|
||||
var messages = _messages.Reader.ReadAllAsync(cancellationToken);
|
||||
if (MessageFilter is not null)
|
||||
{
|
||||
messages = MessageFilter(messages, cancellationToken);
|
||||
}
|
||||
|
||||
_writeTask = SseFormatter.WriteAsync(messages, sseResponseStream, WriteJsonRpcMessageToBuffer, cancellationToken);
|
||||
return _writeTask;
|
||||
}
|
||||
|
||||
public async Task<bool> SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
|
||||
{
|
||||
Throw.IfNull(message);
|
||||
|
||||
using var _ = await _disposeLock.LockAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
if (_disposed)
|
||||
{
|
||||
// Don't throw ObjectDisposedException here; just return false to indicate the message wasn't sent.
|
||||
// The calling transport can determine what to do in this case (drop the message, or fall back to another transport).
|
||||
return false;
|
||||
}
|
||||
|
||||
// Emit redundant "event: message" lines for better compatibility with other SDKs.
|
||||
await _messages.Writer.WriteAsync(new SseItem<JsonRpcMessage?>(message, SseParser.EventTypeDefault), cancellationToken).ConfigureAwait(false);
|
||||
return true;
|
||||
}
|
||||
|
||||
public async ValueTask DisposeAsync()
|
||||
{
|
||||
using var _ = await _disposeLock.LockAsync().ConfigureAwait(false);
|
||||
|
||||
if (_disposed)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
_messages.Writer.Complete();
|
||||
try
|
||||
{
|
||||
if (_writeTask is not null)
|
||||
{
|
||||
await _writeTask.ConfigureAwait(false);
|
||||
}
|
||||
}
|
||||
catch (OperationCanceledException) when (_writeCancellationToken?.IsCancellationRequested == true)
|
||||
{
|
||||
// Ignore exceptions caused by intentional cancellation during shutdown.
|
||||
}
|
||||
finally
|
||||
{
|
||||
_jsonWriter?.Dispose();
|
||||
_disposed = true;
|
||||
}
|
||||
}
|
||||
|
||||
private void WriteJsonRpcMessageToBuffer(SseItem<JsonRpcMessage?> item, IBufferWriter<byte> writer)
|
||||
{
|
||||
if (item.EventType == "endpoint" && messageEndpoint is not null)
|
||||
{
|
||||
writer.Write(Encoding.UTF8.GetBytes(messageEndpoint));
|
||||
return;
|
||||
}
|
||||
|
||||
JsonSerializer.Serialize(GetUtf8JsonWriter(writer), item.Data, McpJsonUtilities.JsonContext.Default.JsonRpcMessage!);
|
||||
}
|
||||
|
||||
private Utf8JsonWriter GetUtf8JsonWriter(IBufferWriter<byte> writer)
|
||||
{
|
||||
if (_jsonWriter is null)
|
||||
{
|
||||
_jsonWriter = new Utf8JsonWriter(writer);
|
||||
}
|
||||
else
|
||||
{
|
||||
_jsonWriter.Reset(writer);
|
||||
}
|
||||
|
||||
return _jsonWriter;
|
||||
}
|
||||
}
|
||||
|
|
@ -1,9 +1,6 @@
|
|||
using ModelContextProtocol.Protocol;
|
||||
using ModelContextProtocol.Protocol;
|
||||
using System.Diagnostics;
|
||||
using System.IO.Pipelines;
|
||||
using System.Net.ServerSentEvents;
|
||||
using System.Runtime.CompilerServices;
|
||||
using System.Security.Claims;
|
||||
using System.Text.Json;
|
||||
using System.Threading.Channels;
|
||||
|
||||
|
|
@ -13,17 +10,25 @@ namespace ModelContextProtocol.Server;
|
|||
/// Handles processing the request/response body pairs for the Streamable HTTP transport.
|
||||
/// This is typically used via <see cref="JsonRpcMessageContext.RelatedTransport"/>.
|
||||
/// </summary>
|
||||
internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport parentTransport, Stream responseStream) : ITransport
|
||||
internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport parentTransport, Stream responseStream, CancellationToken sessionCancellationToken) : ITransport
|
||||
{
|
||||
private readonly SseWriter _sseWriter = new();
|
||||
private readonly SemaphoreSlim _messageLock = new(1, 1);
|
||||
private readonly TaskCompletionSource<bool> _httpResponseTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
private readonly SseEventWriter _httpSseWriter = new(responseStream);
|
||||
|
||||
private TaskCompletionSource<bool>? _storeStreamTcs;
|
||||
private ISseEventStreamWriter? _storeSseWriter;
|
||||
|
||||
private RequestId _pendingRequest;
|
||||
private bool _finalResponseMessageSent;
|
||||
private bool _httpResponseCompleted;
|
||||
|
||||
public ChannelReader<JsonRpcMessage> MessageReader => throw new NotSupportedException("JsonRpcMessage.Context.RelatedTransport should only be used for sending messages.");
|
||||
|
||||
string? ITransport.SessionId => parentTransport.SessionId;
|
||||
|
||||
/// <returns>
|
||||
/// True, if data was written to the respond body.
|
||||
/// True, if data was written to the response body.
|
||||
/// False, if nothing was written because the request body did not contain any <see cref="JsonRpcRequest"/> messages to respond to.
|
||||
/// The HTTP application should typically respond with an empty "202 Accepted" response in this scenario.
|
||||
/// </returns>
|
||||
|
|
@ -35,11 +40,11 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
|
|||
{
|
||||
_pendingRequest = request.Id;
|
||||
|
||||
// Invoke the initialize request callback if applicable.
|
||||
if (parentTransport.OnInitRequestReceived is { } onInitRequest && request.Method == RequestMethods.Initialize)
|
||||
// Invoke the initialize request handler if applicable.
|
||||
if (request.Method == RequestMethods.Initialize)
|
||||
{
|
||||
var initializeRequest = JsonSerializer.Deserialize(request.Params, McpJsonUtilities.JsonContext.Default.InitializeRequestParams);
|
||||
await onInitRequest(initializeRequest).ConfigureAwait(false);
|
||||
await parentTransport.HandleInitRequestAsync(initializeRequest).ConfigureAwait(false);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -51,15 +56,28 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
|
|||
message.Context.ExecutionContext = ExecutionContext.Capture();
|
||||
}
|
||||
|
||||
await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
if (_pendingRequest.Id is null)
|
||||
{
|
||||
await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
return false;
|
||||
}
|
||||
|
||||
_sseWriter.MessageFilter = StopOnFinalResponseFilter;
|
||||
await _sseWriter.WriteAllAsync(responseStream, cancellationToken).ConfigureAwait(false);
|
||||
using (await _messageLock.LockAsync(cancellationToken).ConfigureAwait(false))
|
||||
{
|
||||
var primingItem = await TryStartSseEventStreamAsync(_pendingRequest).ConfigureAwait(false);
|
||||
if (primingItem.HasValue)
|
||||
{
|
||||
await _httpSseWriter.WriteAsync(primingItem.Value, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
// Ensure that we've sent the priming event before processing the incoming request.
|
||||
await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
// Wait for the response to be written before returning from the handler.
|
||||
// This keeps the HTTP response open until the final response message is sent.
|
||||
await _httpResponseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
|
|
@ -72,31 +90,136 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
|
|||
throw new InvalidOperationException("Server to client requests are not supported in stateless mode.");
|
||||
}
|
||||
|
||||
bool isAccepted = await _sseWriter.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
if (!isAccepted)
|
||||
using var _ = await _messageLock.LockAsync().ConfigureAwait(false);
|
||||
|
||||
try
|
||||
{
|
||||
// The underlying writer didn't accept the message because the underlying request has completed.
|
||||
// Rather than drop the message, fall back to sending it via the parent transport.
|
||||
await parentTransport.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
if (_finalResponseMessageSent)
|
||||
{
|
||||
// The final response message has already been sent.
|
||||
// Rather than drop the message, fall back to sending it via the parent transport.
|
||||
await parentTransport.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
return;
|
||||
}
|
||||
|
||||
var item = new SseItem<JsonRpcMessage?>(message, SseParser.EventTypeDefault);
|
||||
|
||||
if (_storeSseWriter is not null)
|
||||
{
|
||||
item = await _storeSseWriter.WriteEventAsync(item, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
if (!_httpResponseCompleted)
|
||||
{
|
||||
// Only write the message to the response if the response has not completed.
|
||||
|
||||
try
|
||||
{
|
||||
await _httpSseWriter.WriteAsync(item, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
catch (Exception ex) when (!cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
_httpResponseTcs.TrySetException(ex);
|
||||
}
|
||||
}
|
||||
}
|
||||
finally
|
||||
{
|
||||
// Complete the response if this is the final message.
|
||||
if ((message is JsonRpcResponse or JsonRpcError) && ((JsonRpcMessageWithId)message).Id == _pendingRequest)
|
||||
{
|
||||
_finalResponseMessageSent = true;
|
||||
_httpResponseTcs.TrySetResult(true);
|
||||
_storeStreamTcs?.TrySetResult(true);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public async ValueTask EnablePollingAsync(TimeSpan retryInterval, CancellationToken cancellationToken)
|
||||
{
|
||||
if (parentTransport.Stateless)
|
||||
{
|
||||
throw new InvalidOperationException("Polling is not supported in stateless mode.");
|
||||
}
|
||||
|
||||
using var _ = await _messageLock.LockAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
if (_storeSseWriter is null)
|
||||
{
|
||||
throw new InvalidOperationException($"Polling requires an event stream store to be configured.");
|
||||
}
|
||||
|
||||
// Send the priming event with the new retry interval.
|
||||
var primingItem = await _storeSseWriter.WriteEventAsync(
|
||||
sseItem: new SseItem<JsonRpcMessage?>() { ReconnectionInterval = retryInterval },
|
||||
cancellationToken)
|
||||
.ConfigureAwait(false);
|
||||
|
||||
// Write to the response stream if it still exists.
|
||||
if (!_httpResponseCompleted)
|
||||
{
|
||||
await _httpSseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
// 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 _storeSseWriter.SetModeAsync(SseEventStreamMode.Polling, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// Signal completion so HandlePostAsync can return.
|
||||
_httpResponseTcs.TrySetResult(true);
|
||||
}
|
||||
|
||||
private async ValueTask<SseItem<JsonRpcMessage?>?> TryStartSseEventStreamAsync(RequestId requestId)
|
||||
{
|
||||
Debug.Assert(_storeSseWriter is null);
|
||||
|
||||
_storeSseWriter = await parentTransport.TryCreateEventStreamAsync(
|
||||
streamId: requestId.Id!.ToString()!,
|
||||
cancellationToken: sessionCancellationToken)
|
||||
.ConfigureAwait(false);
|
||||
|
||||
if (_storeSseWriter is null)
|
||||
{
|
||||
return null;
|
||||
}
|
||||
|
||||
_storeStreamTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
_ = HandleStoreStreamDisposalAsync(_storeStreamTcs.Task);
|
||||
|
||||
return await _storeSseWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), sessionCancellationToken).ConfigureAwait(false);
|
||||
|
||||
async Task HandleStoreStreamDisposalAsync(Task streamTask)
|
||||
{
|
||||
try
|
||||
{
|
||||
await streamTask.WaitAsync(sessionCancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
finally
|
||||
{
|
||||
using var _ = await _messageLock.LockAsync().ConfigureAwait(false);
|
||||
|
||||
await _storeSseWriter!.DisposeAsync().ConfigureAwait(false);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public async ValueTask DisposeAsync()
|
||||
{
|
||||
await _sseWriter.DisposeAsync().ConfigureAwait(false);
|
||||
}
|
||||
using var _ = await _messageLock.LockAsync().ConfigureAwait(false);
|
||||
|
||||
private async IAsyncEnumerable<SseItem<JsonRpcMessage?>> StopOnFinalResponseFilter(IAsyncEnumerable<SseItem<JsonRpcMessage?>> messages, [EnumeratorCancellation] CancellationToken cancellationToken)
|
||||
{
|
||||
await foreach (var message in messages.WithCancellation(cancellationToken))
|
||||
if (_httpResponseCompleted)
|
||||
{
|
||||
yield return message;
|
||||
|
||||
if (message.Data is JsonRpcResponse or JsonRpcError && ((JsonRpcMessageWithId)message.Data).Id == _pendingRequest)
|
||||
{
|
||||
// Complete the SSE response stream now that all pending requests have been processed.
|
||||
break;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
_httpResponseCompleted = true;
|
||||
|
||||
_httpResponseTcs.TrySetResult(true);
|
||||
|
||||
_httpSseWriter.Dispose();
|
||||
|
||||
// Don't dispose the event stream writer here, as we may continue to write to the event store
|
||||
// after disposal if there are pending messages.
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
using ModelContextProtocol.Protocol;
|
||||
using System.IO.Pipelines;
|
||||
using System.Diagnostics;
|
||||
using System.Diagnostics.CodeAnalysis;
|
||||
using System.Net.ServerSentEvents;
|
||||
using System.Security.Claims;
|
||||
using System.Threading.Channels;
|
||||
|
||||
|
|
@ -21,24 +23,27 @@ namespace ModelContextProtocol.Server;
|
|||
/// </remarks>
|
||||
public sealed class StreamableHttpServerTransport : ITransport
|
||||
{
|
||||
// For JsonRpcMessages without a RelatedTransport, we don't want to block just because the client didn't make a GET request to handle unsolicited messages.
|
||||
private readonly SseWriter _sseWriter = new(channelOptions: new BoundedChannelOptions(1)
|
||||
{
|
||||
SingleReader = true,
|
||||
SingleWriter = false,
|
||||
FullMode = BoundedChannelFullMode.DropOldest,
|
||||
});
|
||||
/// <summary>
|
||||
/// The stream ID used for unsolicited messages sent via the standalone GET SSE stream.
|
||||
/// </summary>
|
||||
public static readonly string UnsolicitedMessageStreamId = "__get__";
|
||||
|
||||
private readonly Channel<JsonRpcMessage> _incomingChannel = Channel.CreateBounded<JsonRpcMessage>(new BoundedChannelOptions(1)
|
||||
{
|
||||
SingleReader = true,
|
||||
SingleWriter = false,
|
||||
});
|
||||
private readonly CancellationTokenSource _disposeCts = new();
|
||||
private readonly CancellationTokenSource _transportDisposedCts = new();
|
||||
private readonly SemaphoreSlim _unsolicitedMessageLock = new(1, 1);
|
||||
|
||||
private int _getRequestStarted;
|
||||
private SseEventWriter? _httpSseWriter;
|
||||
private ISseEventStreamWriter? _storeSseWriter;
|
||||
private TaskCompletionSource<bool>? _httpResponseTcs;
|
||||
private bool _getHttpRequestStarted;
|
||||
private bool _getHttpResponseCompleted;
|
||||
|
||||
/// <inheritdoc/>
|
||||
public string? SessionId { get; set; }
|
||||
public string? SessionId { get; init; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or initializes a value that indicates whether the transport should be in stateless mode that does not require all requests for a given session
|
||||
|
|
@ -59,15 +64,30 @@ public sealed class StreamableHttpServerTransport : ITransport
|
|||
public bool FlowExecutionContextFromRequests { get; init; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets a callback to be invoked before handling the initialize request.
|
||||
/// Gets or sets the event store for resumability support.
|
||||
/// When set, events are stored and can be replayed when clients reconnect with a Last-Event-ID header.
|
||||
/// </summary>
|
||||
public Func<InitializeRequestParams?, ValueTask>? OnInitRequestReceived { get; set; }
|
||||
public ISseEventStreamStore? EventStreamStore { get; init; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the negotiated protocol version for this session.
|
||||
/// </summary>
|
||||
internal string? NegotiatedProtocolVersion { get; private set; }
|
||||
|
||||
/// <inheritdoc/>
|
||||
public ChannelReader<JsonRpcMessage> MessageReader => _incomingChannel.Reader;
|
||||
|
||||
internal ChannelWriter<JsonRpcMessage> MessageWriter => _incomingChannel.Writer;
|
||||
|
||||
/// <summary>
|
||||
/// Handles the initialize request by capturing the protocol version and invoking the user callback.
|
||||
/// </summary>
|
||||
internal async ValueTask HandleInitRequestAsync(InitializeRequestParams? initParams)
|
||||
{
|
||||
// Capture the negotiated protocol version for resumability checks
|
||||
NegotiatedProtocolVersion = initParams?.ProtocolVersion;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Handles an optional SSE GET request a client using the Streamable HTTP transport might make by
|
||||
/// writing any unsolicited JSON-RPC messages sent via <see cref="SendMessageAsync"/>
|
||||
|
|
@ -78,8 +98,7 @@ public sealed class StreamableHttpServerTransport : ITransport
|
|||
/// <returns>A task representing the send loop that writes JSON-RPC messages to the SSE response stream.</returns>
|
||||
/// <exception cref="ArgumentNullException"><paramref name="sseResponseStream"/> is <see langword="null"/>.</exception>
|
||||
/// <exception cref="InvalidOperationException">
|
||||
/// <see cref="Stateless"/> is <see langword="true"/> and GET requests are not supported in stateless mode,
|
||||
/// or a GET request has already been started for this session.
|
||||
/// <see cref="Stateless"/> is <see langword="true"/> and GET requests are not supported in stateless mode.
|
||||
/// </exception>
|
||||
public async Task HandleGetRequestAsync(Stream sseResponseStream, CancellationToken cancellationToken = default)
|
||||
{
|
||||
|
|
@ -90,13 +109,27 @@ public sealed class StreamableHttpServerTransport : ITransport
|
|||
throw new InvalidOperationException("GET requests are not supported in stateless mode.");
|
||||
}
|
||||
|
||||
if (Interlocked.Exchange(ref _getRequestStarted, 1) == 1)
|
||||
using (await _unsolicitedMessageLock.LockAsync(cancellationToken).ConfigureAwait(false))
|
||||
{
|
||||
throw new InvalidOperationException("Session resumption is not yet supported. Please start a new session.");
|
||||
if (_getHttpRequestStarted)
|
||||
{
|
||||
throw new InvalidOperationException("Session resumption is not yet supported. Please start a new session.");
|
||||
}
|
||||
|
||||
_getHttpRequestStarted = true;
|
||||
_httpSseWriter = new SseEventWriter(sseResponseStream);
|
||||
_httpResponseTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
_storeSseWriter = await TryCreateEventStreamAsync(streamId: UnsolicitedMessageStreamId, cancellationToken).ConfigureAwait(false);
|
||||
if (_storeSseWriter is not null)
|
||||
{
|
||||
var primingItem = await _storeSseWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), cancellationToken).ConfigureAwait(false);
|
||||
await _httpSseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
}
|
||||
|
||||
// We do not need to reference _disposeCts like in HandlePostRequest, because the session ending completes the _sseWriter gracefully.
|
||||
await _sseWriter.WriteAllAsync(sseResponseStream, cancellationToken).ConfigureAwait(false);
|
||||
// Wait for the response to be written before returning from the handler.
|
||||
// This keeps the HTTP response open until the final response message is sent.
|
||||
await _httpResponseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
|
|
@ -122,9 +155,15 @@ public sealed class StreamableHttpServerTransport : ITransport
|
|||
Throw.IfNull(message);
|
||||
Throw.IfNull(responseStream);
|
||||
|
||||
using var postCts = CancellationTokenSource.CreateLinkedTokenSource(_disposeCts.Token, cancellationToken);
|
||||
await using var postTransport = new StreamableHttpPostTransport(this, responseStream);
|
||||
return await postTransport.HandlePostAsync(message, postCts.Token).ConfigureAwait(false);
|
||||
var postTransport = new StreamableHttpPostTransport(this, responseStream, _transportDisposedCts.Token);
|
||||
using var postCts = CancellationTokenSource.CreateLinkedTokenSource(_transportDisposedCts.Token, cancellationToken);
|
||||
await using (postTransport.ConfigureAwait(false))
|
||||
{
|
||||
return await postTransport.HandlePostAsync(
|
||||
message,
|
||||
cancellationToken: postCts.Token)
|
||||
.ConfigureAwait(false);
|
||||
}
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
|
|
@ -137,28 +176,94 @@ public sealed class StreamableHttpServerTransport : ITransport
|
|||
throw new InvalidOperationException("Unsolicited server to client messages are not supported in stateless mode.");
|
||||
}
|
||||
|
||||
// If the underlying writer has been disposed, just drop the message.
|
||||
await _sseWriter.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
using var _ = await _unsolicitedMessageLock.LockAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
if (!_getHttpRequestStarted)
|
||||
{
|
||||
// Clients are not required to make a GET request for unsolicited messages.
|
||||
// If no GET request has been made, drop the message.
|
||||
return;
|
||||
}
|
||||
|
||||
Debug.Assert(_httpSseWriter is not null);
|
||||
Debug.Assert(_httpResponseTcs is not null);
|
||||
|
||||
var item = SseItem.Message(message);
|
||||
|
||||
if (_storeSseWriter is not null)
|
||||
{
|
||||
item = await _storeSseWriter.WriteEventAsync(item, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
if (!_getHttpResponseCompleted)
|
||||
{
|
||||
// Only write the message to the response if the response has not completed.
|
||||
|
||||
try
|
||||
{
|
||||
await _httpSseWriter!.WriteAsync(item, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
catch (Exception ex) when (!cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
_httpResponseTcs!.TrySetException(ex);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public async ValueTask DisposeAsync()
|
||||
{
|
||||
using var _ = await _unsolicitedMessageLock.LockAsync().ConfigureAwait(false);
|
||||
|
||||
if (_getHttpResponseCompleted)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
_getHttpResponseCompleted = true;
|
||||
|
||||
try
|
||||
{
|
||||
_incomingChannel.Writer.TryComplete();
|
||||
await _disposeCts.CancelAsync();
|
||||
await _transportDisposedCts.CancelAsync().ConfigureAwait(false);
|
||||
}
|
||||
finally
|
||||
{
|
||||
try
|
||||
{
|
||||
await _sseWriter.DisposeAsync().ConfigureAwait(false);
|
||||
_httpResponseTcs?.TrySetResult(true);
|
||||
_httpSseWriter?.Dispose();
|
||||
|
||||
if (_storeSseWriter is not null)
|
||||
{
|
||||
await _storeSseWriter.DisposeAsync().ConfigureAwait(false);
|
||||
}
|
||||
}
|
||||
finally
|
||||
{
|
||||
_disposeCts.Dispose();
|
||||
_transportDisposedCts.Dispose();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
internal async ValueTask<ISseEventStreamWriter?> TryCreateEventStreamAsync(string streamId, CancellationToken cancellationToken)
|
||||
{
|
||||
if (EventStreamStore is null || !McpSessionHandler.SupportsPrimingEvent(NegotiatedProtocolVersion))
|
||||
{
|
||||
return null;
|
||||
}
|
||||
|
||||
// We use the 'Streaming' stream mode so that in the case of an unexpected network disconnection,
|
||||
// the client can continue reading the remaining messages in a single, streamed response.
|
||||
const SseEventStreamMode Mode = SseEventStreamMode.Streaming;
|
||||
|
||||
var sseEventStreamWriter = await EventStreamStore.CreateStreamAsync(new SseEventStreamOptions
|
||||
{
|
||||
SessionId = SessionId ?? Guid.NewGuid().ToString("N"),
|
||||
StreamId = streamId,
|
||||
Mode = Mode,
|
||||
}, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
return sseEventStreamWriter;
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue