Add server-side Streamable HTTP transport support (#330)
- Changes IJsonRpcMessage to an abstract base class so RelatedTransport will always be available - Streamable HTTP supports multiple concurrent HTTP request with their own indpendent SSE response streams - RelatedTransport indicates the source or destination of the JsonRpcMessage - Changes the default RequestId to a JSON number for better compatibility with MCP servers in the wild
This commit is contained in:
Stephen Halter
GitHub
parent
7d77be68cf
commit
7dcba357d4
1 changed files with 1637 additions and 395 deletions
|
|
@ -1,10 +1,6 @@
|
|||
using System.Text;
|
||||
using System.Buffers;
|
||||
using System.Net.ServerSentEvents;
|
||||
using System.Text.Json;
|
||||
using System.Threading.Channels;
|
||||
using ModelContextProtocol.Protocol.Messages;
|
||||
using ModelContextProtocol.Utils.Json;
|
||||
using ModelContextProtocol.Utils;
|
||||
using System.Threading.Channels;
|
||||
|
||||
namespace ModelContextProtocol.Protocol.Transport;
|
||||
|
||||
|
|
@ -22,16 +18,22 @@ namespace ModelContextProtocol.Protocol.Transport;
|
|||
/// such as when streaming completion results or providing progress updates during long-running operations.
|
||||
/// </para>
|
||||
/// </remarks>
|
||||
public sealed class SseResponseStreamTransport(Stream sseResponseStream, string messageEndpoint = "/message") : ITransport
|
||||
/// <param name="sseResponseStream">The response stream to write MCP JSON-RPC messages as SSE events to.</param>
|
||||
/// <param name="messageEndpoint">
|
||||
/// The relative or absolute URI the client should use to post MCP JSON-RPC messages for this session.
|
||||
/// These messages should be passed to <see cref="OnMessageReceivedAsync(JsonRpcMessage, CancellationToken)"/>.
|
||||
/// Defaults to "/message".
|
||||
/// </param>
|
||||
public sealed class SseResponseStreamTransport(Stream sseResponseStream, string? messageEndpoint = "/message") : ITransport
|
||||
{
|
||||
private readonly Channel<IJsonRpcMessage> _incomingChannel = CreateBoundedChannel<IJsonRpcMessage>();
|
||||
private readonly Channel<SseItem<IJsonRpcMessage?>> _outgoingSseChannel = CreateBoundedChannel<SseItem<IJsonRpcMessage?>>();
|
||||
private readonly SseWriter _sseWriter = new(messageEndpoint);
|
||||
private readonly Channel<JsonRpcMessage> _incomingChannel = Channel.CreateBounded<JsonRpcMessage>(new BoundedChannelOptions(1)
|
||||
{
|
||||
SingleReader = true,
|
||||
SingleWriter = false,
|
||||
});
|
||||
|
||||
private Task? _sseWriteTask;
|
||||
private Utf8JsonWriter? _jsonWriter;
|
||||
|
||||
/// <inheritdoc />
|
||||
public bool IsConnected { get; private set; }
|
||||
private bool _isConnected;
|
||||
|
||||
/// <summary>
|
||||
/// Starts the transport and writes the JSON-RPC messages sent via <see cref="SendMessageAsync"/>
|
||||
|
|
@ -39,54 +41,27 @@ public sealed class SseResponseStreamTransport(Stream sseResponseStream, string
|
|||
/// </summary>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
|
||||
/// <returns>A task representing the send loop that writes JSON-RPC messages to the SSE response stream.</returns>
|
||||
public Task RunAsync(CancellationToken cancellationToken)
|
||||
public async Task RunAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
// 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 (!_outgoingSseChannel.Writer.TryWrite(new SseItem<IJsonRpcMessage?>(null, "endpoint")))
|
||||
{
|
||||
throw new InvalidOperationException($"You must call ${nameof(RunAsync)} before calling ${nameof(SendMessageAsync)}.");
|
||||
}
|
||||
|
||||
IsConnected = true;
|
||||
|
||||
var sseItems = _outgoingSseChannel.Reader.ReadAllAsync(cancellationToken);
|
||||
return _sseWriteTask = SseFormatter.WriteAsync(sseItems, sseResponseStream, WriteJsonRpcMessageToBuffer, cancellationToken);
|
||||
}
|
||||
|
||||
private void WriteJsonRpcMessageToBuffer(SseItem<IJsonRpcMessage?> item, IBufferWriter<byte> writer)
|
||||
{
|
||||
if (item.EventType == "endpoint")
|
||||
{
|
||||
writer.Write(Encoding.UTF8.GetBytes(messageEndpoint));
|
||||
return;
|
||||
}
|
||||
|
||||
JsonSerializer.Serialize(GetUtf8JsonWriter(writer), item.Data, McpJsonUtilities.JsonContext.Default.IJsonRpcMessage!);
|
||||
_isConnected = true;
|
||||
await _sseWriter.WriteAllAsync(sseResponseStream, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public ChannelReader<IJsonRpcMessage> MessageReader => _incomingChannel.Reader;
|
||||
public ChannelReader<JsonRpcMessage> MessageReader => _incomingChannel.Reader;
|
||||
|
||||
/// <inheritdoc/>
|
||||
public ValueTask DisposeAsync()
|
||||
public async ValueTask DisposeAsync()
|
||||
{
|
||||
IsConnected = false;
|
||||
_isConnected = false;
|
||||
_incomingChannel.Writer.TryComplete();
|
||||
_outgoingSseChannel.Writer.TryComplete();
|
||||
return new ValueTask(_sseWriteTask ?? Task.CompletedTask);
|
||||
await _sseWriter.DisposeAsync().ConfigureAwait(false);
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
|
||||
public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
|
||||
{
|
||||
if (!IsConnected)
|
||||
{
|
||||
throw new InvalidOperationException($"Transport is not connected. Make sure to call {nameof(RunAsync)} first.");
|
||||
}
|
||||
|
||||
// Emit redundant "event: message" lines for better compatibility with other SDKs.
|
||||
await _outgoingSseChannel.Writer.WriteAsync(new SseItem<IJsonRpcMessage?>(message, SseParser.EventTypeDefault), cancellationToken).ConfigureAwait(false);
|
||||
await _sseWriter.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
|
|
@ -111,34 +86,15 @@ public sealed class SseResponseStreamTransport(Stream sseResponseStream, string
|
|||
/// sequencing of operations in the transport lifecycle.
|
||||
/// </para>
|
||||
/// </remarks>
|
||||
public async Task OnMessageReceivedAsync(IJsonRpcMessage message, CancellationToken cancellationToken)
|
||||
public async Task OnMessageReceivedAsync(JsonRpcMessage message, CancellationToken cancellationToken)
|
||||
{
|
||||
if (!IsConnected)
|
||||
Throw.IfNull(message);
|
||||
|
||||
if (!_isConnected)
|
||||
{
|
||||
throw new InvalidOperationException($"Transport is not connected. Make sure to call {nameof(RunAsync)} first.");
|
||||
}
|
||||
|
||||
await _incomingChannel.Writer.WriteAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
private static Channel<T> CreateBoundedChannel<T>(int capacity = 1) =>
|
||||
Channel.CreateBounded<T>(new BoundedChannelOptions(capacity)
|
||||
{
|
||||
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