Add middleware and authz support for server-side handlers (#733)

This commit is contained in:
Stephen Halter 2025-09-10 10:07:45 -07:00 • committed by GitHub
commit d344c65120
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1,7 +1,9 @@
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;
@ -9,14 +11,14 @@ namespace ModelContextProtocol.Server;
/// <summary>
/// Handles processing the request/response body pairs for the Streamable HTTP transport.
/// This is typically used via <see cref="JsonRpcMessage.RelatedTransport"/>.
/// This is typically used via <see cref="JsonRpcMessageContext.RelatedTransport"/>.
/// </summary>
internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport parentTransport, IDuplexPipe httpBodies) : ITransport
internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport parentTransport, Stream responseStream) : ITransport
{
private readonly SseWriter _sseWriter = new();
private RequestId _pendingRequest;
public ChannelReader<JsonRpcMessage> MessageReader => throw new NotSupportedException("JsonRpcMessage.RelatedTransport should only be used for sending messages.");
public ChannelReader<JsonRpcMessage> MessageReader => throw new NotSupportedException("JsonRpcMessage.Context.RelatedTransport should only be used for sending messages.");
string? ITransport.SessionId => parentTransport.SessionId;
@ -25,11 +27,31 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
/// 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>
public async ValueTask<bool> RunAsync(CancellationToken cancellationToken)
public async ValueTask<bool> HandlePostAsync(JsonRpcMessage message, CancellationToken cancellationToken)
{
var message = await JsonSerializer.DeserializeAsync(httpBodies.Input.AsStream(),
McpJsonUtilities.JsonContext.Default.JsonRpcMessage, cancellationToken).ConfigureAwait(false);
await OnMessageReceivedAsync(message, cancellationToken).ConfigureAwait(false);
Debug.Assert(_pendingRequest.Id is null);
if (message is JsonRpcRequest request)
{
_pendingRequest = request.Id;
// Invoke the initialize request callback if applicable.
if (parentTransport.OnInitRequestReceived is { } onInitRequest && request.Method == RequestMethods.Initialize)
{
var initializeRequest = JsonSerializer.Deserialize(request.Params, McpJsonUtilities.JsonContext.Default.InitializeRequestParams);
await onInitRequest(initializeRequest).ConfigureAwait(false);
}
}
message.Context ??= new JsonRpcMessageContext();
message.Context.RelatedTransport = this;
if (parentTransport.FlowExecutionContextFromRequests)
{
message.Context.ExecutionContext = ExecutionContext.Capture();
}
await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);
if (_pendingRequest.Id is null)
{
@ -37,12 +59,14 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
}
_sseWriter.MessageFilter = StopOnFinalResponseFilter;
await _sseWriter.WriteAllAsync(httpBodies.Output.AsStream(), cancellationToken).ConfigureAwait(false);
await _sseWriter.WriteAllAsync(responseStream, cancellationToken).ConfigureAwait(false);
return true;
}
public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
{
Throw.IfNull(message);
if (parentTransport.Stateless && message is JsonRpcRequest)
{
throw new InvalidOperationException("Server to client requests are not supported in stateless mode.");
@ -69,33 +93,4 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
}
}
}
private async ValueTask OnMessageReceivedAsync(JsonRpcMessage? message, CancellationToken cancellationToken)
{
if (message is null)
{
throw new InvalidOperationException("Received invalid null message.");
}
if (message is JsonRpcRequest request)
{
_pendingRequest = request.Id;
// Invoke the initialize request callback if applicable.
if (parentTransport.OnInitRequestReceived is { } onInitRequest && request.Method == RequestMethods.Initialize)
{
var initializeRequest = JsonSerializer.Deserialize(request.Params, McpJsonUtilities.JsonContext.Default.InitializeRequestParams);
await onInitRequest(initializeRequest).ConfigureAwait(false);
}
}
message.RelatedTransport = this;
if (parentTransport.FlowExecutionContextFromRequests)
{
message.ExecutionContext = ExecutionContext.Capture();
}
await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);
}
}