Use Kestrel for all in-memory HTTP tests (#225)
* 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:
parent
848c9df433
commit
faf12b6a94
17 changed files with 631 additions and 642 deletions
|
|
@ -1,8 +1,10 @@
|
||||||
using Microsoft.AspNetCore.Http;
|
using Microsoft.AspNetCore.Http;
|
||||||
using Microsoft.AspNetCore.Http.Features;
|
using Microsoft.AspNetCore.Http.Features;
|
||||||
using Microsoft.AspNetCore.Routing;
|
using Microsoft.AspNetCore.Routing;
|
||||||
|
using Microsoft.AspNetCore.Routing.Patterns;
|
||||||
using Microsoft.AspNetCore.WebUtilities;
|
using Microsoft.AspNetCore.WebUtilities;
|
||||||
using Microsoft.Extensions.DependencyInjection;
|
using Microsoft.Extensions.DependencyInjection;
|
||||||
|
using Microsoft.Extensions.Hosting;
|
||||||
using Microsoft.Extensions.Logging;
|
using Microsoft.Extensions.Logging;
|
||||||
using Microsoft.Extensions.Options;
|
using Microsoft.Extensions.Options;
|
||||||
using ModelContextProtocol.Protocol.Messages;
|
using ModelContextProtocol.Protocol.Messages;
|
||||||
|
|
@ -10,6 +12,7 @@ using ModelContextProtocol.Protocol.Transport;
|
||||||
using ModelContextProtocol.Server;
|
using ModelContextProtocol.Server;
|
||||||
using ModelContextProtocol.Utils.Json;
|
using ModelContextProtocol.Utils.Json;
|
||||||
using System.Collections.Concurrent;
|
using System.Collections.Concurrent;
|
||||||
|
using System.Diagnostics.CodeAnalysis;
|
||||||
using System.Security.Cryptography;
|
using System.Security.Cryptography;
|
||||||
|
|
||||||
namespace Microsoft.AspNetCore.Builder;
|
namespace Microsoft.AspNetCore.Builder;
|
||||||
|
|
@ -23,25 +26,54 @@ public static class McpEndpointRouteBuilderExtensions
|
||||||
/// Sets up endpoints for handling MCP HTTP Streaming transport.
|
/// Sets up endpoints for handling MCP HTTP Streaming transport.
|
||||||
/// </summary>
|
/// </summary>
|
||||||
/// <param name="endpoints">The web application to attach MCP HTTP endpoints.</param>
|
/// <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>
|
/// <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);
|
ConcurrentDictionary<string, SseResponseStreamTransport> _sessions = new(StringComparer.Ordinal);
|
||||||
|
|
||||||
var loggerFactory = endpoints.ServiceProvider.GetRequiredService<ILoggerFactory>();
|
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 =>
|
routeGroup.MapGet("/sse", async context =>
|
||||||
{
|
{
|
||||||
var response = context.Response;
|
// If the server is shutting down, we need to cancel all SSE connections immediately without waiting for HostOptions.ShutdownTimeout
|
||||||
var requestAborted = context.RequestAborted;
|
// 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.ContentType = "text/event-stream";
|
||||||
response.Headers.CacheControl = "no-cache,no-store";
|
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();
|
var sessionId = MakeNewSessionId();
|
||||||
await using var transport = new SseResponseStreamTransport(response.Body, $"/message?sessionId={sessionId}");
|
await using var transport = new SseResponseStreamTransport(response.Body, $"/message?sessionId={sessionId}");
|
||||||
if (!_sessions.TryAdd(sessionId, transport))
|
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.");
|
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
|
try
|
||||||
{
|
{
|
||||||
// Make sure we disable all response buffering for SSE
|
var transportTask = transport.RunAsync(cancellationToken);
|
||||||
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);
|
|
||||||
|
|
||||||
try
|
try
|
||||||
{
|
{
|
||||||
runSession ??= RunSession;
|
await using var mcpServer = McpServerFactory.Create(transport, options, loggerFactory, endpoints.ServiceProvider);
|
||||||
await runSession(context, server, requestAborted);
|
context.Features.Set(mcpServer);
|
||||||
|
|
||||||
|
runSessionAsync ??= RunSession;
|
||||||
|
await runSessionAsync(context, mcpServer, cancellationToken);
|
||||||
}
|
}
|
||||||
finally
|
finally
|
||||||
{
|
{
|
||||||
|
|
@ -69,7 +106,7 @@ public static class McpEndpointRouteBuilderExtensions
|
||||||
await transportTask;
|
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,
|
// RequestAborted always triggers when the client disconnects before a complete response body is written,
|
||||||
// but this is how SSE connections are typically closed.
|
// but this is how SSE connections are typically closed.
|
||||||
|
|
|
||||||
|
|
@ -106,7 +106,10 @@ internal sealed class McpClient : McpEndpoint, IMcpClient
|
||||||
{
|
{
|
||||||
// Connect transport
|
// Connect transport
|
||||||
_sessionTransport = await _clientTransport.ConnectAsync(cancellationToken).ConfigureAwait(false);
|
_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
|
// Perform initialization sequence
|
||||||
using var initializationCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
using var initializationCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
||||||
|
|
|
||||||
|
|
@ -151,16 +151,14 @@ internal sealed class SseClientSessionTransport : TransportBase
|
||||||
{
|
{
|
||||||
try
|
try
|
||||||
{
|
{
|
||||||
if (!_connectionCts.IsCancellationRequested)
|
await _connectionCts.CancelAsync().ConfigureAwait(false);
|
||||||
{
|
|
||||||
await _connectionCts.CancelAsync().ConfigureAwait(false);
|
|
||||||
_connectionCts.Dispose();
|
|
||||||
}
|
|
||||||
|
|
||||||
if (_receiveTask != null)
|
if (_receiveTask != null)
|
||||||
{
|
{
|
||||||
await _receiveTask.ConfigureAwait(false);
|
await _receiveTask.ConfigureAwait(false);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
_connectionCts.Dispose();
|
||||||
}
|
}
|
||||||
finally
|
finally
|
||||||
{
|
{
|
||||||
|
|
|
||||||
|
|
@ -18,6 +18,8 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
|
||||||
Version = DefaultAssemblyName.Version?.ToString() ?? "1.0.0",
|
Version = DefaultAssemblyName.Version?.ToString() ?? "1.0.0",
|
||||||
};
|
};
|
||||||
|
|
||||||
|
private readonly ITransport _sessionTransport;
|
||||||
|
|
||||||
private readonly EventHandler? _toolsChangedDelegate;
|
private readonly EventHandler? _toolsChangedDelegate;
|
||||||
private readonly EventHandler? _promptsChangedDelegate;
|
private readonly EventHandler? _promptsChangedDelegate;
|
||||||
|
|
||||||
|
|
@ -41,6 +43,7 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
|
||||||
|
|
||||||
options ??= new();
|
options ??= new();
|
||||||
|
|
||||||
|
_sessionTransport = transport;
|
||||||
ServerOptions = options;
|
ServerOptions = options;
|
||||||
Services = serviceProvider;
|
Services = serviceProvider;
|
||||||
_endpointName = $"Server ({options.ServerInfo?.Name ?? DefaultImplementation.Name} {options.ServerInfo?.Version ?? DefaultImplementation.Version})";
|
_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;
|
prompts.Changed += _promptsChangedDelegate;
|
||||||
}
|
}
|
||||||
|
|
||||||
// And start the session.
|
// And initialize the session.
|
||||||
StartSession(transport);
|
InitializeSession(transport);
|
||||||
}
|
}
|
||||||
|
|
||||||
public ServerCapabilities? ServerCapabilities { get; set; }
|
public ServerCapabilities? ServerCapabilities { get; set; }
|
||||||
|
|
@ -112,9 +115,8 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
|
||||||
|
|
||||||
try
|
try
|
||||||
{
|
{
|
||||||
using var _ = cancellationToken.Register(static s => ((McpServer)s!).CancelSession(), this);
|
StartSession(_sessionTransport, fullSessionCancellationToken: cancellationToken);
|
||||||
// The McpServer ctor always calls StartSession, so MessageProcessingTask is always set.
|
await MessageProcessingTask.ConfigureAwait(false);
|
||||||
await MessageProcessingTask!.ConfigureAwait(false);
|
|
||||||
}
|
}
|
||||||
finally
|
finally
|
||||||
{
|
{
|
||||||
|
|
|
||||||
|
|
@ -5,6 +5,7 @@ using ModelContextProtocol.Protocol.Messages;
|
||||||
using ModelContextProtocol.Protocol.Transport;
|
using ModelContextProtocol.Protocol.Transport;
|
||||||
using ModelContextProtocol.Server;
|
using ModelContextProtocol.Server;
|
||||||
using ModelContextProtocol.Utils;
|
using ModelContextProtocol.Utils;
|
||||||
|
using System.Diagnostics;
|
||||||
using System.Diagnostics.CodeAnalysis;
|
using System.Diagnostics.CodeAnalysis;
|
||||||
using System.Reflection;
|
using System.Reflection;
|
||||||
|
|
||||||
|
|
@ -62,12 +63,16 @@ internal abstract class McpEndpoint : IAsyncDisposable
|
||||||
/// </summary>
|
/// </summary>
|
||||||
protected Task? MessageProcessingTask { get; private set; }
|
protected Task? MessageProcessingTask { get; private set; }
|
||||||
|
|
||||||
[MemberNotNull(nameof(MessageProcessingTask))]
|
protected void InitializeSession(ITransport sessionTransport)
|
||||||
protected void StartSession(ITransport sessionTransport)
|
|
||||||
{
|
{
|
||||||
_sessionCts = new CancellationTokenSource();
|
|
||||||
_session = new McpSession(this is IMcpServer, sessionTransport, EndpointName, RequestHandlers, NotificationHandlers, _logger);
|
_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();
|
protected void CancelSession() => _sessionCts?.Cancel();
|
||||||
|
|
@ -122,5 +127,5 @@ internal abstract class McpEndpoint : IAsyncDisposable
|
||||||
}
|
}
|
||||||
|
|
||||||
protected McpSession GetSessionOrThrow()
|
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.");
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,5 @@
|
||||||
using ModelContextProtocol.Protocol.Types;
|
using Microsoft.AspNetCore.Connections;
|
||||||
|
using ModelContextProtocol.Protocol.Types;
|
||||||
using ModelContextProtocol.Server;
|
using ModelContextProtocol.Server;
|
||||||
using Serilog;
|
using Serilog;
|
||||||
using System.Text;
|
using System.Text;
|
||||||
|
|
@ -372,18 +373,34 @@ public class Program
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
public static async Task MainAsync(string[] args, ILoggerProvider? loggerProvider = null, CancellationToken cancellationToken = default)
|
public static async Task MainAsync(string[] args, ILoggerProvider? loggerProvider = null, IConnectionListenerFactory? kestrelTransport = null, CancellationToken cancellationToken = default)
|
||||||
{
|
{
|
||||||
Console.WriteLine("Starting server...");
|
Console.WriteLine("Starting server...");
|
||||||
|
|
||||||
int port = args.Length > 0 && uint.TryParse(args[0], out var parsedPort) ? (int)parsedPort : 3001;
|
var builder = WebApplication.CreateEmptyBuilder(new()
|
||||||
|
|
||||||
var builder = WebApplication.CreateSlimBuilder(args);
|
|
||||||
builder.WebHost.ConfigureKestrel(options =>
|
|
||||||
{
|
{
|
||||||
options.ListenLocalhost(port);
|
Args = args,
|
||||||
});
|
});
|
||||||
|
|
||||||
|
if (kestrelTransport is null)
|
||||||
|
{
|
||||||
|
int port = args.Length > 0 && uint.TryParse(args[0], out var parsedPort) ? (int)parsedPort : 3001;
|
||||||
|
builder.WebHost.ConfigureKestrel(options =>
|
||||||
|
{
|
||||||
|
options.ListenLocalhost(port);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
else
|
||||||
|
{
|
||||||
|
// Add passed-in transport before calling UseKestrelCore() to avoid the SocketsHttpHandler getting added.
|
||||||
|
builder.Services.AddSingleton(kestrelTransport);
|
||||||
|
}
|
||||||
|
|
||||||
|
builder.WebHost.UseKestrelCore();
|
||||||
|
builder.Services.AddLogging();
|
||||||
|
builder.Services.AddRoutingCore();
|
||||||
|
|
||||||
|
builder.Logging.AddConsole();
|
||||||
ConfigureSerilog(builder.Logging);
|
ConfigureSerilog(builder.Logging);
|
||||||
if (loggerProvider is not null)
|
if (loggerProvider is not null)
|
||||||
{
|
{
|
||||||
|
|
@ -393,6 +410,8 @@ public class Program
|
||||||
builder.Services.AddMcpServer(ConfigureOptions);
|
builder.Services.AddMcpServer(ConfigureOptions);
|
||||||
|
|
||||||
var app = builder.Build();
|
var app = builder.Build();
|
||||||
|
app.UseRouting();
|
||||||
|
app.UseEndpoints(_ => { });
|
||||||
|
|
||||||
app.MapMcp();
|
app.MapMcp();
|
||||||
|
|
||||||
|
|
|
||||||
123
tests/ModelContextProtocol.Tests/DockerEverythingServerTests.cs
Normal file
123
tests/ModelContextProtocol.Tests/DockerEverythingServerTests.cs
Normal file
|
|
@ -0,0 +1,123 @@
|
||||||
|
using ModelContextProtocol.Client;
|
||||||
|
using ModelContextProtocol.Protocol.Transport;
|
||||||
|
using ModelContextProtocol.Protocol.Types;
|
||||||
|
using ModelContextProtocol.Tests.Utils;
|
||||||
|
|
||||||
|
namespace ModelContextProtocol.Tests;
|
||||||
|
|
||||||
|
public class DockerEverythingServerTests(ITestOutputHelper testOutputHelper) : LoggedTest(testOutputHelper)
|
||||||
|
{
|
||||||
|
/// <summary>Port number to be grabbed by the next test.</summary>
|
||||||
|
private static int s_nextPort = 3000;
|
||||||
|
|
||||||
|
// If the tests run concurrently against different versions of the runtime, tests can conflict with
|
||||||
|
// each other in the ports set up for interacting with containers. Ensure that such suites running
|
||||||
|
// against different TFMs use different port numbers.
|
||||||
|
private static readonly int s_portOffset = 1000 * (Environment.Version.Major switch
|
||||||
|
{
|
||||||
|
int v when v >= 8 => Environment.Version.Major - 7,
|
||||||
|
_ => 0,
|
||||||
|
});
|
||||||
|
|
||||||
|
private static int CreatePortNumber() => Interlocked.Increment(ref s_nextPort) + s_portOffset;
|
||||||
|
|
||||||
|
public static bool IsDockerAvailable => EverythingSseServerFixture.IsDockerAvailable;
|
||||||
|
|
||||||
|
[Fact(Skip = "docker is not available", SkipUnless = nameof(IsDockerAvailable))]
|
||||||
|
[Trait("Execution", "Manual")]
|
||||||
|
public async Task ConnectAndReceiveMessage_EverythingServerWithSse()
|
||||||
|
{
|
||||||
|
int port = CreatePortNumber();
|
||||||
|
|
||||||
|
await using var fixture = new EverythingSseServerFixture(port);
|
||||||
|
await fixture.StartAsync();
|
||||||
|
|
||||||
|
var defaultOptions = new McpClientOptions
|
||||||
|
{
|
||||||
|
ClientInfo = new() { Name = "IntegrationTestClient", Version = "1.0.0" }
|
||||||
|
};
|
||||||
|
|
||||||
|
var defaultConfig = new McpServerConfig
|
||||||
|
{
|
||||||
|
Id = "everything",
|
||||||
|
Name = "Everything",
|
||||||
|
TransportType = TransportTypes.Sse,
|
||||||
|
TransportOptions = [],
|
||||||
|
Location = $"http://localhost:{port}/sse"
|
||||||
|
};
|
||||||
|
|
||||||
|
// Create client and run tests
|
||||||
|
await using var client = await McpClientFactory.CreateAsync(
|
||||||
|
defaultConfig,
|
||||||
|
defaultOptions,
|
||||||
|
loggerFactory: LoggerFactory,
|
||||||
|
cancellationToken: TestContext.Current.CancellationToken);
|
||||||
|
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
|
||||||
|
|
||||||
|
// assert
|
||||||
|
Assert.NotEmpty(tools);
|
||||||
|
}
|
||||||
|
|
||||||
|
[Fact(Skip = "docker is not available", SkipUnless = nameof(IsDockerAvailable))]
|
||||||
|
[Trait("Execution", "Manual")]
|
||||||
|
public async Task Sampling_Sse_EverythingServer()
|
||||||
|
{
|
||||||
|
int port = CreatePortNumber();
|
||||||
|
|
||||||
|
await using var fixture = new EverythingSseServerFixture(port);
|
||||||
|
await fixture.StartAsync();
|
||||||
|
|
||||||
|
var defaultConfig = new McpServerConfig
|
||||||
|
{
|
||||||
|
Id = "everything",
|
||||||
|
Name = "Everything",
|
||||||
|
TransportType = TransportTypes.Sse,
|
||||||
|
TransportOptions = [],
|
||||||
|
Location = $"http://localhost:{port}/sse"
|
||||||
|
};
|
||||||
|
|
||||||
|
int samplingHandlerCalls = 0;
|
||||||
|
var defaultOptions = new McpClientOptions
|
||||||
|
{
|
||||||
|
Capabilities = new()
|
||||||
|
{
|
||||||
|
Sampling = new()
|
||||||
|
{
|
||||||
|
SamplingHandler = (_, _, _) =>
|
||||||
|
{
|
||||||
|
samplingHandlerCalls++;
|
||||||
|
return Task.FromResult(new CreateMessageResult
|
||||||
|
{
|
||||||
|
Model = "test-model",
|
||||||
|
Role = "assistant",
|
||||||
|
Content = new Content
|
||||||
|
{
|
||||||
|
Type = "text",
|
||||||
|
Text = "Test response"
|
||||||
|
}
|
||||||
|
});
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
await using var client = await McpClientFactory.CreateAsync(
|
||||||
|
defaultConfig,
|
||||||
|
defaultOptions,
|
||||||
|
loggerFactory: LoggerFactory,
|
||||||
|
cancellationToken: TestContext.Current.CancellationToken);
|
||||||
|
|
||||||
|
// Call the server's sampleLLM tool which should trigger our sampling handler
|
||||||
|
var result = await client.CallToolAsync("sampleLLM", new Dictionary<string, object?>
|
||||||
|
{
|
||||||
|
["prompt"] = "Test prompt",
|
||||||
|
["maxTokens"] = 100
|
||||||
|
}, cancellationToken: TestContext.Current.CancellationToken);
|
||||||
|
|
||||||
|
// assert
|
||||||
|
Assert.NotNull(result);
|
||||||
|
var textContent = Assert.Single(result.Content);
|
||||||
|
Assert.Equal("text", textContent.Type);
|
||||||
|
Assert.False(string.IsNullOrEmpty(textContent.Text));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -225,7 +225,7 @@ public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
|
||||||
var releaseHandler = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
|
var releaseHandler = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||||
|
|
||||||
IAsyncDisposable? registration = null;
|
IAsyncDisposable? registration = null;
|
||||||
registration = client.RegisterNotificationHandler(NotificationName, async (notification, cancellationToken) =>
|
await using var _ = registration = client.RegisterNotificationHandler(NotificationName, async (notification, cancellationToken) =>
|
||||||
{
|
{
|
||||||
for (int i = 0; i < numberOfDisposals; i++)
|
for (int i = 0; i < numberOfDisposals; i++)
|
||||||
{
|
{
|
||||||
|
|
@ -240,9 +240,5 @@ public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
|
||||||
|
|
||||||
await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
|
await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
|
||||||
await handlerRunning.Task;
|
await handlerRunning.Task;
|
||||||
|
|
||||||
ValueTask disposal = registration.DisposeAsync();
|
|
||||||
Assert.True(disposal.IsCompletedSuccessfully);
|
|
||||||
await disposal;
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
19
tests/ModelContextProtocol.Tests/Server/MapMcpTests.cs
Normal file
19
tests/ModelContextProtocol.Tests/Server/MapMcpTests.cs
Normal file
|
|
@ -0,0 +1,19 @@
|
||||||
|
using Microsoft.AspNetCore.Builder;
|
||||||
|
using ModelContextProtocol.Tests.Utils;
|
||||||
|
|
||||||
|
namespace ModelContextProtocol.Tests.Server;
|
||||||
|
|
||||||
|
public class MapMcpTests(ITestOutputHelper testOutputHelper) : KestrelInMemoryTest(testOutputHelper)
|
||||||
|
{
|
||||||
|
[Fact]
|
||||||
|
public async Task Allows_Customizing_Route()
|
||||||
|
{
|
||||||
|
await using var app = Builder.Build();
|
||||||
|
app.MapMcp("/mcp");
|
||||||
|
await app.StartAsync(TestContext.Current.CancellationToken);
|
||||||
|
|
||||||
|
using var httpClient = CreateHttpClient();
|
||||||
|
using var response = await httpClient.GetAsync("http://localhost/mcp/sse", HttpCompletionOption.ResponseHeadersRead, TestContext.Current.CancellationToken);
|
||||||
|
response.EnsureSuccessStatusCode();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -1,238 +1,162 @@
|
||||||
using Microsoft.Extensions.Logging;
|
using Microsoft.AspNetCore.Builder;
|
||||||
|
using Microsoft.AspNetCore.Http;
|
||||||
|
using Microsoft.AspNetCore.Http.Features;
|
||||||
|
using Microsoft.AspNetCore.Routing;
|
||||||
|
using Microsoft.Extensions.DependencyInjection;
|
||||||
|
using Microsoft.Extensions.Logging;
|
||||||
|
using Microsoft.Extensions.Options;
|
||||||
using ModelContextProtocol.Client;
|
using ModelContextProtocol.Client;
|
||||||
|
using ModelContextProtocol.Protocol.Messages;
|
||||||
using ModelContextProtocol.Protocol.Transport;
|
using ModelContextProtocol.Protocol.Transport;
|
||||||
using ModelContextProtocol.Protocol.Types;
|
using ModelContextProtocol.Server;
|
||||||
using ModelContextProtocol.Tests.Utils;
|
using ModelContextProtocol.Tests.Utils;
|
||||||
|
using ModelContextProtocol.Utils.Json;
|
||||||
|
|
||||||
namespace ModelContextProtocol.Tests;
|
namespace ModelContextProtocol.Tests;
|
||||||
|
|
||||||
public class SseIntegrationTests(ITestOutputHelper outputHelper) : LoggedTest(outputHelper)
|
public class SseIntegrationTests(ITestOutputHelper outputHelper) : KestrelInMemoryTest(outputHelper)
|
||||||
{
|
{
|
||||||
/// <summary>Port number to be grabbed by the next test.</summary>
|
private McpServerConfig DefaultServerConfig = new()
|
||||||
private static int s_nextPort = 3000;
|
|
||||||
|
|
||||||
// If the tests run concurrently against different versions of the runtime, tests can conflict with
|
|
||||||
// each other in the ports set up for interacting with containers. Ensure that such suites running
|
|
||||||
// against different TFMs use different port numbers.
|
|
||||||
private static readonly int s_portOffset = 1000 * (Environment.Version.Major switch
|
|
||||||
{
|
{
|
||||||
int v when v >= 8 => Environment.Version.Major - 7,
|
Id = "test_server",
|
||||||
_ => 0,
|
Name = "In-memory Test Server",
|
||||||
});
|
TransportType = TransportTypes.Sse,
|
||||||
|
TransportOptions = [],
|
||||||
|
Location = $"http://localhost/sse"
|
||||||
|
};
|
||||||
|
|
||||||
|
private Task<IMcpClient> ConnectMcpClient(HttpClient httpClient, McpClientOptions? clientOptions = null)
|
||||||
|
=> McpClientFactory.CreateAsync(
|
||||||
|
DefaultServerConfig,
|
||||||
|
clientOptions,
|
||||||
|
(_, _) => new SseClientTransport(new(), DefaultServerConfig, httpClient, LoggerFactory),
|
||||||
|
LoggerFactory,
|
||||||
|
TestContext.Current.CancellationToken);
|
||||||
|
|
||||||
private static int CreatePortNumber() => Interlocked.Increment(ref s_nextPort) + s_portOffset;
|
|
||||||
|
|
||||||
[Fact]
|
[Fact]
|
||||||
public async Task ConnectAndReceiveMessage_InMemoryServer()
|
public async Task ConnectAndReceiveMessage_InMemoryServer()
|
||||||
{
|
{
|
||||||
// Arrange
|
await using var app = Builder.Build();
|
||||||
await using InMemoryTestSseServer server = new(CreatePortNumber(), LoggerFactory.CreateLogger<InMemoryTestSseServer>());
|
app.MapMcp();
|
||||||
await server.StartAsync();
|
await app.StartAsync(TestContext.Current.CancellationToken);
|
||||||
|
|
||||||
var defaultConfig = new McpServerConfig
|
using var httpClient = CreateHttpClient();
|
||||||
{
|
await using var mcpClient = await ConnectMcpClient(httpClient);
|
||||||
Id = "test_server",
|
|
||||||
Name = "In-memory Test Server",
|
|
||||||
TransportType = TransportTypes.Sse,
|
|
||||||
TransportOptions = [],
|
|
||||||
Location = $"http://localhost:{server.Port}/sse"
|
|
||||||
};
|
|
||||||
|
|
||||||
// Act
|
|
||||||
await using var client = await McpClientFactory.CreateAsync(
|
|
||||||
defaultConfig,
|
|
||||||
loggerFactory: LoggerFactory,
|
|
||||||
cancellationToken: TestContext.Current.CancellationToken);
|
|
||||||
|
|
||||||
// Wait for SSE connection to be established
|
|
||||||
await server.WaitForConnectionAsync(TimeSpan.FromSeconds(10));
|
|
||||||
|
|
||||||
// Send a test message through POST endpoint
|
// Send a test message through POST endpoint
|
||||||
await client.SendNotificationAsync("test/message", new { message = "Hello, SSE!" }, cancellationToken: TestContext.Current.CancellationToken);
|
await mcpClient.SendNotificationAsync("test/message", new { message = "Hello, SSE!" }, cancellationToken: TestContext.Current.CancellationToken);
|
||||||
|
|
||||||
// Assert
|
|
||||||
Assert.True(true);
|
Assert.True(true);
|
||||||
}
|
}
|
||||||
|
|
||||||
[Fact]
|
|
||||||
[Trait("Execution", "Manual")]
|
|
||||||
public async Task ConnectAndReceiveMessage_EverythingServerWithSse()
|
|
||||||
{
|
|
||||||
Assert.SkipWhen(!EverythingSseServerFixture.IsDockerAvailable, "docker is not available");
|
|
||||||
|
|
||||||
int port = CreatePortNumber();
|
|
||||||
|
|
||||||
await using var fixture = new EverythingSseServerFixture(port);
|
|
||||||
await fixture.StartAsync();
|
|
||||||
|
|
||||||
var defaultOptions = new McpClientOptions
|
|
||||||
{
|
|
||||||
ClientInfo = new() { Name = "IntegrationTestClient", Version = "1.0.0" }
|
|
||||||
};
|
|
||||||
|
|
||||||
var defaultConfig = new McpServerConfig
|
|
||||||
{
|
|
||||||
Id = "everything",
|
|
||||||
Name = "Everything",
|
|
||||||
TransportType = TransportTypes.Sse,
|
|
||||||
TransportOptions = [],
|
|
||||||
Location = $"http://localhost:{port}/sse"
|
|
||||||
};
|
|
||||||
|
|
||||||
// Create client and run tests
|
|
||||||
await using var client = await McpClientFactory.CreateAsync(
|
|
||||||
defaultConfig,
|
|
||||||
defaultOptions,
|
|
||||||
loggerFactory: LoggerFactory,
|
|
||||||
cancellationToken: TestContext.Current.CancellationToken);
|
|
||||||
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
|
|
||||||
|
|
||||||
// assert
|
|
||||||
Assert.NotEmpty(tools);
|
|
||||||
}
|
|
||||||
|
|
||||||
[Fact]
|
|
||||||
[Trait("Execution", "Manual")]
|
|
||||||
public async Task Sampling_Sse_EverythingServer()
|
|
||||||
{
|
|
||||||
Assert.SkipWhen(!EverythingSseServerFixture.IsDockerAvailable, "docker is not available");
|
|
||||||
|
|
||||||
int port = CreatePortNumber();
|
|
||||||
|
|
||||||
await using var fixture = new EverythingSseServerFixture(port);
|
|
||||||
await fixture.StartAsync();
|
|
||||||
|
|
||||||
var defaultConfig = new McpServerConfig
|
|
||||||
{
|
|
||||||
Id = "everything",
|
|
||||||
Name = "Everything",
|
|
||||||
TransportType = TransportTypes.Sse,
|
|
||||||
TransportOptions = [],
|
|
||||||
Location = $"http://localhost:{port}/sse"
|
|
||||||
};
|
|
||||||
|
|
||||||
int samplingHandlerCalls = 0;
|
|
||||||
var defaultOptions = new McpClientOptions
|
|
||||||
{
|
|
||||||
Capabilities = new()
|
|
||||||
{
|
|
||||||
Sampling = new()
|
|
||||||
{
|
|
||||||
SamplingHandler = (_, _, _) =>
|
|
||||||
{
|
|
||||||
samplingHandlerCalls++;
|
|
||||||
return Task.FromResult(new CreateMessageResult
|
|
||||||
{
|
|
||||||
Model = "test-model",
|
|
||||||
Role = "assistant",
|
|
||||||
Content = new Content
|
|
||||||
{
|
|
||||||
Type = "text",
|
|
||||||
Text = "Test response"
|
|
||||||
}
|
|
||||||
});
|
|
||||||
},
|
|
||||||
},
|
|
||||||
},
|
|
||||||
};
|
|
||||||
|
|
||||||
await using var client = await McpClientFactory.CreateAsync(
|
|
||||||
defaultConfig,
|
|
||||||
defaultOptions,
|
|
||||||
loggerFactory: LoggerFactory,
|
|
||||||
cancellationToken: TestContext.Current.CancellationToken);
|
|
||||||
|
|
||||||
// Call the server's sampleLLM tool which should trigger our sampling handler
|
|
||||||
var result = await client.CallToolAsync("sampleLLM", new Dictionary<string, object?>
|
|
||||||
{
|
|
||||||
["prompt"] = "Test prompt",
|
|
||||||
["maxTokens"] = 100
|
|
||||||
}, cancellationToken: TestContext.Current.CancellationToken);
|
|
||||||
|
|
||||||
// assert
|
|
||||||
Assert.NotNull(result);
|
|
||||||
var textContent = Assert.Single(result.Content);
|
|
||||||
Assert.Equal("text", textContent.Type);
|
|
||||||
Assert.False(string.IsNullOrEmpty(textContent.Text));
|
|
||||||
}
|
|
||||||
|
|
||||||
[Fact]
|
[Fact]
|
||||||
public async Task ConnectAndReceiveMessage_InMemoryServer_WithFullEndpointEventUri()
|
public async Task ConnectAndReceiveMessage_InMemoryServer_WithFullEndpointEventUri()
|
||||||
{
|
{
|
||||||
// Arrange
|
await using var app = Builder.Build();
|
||||||
await using InMemoryTestSseServer server = new(CreatePortNumber(), LoggerFactory.CreateLogger<InMemoryTestSseServer>());
|
MapAbsoluteEndpointUriMcp(app);
|
||||||
server.UseFullUrlForEndpointEvent = true;
|
await app.StartAsync(TestContext.Current.CancellationToken);
|
||||||
await server.StartAsync();
|
|
||||||
|
|
||||||
var defaultConfig = new McpServerConfig
|
using var httpClient = CreateHttpClient();
|
||||||
{
|
await using var mcpClient = await ConnectMcpClient(httpClient);
|
||||||
Id = "test_server",
|
|
||||||
Name = "In-memory Test Server",
|
|
||||||
TransportType = TransportTypes.Sse,
|
|
||||||
TransportOptions = [],
|
|
||||||
Location = $"http://localhost:{server.Port}/sse"
|
|
||||||
};
|
|
||||||
|
|
||||||
// Act
|
|
||||||
await using var client = await McpClientFactory.CreateAsync(
|
|
||||||
defaultConfig,
|
|
||||||
loggerFactory: LoggerFactory,
|
|
||||||
cancellationToken: TestContext.Current.CancellationToken);
|
|
||||||
|
|
||||||
// Wait for SSE connection to be established
|
|
||||||
await server.WaitForConnectionAsync(TimeSpan.FromSeconds(10));
|
|
||||||
|
|
||||||
// Send a test message through POST endpoint
|
// Send a test message through POST endpoint
|
||||||
await client.SendNotificationAsync("test/message", new { message = "Hello, SSE!" }, cancellationToken: TestContext.Current.CancellationToken);
|
await mcpClient.SendNotificationAsync("test/message", new { message = "Hello, SSE!" }, cancellationToken: TestContext.Current.CancellationToken);
|
||||||
|
|
||||||
// Assert
|
|
||||||
Assert.True(true);
|
Assert.True(true);
|
||||||
}
|
}
|
||||||
|
|
||||||
[Fact]
|
[Fact]
|
||||||
public async Task ConnectAndReceiveNotification_InMemoryServer()
|
public async Task ConnectAndReceiveNotification_InMemoryServer()
|
||||||
{
|
{
|
||||||
// Arrange
|
|
||||||
await using InMemoryTestSseServer server = new(CreatePortNumber(), LoggerFactory.CreateLogger<InMemoryTestSseServer>());
|
|
||||||
await server.StartAsync();
|
|
||||||
|
|
||||||
var defaultConfig = new McpServerConfig
|
|
||||||
{
|
|
||||||
Id = "test_server",
|
|
||||||
Name = "In-memory Test Server",
|
|
||||||
TransportType = TransportTypes.Sse,
|
|
||||||
TransportOptions = [],
|
|
||||||
Location = $"http://localhost:{server.Port}/sse"
|
|
||||||
};
|
|
||||||
|
|
||||||
// Act
|
|
||||||
var receivedNotification = new TaskCompletionSource<string?>();
|
var receivedNotification = new TaskCompletionSource<string?>();
|
||||||
await using var client = await McpClientFactory.CreateAsync(
|
|
||||||
defaultConfig,
|
await using var app = Builder.Build();
|
||||||
new()
|
app.MapMcp(runSessionAsync: (httpContext, mcpServer, cancellationToken) =>
|
||||||
|
{
|
||||||
|
mcpServer.RegisterNotificationHandler("test/notification", async (notification, cancellationToken) =>
|
||||||
{
|
{
|
||||||
Capabilities = new()
|
Assert.Equal("Hello from client!", notification.Params?["message"]?.GetValue<string>());
|
||||||
{
|
await mcpServer.SendNotificationAsync("test/notification", new { message = "Hello from server!" }, cancellationToken: cancellationToken);
|
||||||
NotificationHandlers = [new("test/notification", (notification, cancellationToken) =>
|
});
|
||||||
{
|
return mcpServer.RunAsync(cancellationToken);
|
||||||
var msg = notification.Params?["message"]?.GetValue<string>();
|
});
|
||||||
receivedNotification.SetResult(msg);
|
await app.StartAsync(TestContext.Current.CancellationToken);
|
||||||
|
|
||||||
return Task.CompletedTask;
|
using var httpClient = CreateHttpClient();
|
||||||
})],
|
await using var mcpClient = await ConnectMcpClient(httpClient);
|
||||||
},
|
|
||||||
},
|
|
||||||
loggerFactory: LoggerFactory,
|
|
||||||
cancellationToken: TestContext.Current.CancellationToken);
|
|
||||||
|
|
||||||
// Wait for SSE connection to be established
|
mcpClient.RegisterNotificationHandler("test/notification", (args, ca) =>
|
||||||
await server.WaitForConnectionAsync(TimeSpan.FromSeconds(10));
|
{
|
||||||
|
var msg = args.Params?["message"]?.GetValue<string>();
|
||||||
|
receivedNotification.SetResult(msg);
|
||||||
|
return Task.CompletedTask;
|
||||||
|
});
|
||||||
|
|
||||||
// Act
|
// Send a test message through POST endpoint
|
||||||
await server.SendTestNotificationAsync("Hello from server!");
|
await mcpClient.SendNotificationAsync("test/notification", new { message = "Hello from client!" }, cancellationToken: TestContext.Current.CancellationToken);
|
||||||
|
|
||||||
// Assert
|
|
||||||
var message = await receivedNotification.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken);
|
var message = await receivedNotification.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken);
|
||||||
Assert.Equal("Hello from server!", message);
|
Assert.Equal("Hello from server!", message);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private static void MapAbsoluteEndpointUriMcp(IEndpointRouteBuilder endpoints)
|
||||||
|
{
|
||||||
|
var loggerFactory = endpoints.ServiceProvider.GetRequiredService<ILoggerFactory>();
|
||||||
|
var optionsSnapshot = endpoints.ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>();
|
||||||
|
|
||||||
|
var routeGroup = endpoints.MapGroup("");
|
||||||
|
SseResponseStreamTransport? session = null;
|
||||||
|
|
||||||
|
routeGroup.MapGet("/sse", async context =>
|
||||||
|
{
|
||||||
|
var response = context.Response;
|
||||||
|
var requestAborted = context.RequestAborted;
|
||||||
|
|
||||||
|
response.Headers.ContentType = "text/event-stream";
|
||||||
|
|
||||||
|
await using var transport = new SseResponseStreamTransport(response.Body, "http://localhost/message");
|
||||||
|
session = transport;
|
||||||
|
|
||||||
|
try
|
||||||
|
{
|
||||||
|
var transportTask = transport.RunAsync(cancellationToken: requestAborted);
|
||||||
|
await using var server = McpServerFactory.Create(transport, optionsSnapshot.Value, loggerFactory, endpoints.ServiceProvider);
|
||||||
|
|
||||||
|
try
|
||||||
|
{
|
||||||
|
await server.RunAsync(requestAborted);
|
||||||
|
}
|
||||||
|
finally
|
||||||
|
{
|
||||||
|
await transport.DisposeAsync();
|
||||||
|
await transportTask;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
catch (OperationCanceledException) when (requestAborted.IsCancellationRequested)
|
||||||
|
{
|
||||||
|
// RequestAborted always triggers when the client disconnects before a complete response body is written,
|
||||||
|
// but this is how SSE connections are typically closed.
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
routeGroup.MapPost("/message", async context =>
|
||||||
|
{
|
||||||
|
if (session is null)
|
||||||
|
{
|
||||||
|
await Results.BadRequest("Session not started.").ExecuteAsync(context);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
var message = (IJsonRpcMessage?)await context.Request.ReadFromJsonAsync(McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(IJsonRpcMessage)), context.RequestAborted);
|
||||||
|
if (message is null)
|
||||||
|
{
|
||||||
|
await Results.BadRequest("No message in request body.").ExecuteAsync(context);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
await session.OnMessageReceivedAsync(message, context.RequestAborted);
|
||||||
|
context.Response.StatusCode = StatusCodes.Status202Accepted;
|
||||||
|
await context.Response.WriteAsync("Accepted");
|
||||||
|
});
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,6 @@
|
||||||
using ModelContextProtocol.Protocol.Transport;
|
using Microsoft.Extensions.Logging;
|
||||||
|
using ModelContextProtocol.Client;
|
||||||
|
using ModelContextProtocol.Protocol.Transport;
|
||||||
using ModelContextProtocol.Test.Utils;
|
using ModelContextProtocol.Test.Utils;
|
||||||
using ModelContextProtocol.Tests.Utils;
|
using ModelContextProtocol.Tests.Utils;
|
||||||
using ModelContextProtocol.TestSseServer;
|
using ModelContextProtocol.TestSseServer;
|
||||||
|
|
@ -7,28 +9,52 @@ namespace ModelContextProtocol.Tests;
|
||||||
|
|
||||||
public class SseServerIntegrationTestFixture : IAsyncDisposable
|
public class SseServerIntegrationTestFixture : IAsyncDisposable
|
||||||
{
|
{
|
||||||
|
private readonly KestrelInMemoryTransport _inMemoryTransport = new();
|
||||||
|
|
||||||
private readonly Task _serverTask;
|
private readonly Task _serverTask;
|
||||||
private readonly CancellationTokenSource _stopCts = new();
|
private readonly CancellationTokenSource _stopCts = new();
|
||||||
|
|
||||||
|
// XUnit's ITestOutputHelper is created per test, while this fixture is used for
|
||||||
|
// multiple tests, so this dispatches the output to the current test.
|
||||||
private readonly DelegatingTestOutputHelper _delegatingTestOutputHelper = new();
|
private readonly DelegatingTestOutputHelper _delegatingTestOutputHelper = new();
|
||||||
|
|
||||||
public McpServerConfig DefaultConfig { get; }
|
private McpServerConfig DefaultServerConfig { get; } = new McpServerConfig
|
||||||
|
{
|
||||||
|
Id = "test_server",
|
||||||
|
Name = "TestServer",
|
||||||
|
TransportType = TransportTypes.Sse,
|
||||||
|
TransportOptions = [],
|
||||||
|
Location = $"http://localhost/sse"
|
||||||
|
};
|
||||||
|
|
||||||
public SseServerIntegrationTestFixture()
|
public SseServerIntegrationTestFixture()
|
||||||
{
|
{
|
||||||
// Ensure that test suites running against different TFMs and possibly concurrently use different port numbers.
|
var socketsHttpHandler = new SocketsHttpHandler()
|
||||||
int port = 3001 + Environment.Version.Major;
|
|
||||||
|
|
||||||
DefaultConfig = new McpServerConfig
|
|
||||||
{
|
{
|
||||||
Id = "test_server",
|
ConnectCallback = (context, token) =>
|
||||||
Name = "TestServer",
|
{
|
||||||
TransportType = TransportTypes.Sse,
|
var connection = _inMemoryTransport.CreateConnection();
|
||||||
TransportOptions = [],
|
return new(connection.ClientStream);
|
||||||
Location = $"http://localhost:{port}/sse"
|
},
|
||||||
};
|
};
|
||||||
|
|
||||||
_serverTask = Program.MainAsync([port.ToString()], new XunitLoggerProvider(_delegatingTestOutputHelper), _stopCts.Token);
|
HttpClient = new HttpClient(socketsHttpHandler)
|
||||||
|
{
|
||||||
|
BaseAddress = new Uri(DefaultServerConfig.Location),
|
||||||
|
};
|
||||||
|
_serverTask = Program.MainAsync([], new XunitLoggerProvider(_delegatingTestOutputHelper), _inMemoryTransport, _stopCts.Token);
|
||||||
|
}
|
||||||
|
|
||||||
|
public HttpClient HttpClient { get; }
|
||||||
|
|
||||||
|
public Task<IMcpClient> ConnectMcpClientAsync(McpClientOptions? options, ILoggerFactory loggerFactory)
|
||||||
|
{
|
||||||
|
return McpClientFactory.CreateAsync(
|
||||||
|
DefaultServerConfig,
|
||||||
|
options,
|
||||||
|
(_, _) => new SseClientTransport(new(), DefaultServerConfig, HttpClient, loggerFactory),
|
||||||
|
loggerFactory,
|
||||||
|
TestContext.Current.CancellationToken);
|
||||||
}
|
}
|
||||||
|
|
||||||
public void Initialize(ITestOutputHelper output)
|
public void Initialize(ITestOutputHelper output)
|
||||||
|
|
@ -44,7 +70,10 @@ public class SseServerIntegrationTestFixture : IAsyncDisposable
|
||||||
public async ValueTask DisposeAsync()
|
public async ValueTask DisposeAsync()
|
||||||
{
|
{
|
||||||
_delegatingTestOutputHelper.CurrentTestOutputHelper = null;
|
_delegatingTestOutputHelper.CurrentTestOutputHelper = null;
|
||||||
|
|
||||||
|
HttpClient.Dispose();
|
||||||
_stopCts.Cancel();
|
_stopCts.Cancel();
|
||||||
|
|
||||||
try
|
try
|
||||||
{
|
{
|
||||||
await _serverTask.ConfigureAwait(false);
|
await _serverTask.ConfigureAwait(false);
|
||||||
|
|
@ -52,6 +81,7 @@ public class SseServerIntegrationTestFixture : IAsyncDisposable
|
||||||
catch (OperationCanceledException)
|
catch (OperationCanceledException)
|
||||||
{
|
{
|
||||||
}
|
}
|
||||||
|
|
||||||
_stopCts.Dispose();
|
_stopCts.Dispose();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -25,19 +25,9 @@ public class SseServerIntegrationTests : LoggedTest, IClassFixture<SseServerInte
|
||||||
|
|
||||||
private Task<IMcpClient> GetClientAsync(McpClientOptions? options = null)
|
private Task<IMcpClient> GetClientAsync(McpClientOptions? options = null)
|
||||||
{
|
{
|
||||||
return McpClientFactory.CreateAsync(
|
return _fixture.ConnectMcpClientAsync(options, LoggerFactory);
|
||||||
_fixture.DefaultConfig,
|
|
||||||
options,
|
|
||||||
loggerFactory: LoggerFactory,
|
|
||||||
cancellationToken: TestContext.Current.CancellationToken);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
private HttpClient GetHttpClient() =>
|
|
||||||
new()
|
|
||||||
{
|
|
||||||
BaseAddress = new(_fixture.DefaultConfig.Location!),
|
|
||||||
};
|
|
||||||
|
|
||||||
[Fact]
|
[Fact]
|
||||||
public async Task ConnectAndPing_Sse_TestServer()
|
public async Task ConnectAndPing_Sse_TestServer()
|
||||||
{
|
{
|
||||||
|
|
@ -283,8 +273,7 @@ public class SseServerIntegrationTests : LoggedTest, IClassFixture<SseServerInte
|
||||||
[Fact]
|
[Fact]
|
||||||
public async Task EventSourceResponse_Includes_ExpectedHeaders()
|
public async Task EventSourceResponse_Includes_ExpectedHeaders()
|
||||||
{
|
{
|
||||||
using var httpClient = GetHttpClient();
|
using var sseResponse = await _fixture.HttpClient.GetAsync("", HttpCompletionOption.ResponseHeadersRead, TestContext.Current.CancellationToken);
|
||||||
using var sseResponse = await httpClient.GetAsync("", HttpCompletionOption.ResponseHeadersRead, TestContext.Current.CancellationToken);
|
|
||||||
|
|
||||||
sseResponse.EnsureSuccessStatusCode();
|
sseResponse.EnsureSuccessStatusCode();
|
||||||
|
|
||||||
|
|
@ -300,8 +289,7 @@ public class SseServerIntegrationTests : LoggedTest, IClassFixture<SseServerInte
|
||||||
{
|
{
|
||||||
// Simulate our own MCP client handshake using a plain HttpClient so we can look for "event: message"
|
// Simulate our own MCP client handshake using a plain HttpClient so we can look for "event: message"
|
||||||
// in the raw SSE response stream which is not exposed by the real MCP client.
|
// in the raw SSE response stream which is not exposed by the real MCP client.
|
||||||
using var httpClient = GetHttpClient();
|
await using var sseResponse = await _fixture.HttpClient.GetStreamAsync("", TestContext.Current.CancellationToken);
|
||||||
await using var sseResponse = await httpClient.GetStreamAsync("", TestContext.Current.CancellationToken);
|
|
||||||
using var streamReader = new StreamReader(sseResponse);
|
using var streamReader = new StreamReader(sseResponse);
|
||||||
|
|
||||||
var endpointEvent = await streamReader.ReadLineAsync(TestContext.Current.CancellationToken);
|
var endpointEvent = await streamReader.ReadLineAsync(TestContext.Current.CancellationToken);
|
||||||
|
|
@ -317,7 +305,7 @@ public class SseServerIntegrationTests : LoggedTest, IClassFixture<SseServerInte
|
||||||
""";
|
""";
|
||||||
using (var initializeRequestBody = new StringContent(initializeRequest, Encoding.UTF8, "application/json"))
|
using (var initializeRequestBody = new StringContent(initializeRequest, Encoding.UTF8, "application/json"))
|
||||||
{
|
{
|
||||||
var response = await httpClient.PostAsync(messageEndpoint, initializeRequestBody, TestContext.Current.CancellationToken);
|
var response = await _fixture.HttpClient.PostAsync(messageEndpoint, initializeRequestBody, TestContext.Current.CancellationToken);
|
||||||
Assert.Equal(HttpStatusCode.Accepted, response.StatusCode);
|
Assert.Equal(HttpStatusCode.Accepted, response.StatusCode);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -326,7 +314,7 @@ public class SseServerIntegrationTests : LoggedTest, IClassFixture<SseServerInte
|
||||||
""";
|
""";
|
||||||
using (var initializedNotificationBody = new StringContent(initializedNotification, Encoding.UTF8, "application/json"))
|
using (var initializedNotificationBody = new StringContent(initializedNotification, Encoding.UTF8, "application/json"))
|
||||||
{
|
{
|
||||||
var response = await httpClient.PostAsync(messageEndpoint, initializedNotificationBody, TestContext.Current.CancellationToken);
|
var response = await _fixture.HttpClient.PostAsync(messageEndpoint, initializedNotificationBody, TestContext.Current.CancellationToken);
|
||||||
Assert.Equal(HttpStatusCode.Accepted, response.StatusCode);
|
Assert.Equal(HttpStatusCode.Accepted, response.StatusCode);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,27 @@
|
||||||
|
using ModelContextProtocol.Protocol.Transport;
|
||||||
|
using ModelContextProtocol.Tests.Utils;
|
||||||
|
using System.IO.Pipelines;
|
||||||
|
|
||||||
|
namespace ModelContextProtocol.Tests.Transport;
|
||||||
|
|
||||||
|
public class SseResponseStreamTransportTests(ITestOutputHelper testOutputHelper) : LoggedTest(testOutputHelper)
|
||||||
|
{
|
||||||
|
[Fact]
|
||||||
|
public async Task Can_Customize_MessageEndpoint()
|
||||||
|
{
|
||||||
|
var responsePipe = new Pipe();
|
||||||
|
|
||||||
|
await using var transport = new SseResponseStreamTransport(responsePipe.Writer.AsStream(), "/my-message-endpoint");
|
||||||
|
var transportRunTask = transport.RunAsync(TestContext.Current.CancellationToken);
|
||||||
|
|
||||||
|
using var responseStreamReader = new StreamReader(responsePipe.Reader.AsStream());
|
||||||
|
var firstLine = await responseStreamReader.ReadLineAsync(TestContext.Current.CancellationToken);
|
||||||
|
Assert.Equal("event: endpoint", firstLine);
|
||||||
|
|
||||||
|
var secondLine = await responseStreamReader.ReadLineAsync(TestContext.Current.CancellationToken);
|
||||||
|
Assert.Equal("data: /my-message-endpoint", secondLine);
|
||||||
|
|
||||||
|
responsePipe.Reader.Complete();
|
||||||
|
responsePipe.Writer.Complete();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -1,382 +0,0 @@
|
||||||
using System.Collections.Concurrent;
|
|
||||||
using System.Net;
|
|
||||||
using System.Text.Json;
|
|
||||||
using ModelContextProtocol.Protocol.Messages;
|
|
||||||
using Microsoft.Extensions.Logging;
|
|
||||||
using Microsoft.Extensions.Logging.Abstractions;
|
|
||||||
using ModelContextProtocol.Protocol.Types;
|
|
||||||
|
|
||||||
namespace ModelContextProtocol.Tests.Utils;
|
|
||||||
|
|
||||||
public sealed class InMemoryTestSseServer : IAsyncDisposable
|
|
||||||
{
|
|
||||||
private readonly HttpListener _listener;
|
|
||||||
private readonly CancellationTokenSource _cts;
|
|
||||||
private readonly ILogger<InMemoryTestSseServer> _logger;
|
|
||||||
private Task? _serverTask;
|
|
||||||
private readonly TaskCompletionSource _connectionEstablished = new();
|
|
||||||
|
|
||||||
// SSE endpoint for GET
|
|
||||||
private readonly string _endpointPath;
|
|
||||||
// POST endpoint
|
|
||||||
private readonly string _messagePath;
|
|
||||||
|
|
||||||
// Keep track of all open SSE connections (StreamWriters).
|
|
||||||
private readonly ConcurrentBag<StreamWriter> _sseClients = [];
|
|
||||||
|
|
||||||
public InMemoryTestSseServer(int port = 5000, ILogger<InMemoryTestSseServer>? logger = null)
|
|
||||||
{
|
|
||||||
Port = port;
|
|
||||||
|
|
||||||
_listener = new HttpListener();
|
|
||||||
_listener.Prefixes.Add($"http://localhost:{port}/");
|
|
||||||
_cts = new CancellationTokenSource();
|
|
||||||
_logger = logger ?? NullLogger<InMemoryTestSseServer>.Instance;
|
|
||||||
|
|
||||||
_endpointPath = "/sse";
|
|
||||||
_messagePath = "/message";
|
|
||||||
}
|
|
||||||
|
|
||||||
public int Port { get; }
|
|
||||||
|
|
||||||
/// <summary>
|
|
||||||
/// This is to be able to use the full URL for the endpoint event.
|
|
||||||
/// </summary>
|
|
||||||
public bool UseFullUrlForEndpointEvent { get; set; }
|
|
||||||
|
|
||||||
/// <summary>
|
|
||||||
/// Full URL for the SSE endpoint, e.g. "http://localhost:5000/sse"
|
|
||||||
/// </summary>
|
|
||||||
public string SseEndpoint
|
|
||||||
=> $"http://localhost:{_listener.Prefixes.First().Split(':')[2].TrimEnd('/')}{_endpointPath}";
|
|
||||||
|
|
||||||
/// <summary>
|
|
||||||
/// Full URL for the message endpoint, e.g. "http://localhost:5000/message"
|
|
||||||
/// </summary>
|
|
||||||
public string MessageEndpoint
|
|
||||||
=> $"http://localhost:{_listener.Prefixes.First().Split(':')[2].TrimEnd('/')}{_messagePath}";
|
|
||||||
|
|
||||||
/// <summary>
|
|
||||||
/// Starts the server so it can accept incoming connections and POST requests.
|
|
||||||
/// </summary>
|
|
||||||
public async Task StartAsync()
|
|
||||||
{
|
|
||||||
_listener.Start();
|
|
||||||
_serverTask = HandleConnectionsAsync(_cts.Token);
|
|
||||||
|
|
||||||
_logger.LogInformation("Test SSE server started on {Endpoint}", SseEndpoint);
|
|
||||||
await Task.CompletedTask;
|
|
||||||
}
|
|
||||||
|
|
||||||
private async Task HandleConnectionsAsync(CancellationToken ct)
|
|
||||||
{
|
|
||||||
try
|
|
||||||
{
|
|
||||||
while (!ct.IsCancellationRequested)
|
|
||||||
{
|
|
||||||
var context = await _listener.GetContextAsync();
|
|
||||||
_ = Task.Run(() => HandleRequestAsync(context, ct), ct);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
catch (OperationCanceledException)
|
|
||||||
{
|
|
||||||
// Ignore, we are shutting down
|
|
||||||
}
|
|
||||||
catch (Exception ex) when (!ct.IsCancellationRequested)
|
|
||||||
{
|
|
||||||
_logger.LogError(ex, "Error in SSE server connection handling");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
private async Task HandleRequestAsync(HttpListenerContext context, CancellationToken ct)
|
|
||||||
{
|
|
||||||
var request = context.Request;
|
|
||||||
var response = context.Response;
|
|
||||||
|
|
||||||
// Handle SSE endpoint
|
|
||||||
if (request.HttpMethod.Equals("GET", StringComparison.OrdinalIgnoreCase)
|
|
||||||
&& request.Url?.AbsolutePath.Equals(_endpointPath, StringComparison.OrdinalIgnoreCase) == true)
|
|
||||||
{
|
|
||||||
await HandleSseConnectionAsync(context, ct);
|
|
||||||
}
|
|
||||||
// Handle POST /message
|
|
||||||
else if (request.HttpMethod.Equals("POST", StringComparison.OrdinalIgnoreCase)
|
|
||||||
&& request.Url?.AbsolutePath.Equals(_messagePath, StringComparison.OrdinalIgnoreCase) == true)
|
|
||||||
{
|
|
||||||
await HandlePostMessageAsync(context, ct);
|
|
||||||
}
|
|
||||||
else
|
|
||||||
{
|
|
||||||
response.StatusCode = 404;
|
|
||||||
response.Close();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// <summary>
|
|
||||||
/// Handle Server-Sent Events (SSE) connection.
|
|
||||||
/// Send the initial event: endpoint with the full POST URL.
|
|
||||||
/// Keep the connection open until the server is disposed or the client disconnects.
|
|
||||||
/// </summary>
|
|
||||||
private async Task HandleSseConnectionAsync(HttpListenerContext context, CancellationToken ct)
|
|
||||||
{
|
|
||||||
var response = context.Response;
|
|
||||||
response.ContentType = "text/event-stream";
|
|
||||||
response.Headers["Cache-Control"] = "no-cache";
|
|
||||||
// Ensures the response is never chunked away by the framework
|
|
||||||
response.SendChunked = true;
|
|
||||||
response.StatusCode = (int)HttpStatusCode.OK;
|
|
||||||
|
|
||||||
using var writer = new StreamWriter(response.OutputStream);
|
|
||||||
_sseClients.Add(writer);
|
|
||||||
|
|
||||||
// Immediately send the "endpoint" event with the POST URL
|
|
||||||
await writer.WriteLineAsync("event: endpoint");
|
|
||||||
if (UseFullUrlForEndpointEvent)
|
|
||||||
{
|
|
||||||
await writer.WriteLineAsync($"data: {MessageEndpoint}");
|
|
||||||
}
|
|
||||||
else
|
|
||||||
{
|
|
||||||
await writer.WriteLineAsync($"data: {_messagePath}");
|
|
||||||
}
|
|
||||||
await writer.WriteLineAsync(); // blank line to end an SSE message
|
|
||||||
await writer.FlushAsync(ct);
|
|
||||||
|
|
||||||
_logger.LogInformation("New SSE client connected.");
|
|
||||||
_connectionEstablished.TrySetResult(); // Signal connection is ready
|
|
||||||
|
|
||||||
try
|
|
||||||
{
|
|
||||||
// Keep the connection open by "pinging" or just waiting
|
|
||||||
// until the client disconnects or the server is canceled.
|
|
||||||
while (!ct.IsCancellationRequested && response.OutputStream.CanWrite)
|
|
||||||
{
|
|
||||||
_logger.LogDebug("SSE connection alive check");
|
|
||||||
// Optionally do a periodic no-op to keep connection alive:
|
|
||||||
await writer.WriteLineAsync(": keep-alive");
|
|
||||||
await writer.FlushAsync(ct);
|
|
||||||
await Task.Delay(TimeSpan.FromSeconds(5), ct);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
catch (TaskCanceledException)
|
|
||||||
{
|
|
||||||
// This is expected on shutdown
|
|
||||||
_logger.LogInformation("SSE connection cancelled (expected on shutdown)");
|
|
||||||
}
|
|
||||||
catch (IOException)
|
|
||||||
{
|
|
||||||
// Client likely disconnected
|
|
||||||
_logger.LogInformation("SSE client disconnected");
|
|
||||||
}
|
|
||||||
finally
|
|
||||||
{
|
|
||||||
// Remove this writer from bag (we're disposing it anyway)
|
|
||||||
_sseClients.TryTake(out _);
|
|
||||||
_logger.LogInformation("SSE client disconnected.");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Add method to wait for connection
|
|
||||||
public Task WaitForConnectionAsync(TimeSpan timeout) =>
|
|
||||||
_connectionEstablished.Task.WaitAsync(timeout);
|
|
||||||
|
|
||||||
/// <summary>
|
|
||||||
/// Handle POST /message endpoint.
|
|
||||||
/// Echo the content back to the caller and broadcast it over SSE as well.
|
|
||||||
/// </summary>
|
|
||||||
private async Task HandlePostMessageAsync(HttpListenerContext context, CancellationToken cancellationToken)
|
|
||||||
{
|
|
||||||
var request = context.Request;
|
|
||||||
var response = context.Response;
|
|
||||||
|
|
||||||
try
|
|
||||||
{
|
|
||||||
using var reader = new StreamReader(request.InputStream);
|
|
||||||
string content = await reader.ReadToEndAsync(cancellationToken);
|
|
||||||
|
|
||||||
var jsonRpcNotification = JsonSerializer.Deserialize<JsonRpcNotification>(content);
|
|
||||||
if (jsonRpcNotification != null && jsonRpcNotification.Method != RequestMethods.Initialize)
|
|
||||||
{
|
|
||||||
// Test server so just ignore notifications
|
|
||||||
|
|
||||||
// Set status code to 202 Accepted
|
|
||||||
response.StatusCode = 202;
|
|
||||||
|
|
||||||
// Write "accepted" as the response body
|
|
||||||
using var writer = new StreamWriter(response.OutputStream);
|
|
||||||
await writer.WriteAsync("accepted");
|
|
||||||
await writer.FlushAsync(cancellationToken);
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
var jsonRpcRequest = JsonSerializer.Deserialize<JsonRpcRequest>(content);
|
|
||||||
|
|
||||||
if (jsonRpcRequest != null)
|
|
||||||
{
|
|
||||||
if (jsonRpcRequest.Method == RequestMethods.Initialize)
|
|
||||||
{
|
|
||||||
await HandleInitializationRequest(response, jsonRpcRequest);
|
|
||||||
}
|
|
||||||
else
|
|
||||||
{
|
|
||||||
// Set status code to 202 Accepted
|
|
||||||
response.StatusCode = 202;
|
|
||||||
|
|
||||||
// Write "accepted" as the response body
|
|
||||||
using var writer = new StreamWriter(response.OutputStream);
|
|
||||||
await writer.WriteAsync("accepted");
|
|
||||||
await writer.FlushAsync(cancellationToken);
|
|
||||||
|
|
||||||
// Then process the message and send the response via SSE
|
|
||||||
if (jsonRpcRequest != null)
|
|
||||||
{
|
|
||||||
// Process the request and generate a response
|
|
||||||
var responseMessage = await ProcessRequestAsync(jsonRpcRequest);
|
|
||||||
|
|
||||||
// Send the response via the SSE connection instead of HTTP response
|
|
||||||
await SendSseMessageAsync(responseMessage);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
catch (Exception ex)
|
|
||||||
{
|
|
||||||
await SendJsonRpcErrorAsync(response, null, -32603, "Internal error", ex.Message);
|
|
||||||
response.StatusCode = 500;
|
|
||||||
}
|
|
||||||
finally
|
|
||||||
{
|
|
||||||
response.Close();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
private static Task<JsonRpcResponse> ProcessRequestAsync(JsonRpcRequest jsonRpcRequest)
|
|
||||||
{
|
|
||||||
// This is a test server so we just echo back the request
|
|
||||||
return Task.FromResult(new JsonRpcResponse
|
|
||||||
{
|
|
||||||
Id = jsonRpcRequest.Id,
|
|
||||||
Result = jsonRpcRequest.Params!
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
private async Task SendSseMessageAsync(JsonRpcResponse jsonRpcResponse)
|
|
||||||
{
|
|
||||||
// This is a test server so we just send to all connected clients
|
|
||||||
await BroadcastMessageAsync(JsonSerializer.Serialize(jsonRpcResponse));
|
|
||||||
}
|
|
||||||
|
|
||||||
private static async Task SendJsonRpcErrorAsync(HttpListenerResponse response, RequestId? id, int code, string message, string? data = null)
|
|
||||||
{
|
|
||||||
var errorResponse = new JsonRpcError
|
|
||||||
{
|
|
||||||
Id = id ?? new RequestId("error"),
|
|
||||||
JsonRpc = "2.0",
|
|
||||||
Error = new JsonRpcErrorDetail
|
|
||||||
{
|
|
||||||
Code = code,
|
|
||||||
Message = message,
|
|
||||||
Data = data
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
response.StatusCode = 200; // Always 200 for JSON-RPC
|
|
||||||
response.ContentType = "application/json";
|
|
||||||
await using var writer = new StreamWriter(response.OutputStream);
|
|
||||||
await writer.WriteAsync(JsonSerializer.Serialize(errorResponse));
|
|
||||||
}
|
|
||||||
|
|
||||||
private static async Task HandleInitializationRequest(HttpListenerResponse response, JsonRpcRequest jsonRpcRequest)
|
|
||||||
{
|
|
||||||
// We don't need to validate the client's initialization request for the test
|
|
||||||
// Just send back a valid server initialization response
|
|
||||||
var jsonRpcResponse = new JsonRpcResponse()
|
|
||||||
{
|
|
||||||
Id = jsonRpcRequest.Id,
|
|
||||||
Result = JsonSerializer.SerializeToNode(new InitializeResult
|
|
||||||
{
|
|
||||||
ProtocolVersion = "2024-11-05",
|
|
||||||
Capabilities = new(),
|
|
||||||
ServerInfo = new()
|
|
||||||
{
|
|
||||||
Name = "ExampleServer",
|
|
||||||
Version = "1.0.0"
|
|
||||||
}
|
|
||||||
})
|
|
||||||
};
|
|
||||||
|
|
||||||
// Echo back to the HTTP response
|
|
||||||
response.StatusCode = 200;
|
|
||||||
response.ContentType = "application/json";
|
|
||||||
|
|
||||||
await using var writer = new StreamWriter(response.OutputStream);
|
|
||||||
await writer.WriteAsync(JsonSerializer.Serialize(jsonRpcResponse));
|
|
||||||
}
|
|
||||||
|
|
||||||
/// <summary>
|
|
||||||
/// Broadcast a message to all currently connected SSE clients.
|
|
||||||
/// </summary>
|
|
||||||
/// <param name="message">The raw string to send</param>
|
|
||||||
public async Task BroadcastMessageAsync(string message)
|
|
||||||
{
|
|
||||||
foreach (var client in _sseClients.ToArray()) // ToArray to avoid mutation issues
|
|
||||||
{
|
|
||||||
try
|
|
||||||
{
|
|
||||||
// SSE requires "event: <name>" + "data: <payload>" + blank line
|
|
||||||
await client.WriteLineAsync("event: message");
|
|
||||||
await client.WriteLineAsync($"data: {message}");
|
|
||||||
await client.WriteLineAsync();
|
|
||||||
await client.FlushAsync();
|
|
||||||
}
|
|
||||||
catch (IOException)
|
|
||||||
{
|
|
||||||
// Client may have disconnected. We let them get cleaned up on next iteration.
|
|
||||||
}
|
|
||||||
catch (ObjectDisposedException)
|
|
||||||
{
|
|
||||||
// Stream is disposed, ignore.
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
public async ValueTask DisposeAsync()
|
|
||||||
{
|
|
||||||
await _cts.CancelAsync();
|
|
||||||
|
|
||||||
if (_serverTask != null)
|
|
||||||
{
|
|
||||||
try
|
|
||||||
{
|
|
||||||
await Task.WhenAny(_serverTask, Task.Delay(2000));
|
|
||||||
}
|
|
||||||
catch (TaskCanceledException)
|
|
||||||
{
|
|
||||||
// ignore
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
_listener.Close();
|
|
||||||
_cts.Dispose();
|
|
||||||
|
|
||||||
_logger.LogInformation("Test SSE server stopped");
|
|
||||||
}
|
|
||||||
|
|
||||||
/// <summary>
|
|
||||||
/// Send a test notification to all connected SSE clients.
|
|
||||||
/// </summary>
|
|
||||||
/// <param name="content"></param>
|
|
||||||
/// <returns></returns>
|
|
||||||
public async Task SendTestNotificationAsync(string content)
|
|
||||||
{
|
|
||||||
var notification = new JsonRpcNotification
|
|
||||||
{
|
|
||||||
JsonRpc = "2.0",
|
|
||||||
Method = "test/notification",
|
|
||||||
Params = JsonSerializer.SerializeToNode(new { message = content }),
|
|
||||||
};
|
|
||||||
|
|
||||||
var serialized = JsonSerializer.Serialize(notification);
|
|
||||||
await BroadcastMessageAsync(serialized);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
@ -0,0 +1,112 @@
|
||||||
|
using Microsoft.AspNetCore.Connections;
|
||||||
|
using Microsoft.AspNetCore.Http.Features;
|
||||||
|
using System.IO.Pipelines;
|
||||||
|
|
||||||
|
namespace ModelContextProtocol.Tests.Utils;
|
||||||
|
|
||||||
|
public sealed class KestrelInMemoryConnection : ConnectionContext
|
||||||
|
{
|
||||||
|
private readonly Pipe _clientToServerPipe = new();
|
||||||
|
private readonly Pipe _serverToClientPipe = new();
|
||||||
|
private readonly CancellationTokenSource _connectionClosedCts = new CancellationTokenSource();
|
||||||
|
private readonly IFeatureCollection _features = new FeatureCollection();
|
||||||
|
|
||||||
|
public KestrelInMemoryConnection()
|
||||||
|
{
|
||||||
|
ConnectionClosed = _connectionClosedCts.Token;
|
||||||
|
Transport = new DuplexPipe
|
||||||
|
{
|
||||||
|
Input = _clientToServerPipe.Reader,
|
||||||
|
Output = _serverToClientPipe.Writer,
|
||||||
|
};
|
||||||
|
Application = new DuplexPipe
|
||||||
|
{
|
||||||
|
Input = _serverToClientPipe.Reader,
|
||||||
|
Output = _clientToServerPipe.Writer,
|
||||||
|
};
|
||||||
|
ClientStream = new DuplexStream(Application, _connectionClosedCts);
|
||||||
|
}
|
||||||
|
|
||||||
|
public IDuplexPipe Application { get; }
|
||||||
|
public Stream ClientStream { get; }
|
||||||
|
|
||||||
|
public override IDuplexPipe Transport { get; set; }
|
||||||
|
public override string ConnectionId { get; set; } = Guid.NewGuid().ToString("N");
|
||||||
|
|
||||||
|
public override IFeatureCollection Features => _features;
|
||||||
|
|
||||||
|
public override IDictionary<object, object?> Items { get; set; } = new Dictionary<object, object?>();
|
||||||
|
|
||||||
|
public override ValueTask DisposeAsync()
|
||||||
|
{
|
||||||
|
// This is called by Kestrel. The client should dispose the DuplexStream which
|
||||||
|
// completes the other half of these pipes.
|
||||||
|
_serverToClientPipe.Writer.Complete();
|
||||||
|
_serverToClientPipe.Reader.Complete();
|
||||||
|
// Disposing the CTS without waiting for the client would be problematic
|
||||||
|
// except we always dispose the HttpClient before Kestrel in our tests.
|
||||||
|
_connectionClosedCts.Dispose();
|
||||||
|
return base.DisposeAsync();
|
||||||
|
}
|
||||||
|
|
||||||
|
private class DuplexPipe : IDuplexPipe
|
||||||
|
{
|
||||||
|
public required PipeReader Input { get; init; }
|
||||||
|
public required PipeWriter Output { get; init; }
|
||||||
|
}
|
||||||
|
|
||||||
|
private class DuplexStream(IDuplexPipe duplexPipe, CancellationTokenSource connectionClosedCts) : Stream
|
||||||
|
{
|
||||||
|
private readonly Stream _readStream = duplexPipe.Input.AsStream();
|
||||||
|
private readonly Stream _writeStream = duplexPipe.Output.AsStream();
|
||||||
|
|
||||||
|
public override bool CanRead => true;
|
||||||
|
public override bool CanWrite => true;
|
||||||
|
public override bool CanSeek => false;
|
||||||
|
|
||||||
|
public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
|
||||||
|
=> _readStream.ReadAsync(buffer, offset, count, cancellationToken);
|
||||||
|
public override ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
|
||||||
|
=> _readStream.ReadAsync(buffer, cancellationToken);
|
||||||
|
|
||||||
|
public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
|
||||||
|
=> _writeStream.WriteAsync(buffer, offset, count, cancellationToken);
|
||||||
|
public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
|
||||||
|
=> _writeStream.WriteAsync(buffer, cancellationToken);
|
||||||
|
|
||||||
|
public override Task FlushAsync(CancellationToken cancellationToken)
|
||||||
|
=> _writeStream.FlushAsync(cancellationToken);
|
||||||
|
|
||||||
|
protected override void Dispose(bool disposing)
|
||||||
|
{
|
||||||
|
// Signal to the server the the client has closed the connection, and dispose the client-half of the Pipes.
|
||||||
|
connectionClosedCts.Cancel();
|
||||||
|
duplexPipe.Input.Complete();
|
||||||
|
duplexPipe.Output.Complete();
|
||||||
|
}
|
||||||
|
|
||||||
|
// Unsupported stuff
|
||||||
|
public override long Length => throw new NotSupportedException();
|
||||||
|
|
||||||
|
public override long Position
|
||||||
|
{
|
||||||
|
get => throw new NotSupportedException();
|
||||||
|
set => throw new NotSupportedException();
|
||||||
|
}
|
||||||
|
|
||||||
|
public override long Seek(long offset, SeekOrigin origin)
|
||||||
|
{
|
||||||
|
throw new NotSupportedException();
|
||||||
|
}
|
||||||
|
|
||||||
|
public override void SetLength(long value)
|
||||||
|
{
|
||||||
|
throw new NotSupportedException();
|
||||||
|
}
|
||||||
|
|
||||||
|
// Don't bother supporting sync or APM methods. SocketsHttpHandler shouldn't use them.
|
||||||
|
public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException();
|
||||||
|
public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException();
|
||||||
|
public override void Flush() => throw new NotSupportedException();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,38 @@
|
||||||
|
using Microsoft.AspNetCore.Builder;
|
||||||
|
using Microsoft.AspNetCore.Connections;
|
||||||
|
using Microsoft.Extensions.DependencyInjection;
|
||||||
|
using Microsoft.Extensions.DependencyInjection.Extensions;
|
||||||
|
|
||||||
|
namespace ModelContextProtocol.Tests.Utils;
|
||||||
|
|
||||||
|
public class KestrelInMemoryTest : LoggedTest
|
||||||
|
{
|
||||||
|
private readonly KestrelInMemoryTransport _inMemoryTransport = new();
|
||||||
|
|
||||||
|
public KestrelInMemoryTest(ITestOutputHelper testOutputHelper)
|
||||||
|
: base(testOutputHelper)
|
||||||
|
{
|
||||||
|
// Use SlimBuilder instead of EmptyBuilder to avoid having to call UseRouting() and UseEndpoints(_ => { })
|
||||||
|
// or a helper that does the same every test. But clear out the existing socket transport to avoid potential port conflicts.
|
||||||
|
Builder = WebApplication.CreateSlimBuilder();
|
||||||
|
Builder.Services.RemoveAll<IConnectionListenerFactory>();
|
||||||
|
Builder.Services.AddSingleton<IConnectionListenerFactory>(_inMemoryTransport);
|
||||||
|
Builder.Services.AddSingleton(LoggerProvider);
|
||||||
|
}
|
||||||
|
|
||||||
|
public WebApplicationBuilder Builder { get; }
|
||||||
|
|
||||||
|
public HttpClient CreateHttpClient()
|
||||||
|
{
|
||||||
|
var socketsHttpHandler = new SocketsHttpHandler()
|
||||||
|
{
|
||||||
|
ConnectCallback = (context, token) =>
|
||||||
|
{
|
||||||
|
var connection = _inMemoryTransport.CreateConnection();
|
||||||
|
return new(connection.ClientStream);
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
return new HttpClient(socketsHttpHandler);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,50 @@
|
||||||
|
using Microsoft.AspNetCore.Connections;
|
||||||
|
using System.Net;
|
||||||
|
using System.Threading.Channels;
|
||||||
|
|
||||||
|
namespace ModelContextProtocol.Tests.Utils;
|
||||||
|
|
||||||
|
public sealed class KestrelInMemoryTransport : IConnectionListenerFactory, IConnectionListener
|
||||||
|
{
|
||||||
|
private readonly Channel<ConnectionContext> _acceptQueue = Channel.CreateUnbounded<ConnectionContext>();
|
||||||
|
private EndPoint? _endPoint;
|
||||||
|
|
||||||
|
public EndPoint EndPoint => _endPoint ?? throw new InvalidOperationException("EndPoint is not set. Call BindAsync first.");
|
||||||
|
|
||||||
|
public KestrelInMemoryConnection CreateConnection()
|
||||||
|
{
|
||||||
|
var connection = new KestrelInMemoryConnection();
|
||||||
|
_acceptQueue.Writer.TryWrite(connection);
|
||||||
|
return connection;
|
||||||
|
}
|
||||||
|
|
||||||
|
public async ValueTask<ConnectionContext?> AcceptAsync(CancellationToken cancellationToken = default)
|
||||||
|
{
|
||||||
|
if (await _acceptQueue.Reader.WaitToReadAsync(cancellationToken))
|
||||||
|
{
|
||||||
|
while (_acceptQueue.Reader.TryRead(out var item))
|
||||||
|
{
|
||||||
|
return item;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
|
||||||
|
public ValueTask<IConnectionListener> BindAsync(EndPoint endpoint, CancellationToken cancellationToken = default)
|
||||||
|
{
|
||||||
|
_endPoint = endpoint;
|
||||||
|
return new ValueTask<IConnectionListener>(this);
|
||||||
|
}
|
||||||
|
|
||||||
|
public ValueTask DisposeAsync()
|
||||||
|
{
|
||||||
|
return UnbindAsync(default);
|
||||||
|
}
|
||||||
|
|
||||||
|
public ValueTask UnbindAsync(CancellationToken cancellationToken = default)
|
||||||
|
{
|
||||||
|
_acceptQueue.Writer.TryComplete();
|
||||||
|
return default;
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in a new issue