modelcontextprotocol/csharp-sdk
Watch
1
0
Fork
You've already forked csharp-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Use Kestrel for all in-memory HTTP tests (#225)

Browse source 
* Add route pattern parameter to MapMcp
* Add configureOptionsAsync
* Use Kestrel for all in-memory HTTP tests
* Use ApplicationStopping token for SSE responses
* Use RegisterNotificationHandler in MapMcp tests
* Work around SocketsHttpHandler bug where it doesn't call DispseAsync on the stream returned by the ConnectCallback
This commit is contained in:
Stephen Halter 2025-04-06 19:27:17 -07:00 • committed by GitHub
commit faf12b6a94
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 631 additions and 642 deletions
Download patch file Download diff file Expand all files Collapse all files

67
src/ModelContextProtocol.AspNetCore/McpEndpointRouteBuilderExtensions.cs
View file

@ -1,8 +1,10 @@
using Microsoft.AspNetCore.Http;
using Microsoft.AspNetCore.Http.Features;
using Microsoft.AspNetCore.Routing;
using Microsoft.AspNetCore.Routing.Patterns;
using Microsoft.AspNetCore.WebUtilities;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using ModelContextProtocol.Protocol.Messages;
@ -10,6 +12,7 @@ using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Server;
using ModelContextProtocol.Utils.Json;
using System.Collections.Concurrent;
using System.Diagnostics.CodeAnalysis;
using System.Security.Cryptography;
namespace Microsoft.AspNetCore.Builder;
@ -23,25 +26,54 @@ public static class McpEndpointRouteBuilderExtensions
/// Sets up endpoints for handling MCP HTTP Streaming transport.
/// </summary>
/// <param name="endpoints">The web application to attach MCP HTTP endpoints.</param>
/// <param name="runSession">Provides an optional asynchronous callback for handling new MCP sessions.</param>
/// <param name="pattern">The route pattern prefix to map to.</param>
/// <param name="configureOptionsAsync">Configure per-session options.</param>
/// <param name="runSessionAsync">Provides an optional asynchronous callback for handling new MCP sessions.</param>
/// <returns>Returns a builder for configuring additional endpoint conventions like authorization policies.</returns>
public static IEndpointConventionBuilder MapMcp(this IEndpointRouteBuilder endpoints, Func<HttpContext, IMcpServer, CancellationToken, Task>? runSession = null)
public static IEndpointConventionBuilder MapMcp(
this IEndpointRouteBuilder endpoints,
[StringSyntax("Route")] string pattern = "",
Func<HttpContext, McpServerOptions, CancellationToken, Task>? configureOptionsAsync = null,
Func<HttpContext, IMcpServer, CancellationToken, Task>? runSessionAsync = null)
=> endpoints.MapMcp(RoutePatternFactory.Parse(pattern), configureOptionsAsync, runSessionAsync);
/// <summary>
/// Sets up endpoints for handling MCP HTTP Streaming transport.
/// </summary>
/// <param name="endpoints">The web application to attach MCP HTTP endpoints.</param>
/// <param name="pattern">The route pattern prefix to map to.</param>
/// <param name="configureOptionsAsync">Configure per-session options.</param>
/// <param name="runSessionAsync">Provides an optional asynchronous callback for handling new MCP sessions.</param>
/// <returns>Returns a builder for configuring additional endpoint conventions like authorization policies.</returns>
public static IEndpointConventionBuilder MapMcp(this IEndpointRouteBuilder endpoints,
RoutePattern pattern,
Func<HttpContext, McpServerOptions, CancellationToken, Task>? configureOptionsAsync = null,
Func<HttpContext, IMcpServer, CancellationToken, Task>? runSessionAsync = null)
{
ConcurrentDictionary<string, SseResponseStreamTransport> _sessions = new(StringComparer.Ordinal);
var loggerFactory = endpoints.ServiceProvider.GetRequiredService<ILoggerFactory>();
var mcpServerOptions = endpoints.ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>();
var optionsSnapshot = endpoints.ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>();
var optionsFactory = endpoints.ServiceProvider.GetRequiredService<IOptionsFactory<McpServerOptions>>();
var hostApplicationLifetime = endpoints.ServiceProvider.GetRequiredService<IHostApplicationLifetime>();
var routeGroup = endpoints.MapGroup("");
var routeGroup = endpoints.MapGroup(pattern);
routeGroup.MapGet("/sse", async context =>
{
var response = context.Response;
var requestAborted = context.RequestAborted;
// If the server is shutting down, we need to cancel all SSE connections immediately without waiting for HostOptions.ShutdownTimeout
// which defaults to 30 seconds.
using var sseCts = CancellationTokenSource.CreateLinkedTokenSource(context.RequestAborted, hostApplicationLifetime.ApplicationStopping);
var cancellationToken = sseCts.Token;
var response = context.Response;
response.Headers.ContentType = "text/event-stream";
response.Headers.CacheControl = "no-cache,no-store";
// Make sure we disable all response buffering for SSE
context.Response.Headers.ContentEncoding = "identity";
context.Features.GetRequiredFeature<IHttpResponseBodyFeature>().DisableBuffering();
var sessionId = MakeNewSessionId();
await using var transport = new SseResponseStreamTransport(response.Body, $"/message?sessionId={sessionId}");
if (!_sessions.TryAdd(sessionId, transport))
@ -49,19 +81,24 @@ public static class McpEndpointRouteBuilderExtensions
throw new Exception($"Unreachable given good entropy! Session with ID '{sessionId}' has already been created.");
}
var options = optionsSnapshot.Value;
if (configureOptionsAsync is not null)
{
options = optionsFactory.Create(Options.DefaultName);
await configureOptionsAsync.Invoke(context, options, cancellationToken);
}
try
{
// Make sure we disable all response buffering for SSE
context.Response.Headers.ContentEncoding = "identity";
context.Features.GetRequiredFeature<IHttpResponseBodyFeature>().DisableBuffering();
var transportTask = transport.RunAsync(cancellationToken: requestAborted);
await using var server = McpServerFactory.Create(transport, mcpServerOptions.Value, loggerFactory, endpoints.ServiceProvider);
var transportTask = transport.RunAsync(cancellationToken);
try
{
runSession ??= RunSession;
await runSession(context, server, requestAborted);
await using var mcpServer = McpServerFactory.Create(transport, options, loggerFactory, endpoints.ServiceProvider);
context.Features.Set(mcpServer);
runSessionAsync ??= RunSession;
await runSessionAsync(context, mcpServer, cancellationToken);
}
finally
{
@ -69,7 +106,7 @@ public static class McpEndpointRouteBuilderExtensions
await transportTask;
}
}
catch (OperationCanceledException) when (requestAborted.IsCancellationRequested)
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
{
// RequestAborted always triggers when the client disconnects before a complete response body is written,
// but this is how SSE connections are typically closed.

5
src/ModelContextProtocol/Client/McpClient.cs
View file

@ -106,7 +106,10 @@ internal sealed class McpClient : McpEndpoint, IMcpClient
{
// Connect transport
_sessionTransport = await _clientTransport.ConnectAsync(cancellationToken).ConfigureAwait(false);
StartSession(_sessionTransport);
InitializeSession(_sessionTransport);
// We don't want the ConnectAsync token to cancel the session after we've successfully connected.
// The base class handles cleaning up the session in DisposeAsync without our help.
StartSession(_sessionTransport, fullSessionCancellationToken: CancellationToken.None);
// Perform initialization sequence
using var initializationCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);

8
src/ModelContextProtocol/Protocol/Transport/SseClientSessionTransport.cs
View file

@ -151,16 +151,14 @@ internal sealed class SseClientSessionTransport : TransportBase
{
try
{
if (!_connectionCts.IsCancellationRequested)
{
await _connectionCts.CancelAsync().ConfigureAwait(false);
_connectionCts.Dispose();
}
await _connectionCts.CancelAsync().ConfigureAwait(false);
if (_receiveTask != null)
{
await _receiveTask.ConfigureAwait(false);
}
_connectionCts.Dispose();
}
finally
{

12
src/ModelContextProtocol/Server/McpServer.cs
View file

@ -18,6 +18,8 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
Version = DefaultAssemblyName.Version?.ToString() ?? "1.0.0",
};
private readonly ITransport _sessionTransport;
private readonly EventHandler? _toolsChangedDelegate;
private readonly EventHandler? _promptsChangedDelegate;
@ -41,6 +43,7 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
options ??= new();
_sessionTransport = transport;
ServerOptions = options;
Services = serviceProvider;
_endpointName = $"Server ({options.ServerInfo?.Name ?? DefaultImplementation.Name} {options.ServerInfo?.Version ?? DefaultImplementation.Version})";
@ -81,8 +84,8 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
prompts.Changed += _promptsChangedDelegate;
}
// And start the session.
StartSession(transport);
// And initialize the session.
InitializeSession(transport);
}
public ServerCapabilities? ServerCapabilities { get; set; }
@ -112,9 +115,8 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
try
{
using var _ = cancellationToken.Register(static s => ((McpServer)s!).CancelSession(), this);
// The McpServer ctor always calls StartSession, so MessageProcessingTask is always set.
await MessageProcessingTask!.ConfigureAwait(false);
StartSession(_sessionTransport, fullSessionCancellationToken: cancellationToken);
await MessageProcessingTask.ConfigureAwait(false);
}
finally
{

15
src/ModelContextProtocol/Shared/McpEndpoint.cs
View file

@ -5,6 +5,7 @@ using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Server;
using ModelContextProtocol.Utils;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.Reflection;
@ -62,12 +63,16 @@ internal abstract class McpEndpoint : IAsyncDisposable
/// </summary>
protected Task? MessageProcessingTask { get; private set; }
[MemberNotNull(nameof(MessageProcessingTask))]
protected void StartSession(ITransport sessionTransport)
protected void InitializeSession(ITransport sessionTransport)
{
_sessionCts = new CancellationTokenSource();
_session = new McpSession(this is IMcpServer, sessionTransport, EndpointName, RequestHandlers, NotificationHandlers, _logger);
MessageProcessingTask = _session.ProcessMessagesAsync(_sessionCts.Token);
}
[MemberNotNull(nameof(MessageProcessingTask))]
protected void StartSession(ITransport sessionTransport, CancellationToken fullSessionCancellationToken)
{
_sessionCts = CancellationTokenSource.CreateLinkedTokenSource(fullSessionCancellationToken);
MessageProcessingTask = GetSessionOrThrow().ProcessMessagesAsync(_sessionCts.Token);
}
protected void CancelSession() => _sessionCts?.Cancel();
@ -122,5 +127,5 @@ internal abstract class McpEndpoint : IAsyncDisposable
}
protected McpSession GetSessionOrThrow()
=> _session ?? throw new InvalidOperationException($"This should be unreachable from public API! Call {nameof(StartSession)} before sending messages.");
=> _session ?? throw new InvalidOperationException($"This should be unreachable from public API! Call {nameof(InitializeSession)} before sending messages.");
}