Add middleware and authz support for server-side handlers (#733)
This commit is contained in:
parent
0f0aad020b
commit
d344c65120
1 changed files with 2871 additions and 435 deletions
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue