Improve testing for StreamableHttpHandler and IdleTrackingBackgroundService (#345)

* Add idle session pruning tests
* Remove redundant CTS from StreamableHttpHandler
  * StreamableHttpServerTransport.HandlePostRequest already creates the equivelant CTS
    which DeleteRequest_CompletesSession_WhichCancelsLongRunningToolCalls verifies
* Fix accept header validation for GET
* Use InvalidOperationException for transport-not-connected errors in doc comments
* Fix flaky Cancellation_ThrowsCancellationException test
This commit is contained in:
Stephen Halter 2025-04-24 10:49:02 -07:00 • committed by GitHub
commit e3ff552824
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 315 additions and 291 deletions

View file

@ -22,6 +22,10 @@
<DisableTestingPlatformServerCapability>true</DisableTestingPlatformServerCapability>
</PropertyGroup>
<ItemGroup>
<Compile Include="..\Common\**\*.cs" />
</ItemGroup>
<ItemGroup>
<PackageReference Include="coverlet.collector">
<IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
@ -35,6 +39,7 @@
<PackageReference Include="Microsoft.Extensions.AI.OpenAI" />
<PackageReference Include="Microsoft.Extensions.Logging" />
<PackageReference Include="Microsoft.Extensions.Logging.Console" />
<PackageReference Include="Microsoft.Extensions.TimeProvider.Testing" />
<PackageReference Include="Microsoft.NET.Test.Sdk" />
<PackageReference Include="Moq" />
<PackageReference Include="OpenTelemetry" />

View file

@ -1,9 +1,13 @@
using Microsoft.AspNetCore.Builder;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Time.Testing;
using Microsoft.Net.Http.Headers;
using ModelContextProtocol.AspNetCore.Tests.Utils;
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Protocol.Types;
using ModelContextProtocol.Server;
using ModelContextProtocol.Tests.Utils;
using ModelContextProtocol.Utils.Json;
using System.Net;
using System.Net.ServerSentEvents;
@ -27,8 +31,6 @@ public class StreamableHttpTests(ITestOutputHelper outputHelper) : KestrelInMemo
private async Task StartAsync()
{
AddDefaultHttpClientRequestHeaders();
Builder.Services.AddMcpServer(options =>
{
options.ServerInfo = new Implementation
@ -43,6 +45,9 @@ public class StreamableHttpTests(ITestOutputHelper outputHelper) : KestrelInMemo
_app.MapMcp();
await _app.StartAsync(TestContext.Current.CancellationToken);
HttpClient.DefaultRequestHeaders.Accept.Add(new("application/json"));
HttpClient.DefaultRequestHeaders.Accept.Add(new("text/event-stream"));
}
public async ValueTask DisposeAsync()
@ -54,6 +59,31 @@ public class StreamableHttpTests(ITestOutputHelper outputHelper) : KestrelInMemo
base.Dispose();
}
[Fact]
public async Task NegativeNonInfiniteIdleTimeout_Throws_ArgumentOutOfRangeException()
{
Builder.Services.AddMcpServer().WithHttpTransport(options =>
{
options.IdleTimeout = TimeSpan.MinValue;
});
var ex = await Assert.ThrowsAnyAsync<ArgumentOutOfRangeException>(StartAsync);
Assert.Contains("IdleTimeout", ex.Message);
}
[Fact]
public async Task NegativeMaxIdleSessionCount_Throws_ArgumentOutOfRangeException()
{
Builder.Services.AddMcpServer().WithHttpTransport(options =>
{
options.MaxIdleSessionCount = -1;
});
var ex = await Assert.ThrowsAnyAsync<ArgumentOutOfRangeException>(StartAsync);
Assert.Contains("MaxIdleSessionCount", ex.Message);
}
[Fact]
public async Task InitialPostResponse_Includes_McpSessionIdHeader()
{
@ -74,26 +104,16 @@ public class StreamableHttpTests(ITestOutputHelper outputHelper) : KestrelInMemo
Assert.Equal(HttpStatusCode.UnsupportedMediaType, response.StatusCode);
}
[Fact]
public async Task PostRequest_IsNotAcceptable_WithoutApplicationJsonAcceptHeader()
[Theory]
[InlineData("text/event-stream")]
[InlineData("application/json")]
[InlineData("application/json-text/event-stream")]
public async Task PostRequest_IsNotAcceptable_WithSingleAcceptHeader(string singleAcceptValue)
{
await StartAsync();
HttpClient.DefaultRequestHeaders.Accept.Clear();
HttpClient.DefaultRequestHeaders.Accept.Add(new("text/event-stream"));
using var response = await HttpClient.PostAsync("", JsonContent(InitializeRequest), TestContext.Current.CancellationToken);
Assert.Equal(HttpStatusCode.NotAcceptable, response.StatusCode);
}
[Fact]
public async Task PostRequest_IsNotAcceptable_WithoutTextEventStreamAcceptHeader()
{
await StartAsync();
HttpClient.DefaultRequestHeaders.Accept.Clear();
HttpClient.DefaultRequestHeaders.Accept.Add(new("text/json"));
HttpClient.DefaultRequestHeaders.TryAddWithoutValidation(HeaderNames.Accept, singleAcceptValue);
using var response = await HttpClient.PostAsync("", JsonContent(InitializeRequest), TestContext.Current.CancellationToken);
Assert.Equal(HttpStatusCode.NotAcceptable, response.StatusCode);
@ -105,7 +125,7 @@ public class StreamableHttpTests(ITestOutputHelper outputHelper) : KestrelInMemo
await StartAsync();
HttpClient.DefaultRequestHeaders.Accept.Clear();
HttpClient.DefaultRequestHeaders.Accept.Add(new("text/json"));
HttpClient.DefaultRequestHeaders.Accept.Add(new("application/json"));
using var response = await HttpClient.GetAsync("", TestContext.Current.CancellationToken);
Assert.Equal(HttpStatusCode.NotAcceptable, response.StatusCode);
@ -131,7 +151,6 @@ public class StreamableHttpTests(ITestOutputHelper outputHelper) : KestrelInMemo
[Fact]
public async Task InitializeRequest_Matches_CustomRoute()
{
AddDefaultHttpClientRequestHeaders();
Builder.Services.AddMcpServer().WithHttpTransport();
await using var app = Builder.Build();
@ -139,6 +158,8 @@ public class StreamableHttpTests(ITestOutputHelper outputHelper) : KestrelInMemo
await app.StartAsync(TestContext.Current.CancellationToken);
HttpClient.DefaultRequestHeaders.Accept.Add(new("application/json"));
HttpClient.DefaultRequestHeaders.Accept.Add(new("text/event-stream"));
using var response = await HttpClient.PostAsync("/custom-route", JsonContent(InitializeRequest), TestContext.Current.CancellationToken);
Assert.Equal(HttpStatusCode.OK, response.StatusCode);
}
@ -312,10 +333,21 @@ public class StreamableHttpTests(ITestOutputHelper outputHelper) : KestrelInMemo
for (int i = 0; i < longRunningToolTasks.Length; i++)
{
longRunningToolTasks[i] = CallLongRunningToolAsync();
}
var getResponse = await HttpClient.GetAsync("", HttpCompletionOption.ResponseHeadersRead, TestContext.Current.CancellationToken);
for (int i = 0; i < longRunningToolTasks.Length; i++)
{
Assert.False(longRunningToolTasks[i].IsCompleted);
}
await HttpClient.DeleteAsync("", TestContext.Current.CancellationToken);
// Get request should complete gracefully.
var sseResponseBody = await getResponse.Content.ReadAsStringAsync(TestContext.Current.CancellationToken);
Assert.Empty(sseResponseBody);
// Currently, the OCE thrown by the canceled session is unhandled and turned into a 500 error by Kestrel.
// The spec suggests sending CancelledNotifications. That would be good, but we can do that later.
// For now, the important thing is that request completes without indicating success.
@ -361,10 +393,98 @@ public class StreamableHttpTests(ITestOutputHelper outputHelper) : KestrelInMemo
Assert.Equal(11, currentSseItem);
}
private void AddDefaultHttpClientRequestHeaders()
[Fact]
public async Task IdleSessions_ArePruned_AfterIdleTimeout()
{
HttpClient.DefaultRequestHeaders.Accept.Add(new("application/json"));
HttpClient.DefaultRequestHeaders.Accept.Add(new("text/event-stream"));
var fakeTimeProvider = new FakeTimeProvider();
Builder.Services.AddMcpServer().WithHttpTransport(options =>
{
Assert.Equal(TimeSpan.FromHours(2), options.IdleTimeout);
options.TimeProvider = fakeTimeProvider;
});
await StartAsync();
await CallInitializeAndValidateAsync();
await CallEchoAndValidateAsync();
// Add 5 seconds to idle timeout to account for the interval of the PeriodicTimer.
fakeTimeProvider.Advance(TimeSpan.FromHours(2) + TimeSpan.FromSeconds(5));
using var response = await HttpClient.PostAsync("", JsonContent(EchoRequest), TestContext.Current.CancellationToken);
Assert.Equal(HttpStatusCode.NotFound, response.StatusCode);
}
[Fact]
public async Task IdleSessions_AreNotPruned_WithInfiniteIdleTimeoutWhileUnderMaxIdleSessionCount()
{
var fakeTimeProvider = new FakeTimeProvider();
Builder.Services.AddMcpServer().WithHttpTransport(options =>
{
options.IdleTimeout = Timeout.InfiniteTimeSpan;
options.TimeProvider = fakeTimeProvider;
});
await StartAsync();
await CallInitializeAndValidateAsync();
await CallEchoAndValidateAsync();
fakeTimeProvider.Advance(TimeSpan.FromDays(1));
// Echo still works because the session has not been pruned.
await CallEchoAndValidateAsync();
}
[Fact]
public async Task IdleSessionsPastMaxIdleSessionCount_ArePruned_LongestIdleFirstDespiteIdleTimeout()
{
var fakeTimeProvider = new FakeTimeProvider();
Builder.Services.AddMcpServer().WithHttpTransport(options =>
{
options.IdleTimeout = Timeout.InfiniteTimeSpan;
options.MaxIdleSessionCount = 2;
options.TimeProvider = fakeTimeProvider;
});
var mockLoggerProvider = new MockLoggerProvider();
Builder.Logging.AddProvider(mockLoggerProvider);
await StartAsync();
// Start first session.
var firstSessionId = await CallInitializeAndValidateAsync();
// Start a second session to trigger pruning of the original session.
fakeTimeProvider.Advance(TimeSpan.FromTicks(1));
var secondSessionId = await CallInitializeAndValidateAsync();
Assert.NotEqual(firstSessionId, secondSessionId);
// First session ID still works, since we allow up to 2 idle sessions.
fakeTimeProvider.Advance(TimeSpan.FromTicks(1));
SetSessionId(firstSessionId);
await CallEchoAndValidateAsync();
// Start a third session to trigger pruning of the first session.
fakeTimeProvider.Advance(TimeSpan.FromTicks(1));
var thirdSessionId = await CallInitializeAndValidateAsync();
Assert.NotEqual(secondSessionId, thirdSessionId);
// Pruning of the second session results in a 404 since we used the first session more recently.
fakeTimeProvider.Advance(TimeSpan.FromSeconds(10));
SetSessionId(secondSessionId);
using var response = await HttpClient.PostAsync("", JsonContent(EchoRequest), TestContext.Current.CancellationToken);
Assert.Equal(HttpStatusCode.NotFound, response.StatusCode);
// But the first and third session IDs should still work.
SetSessionId(firstSessionId);
await CallEchoAndValidateAsync();
SetSessionId(thirdSessionId);
await CallEchoAndValidateAsync();
var logMessage = Assert.Single(mockLoggerProvider.LogMessages, m => m.LogLevel == LogLevel.Critical);
Assert.StartsWith("Exceeded maximum of 2 idle sessions.", logMessage.Message);
}
private static StringContent JsonContent(string json) => new StringContent(json, Encoding.UTF8, "application/json");
@ -437,7 +557,7 @@ public class StreamableHttpTests(ITestOutputHelper outputHelper) : KestrelInMemo
private string CallToolWithProgressToken(string toolName, string arguments = "{}") =>
Request("tools/call", $$$"""
{"name":"{{{toolName}}}","arguments":{{{arguments}}}, "_meta":{"progressToken": "abc123"}}
{"name":"{{{toolName}}}","arguments":{{{arguments}}},"_meta":{"progressToken":"abc123"}}
""");
private static InitializeResult AssertServerInfo(JsonRpcResponse rpcResponse)
@ -457,13 +577,21 @@ public class StreamableHttpTests(ITestOutputHelper outputHelper) : KestrelInMemo
return callToolResponse;
}
private async Task CallInitializeAndValidateAsync()
private async Task<string> CallInitializeAndValidateAsync()
{
HttpClient.DefaultRequestHeaders.Remove("mcp-session-id");
using var response = await HttpClient.PostAsync("", JsonContent(InitializeRequest), TestContext.Current.CancellationToken);
var rpcResponse = await AssertSingleSseResponseAsync(response);
AssertServerInfo(rpcResponse);
var sessionId = Assert.Single(response.Headers.GetValues("mcp-session-id"));
SetSessionId(sessionId);
return sessionId;
}
private void SetSessionId(string sessionId)
{
HttpClient.DefaultRequestHeaders.Remove("mcp-session-id");
HttpClient.DefaultRequestHeaders.Add("mcp-session-id", sessionId);
}

View file

@ -1,13 +0,0 @@
namespace ModelContextProtocol.Tests.Utils;
public class DelegatingTestOutputHelper : ITestOutputHelper
{
public ITestOutputHelper? CurrentTestOutputHelper { get; set; }
public string Output => CurrentTestOutputHelper?.Output ?? string.Empty;
public void Write(string message) => CurrentTestOutputHelper?.Write(message);
public void Write(string format, params object[] args) => CurrentTestOutputHelper?.Write(format, args);
public void WriteLine(string message) => CurrentTestOutputHelper?.WriteLine(message);
public void WriteLine(string format, params object[] args) => CurrentTestOutputHelper?.WriteLine(format, args);
}

View file

@ -1,30 +0,0 @@
using Microsoft.Extensions.Logging;
namespace ModelContextProtocol.Tests.Utils;
public class LoggedTest : IDisposable
{
private readonly DelegatingTestOutputHelper _delegatingTestOutputHelper;
public LoggedTest(ITestOutputHelper testOutputHelper)
{
_delegatingTestOutputHelper = new()
{
CurrentTestOutputHelper = testOutputHelper,
};
LoggerProvider = new XunitLoggerProvider(_delegatingTestOutputHelper);
LoggerFactory = Microsoft.Extensions.Logging.LoggerFactory.Create(builder =>
{
builder.AddProvider(LoggerProvider);
});
}
public ITestOutputHelper TestOutputHelper => _delegatingTestOutputHelper;
public ILoggerFactory LoggerFactory { get; }
public ILoggerProvider LoggerProvider { get; }
public virtual void Dispose()
{
_delegatingTestOutputHelper.CurrentTestOutputHelper = null;
}
}

View file

@ -1,20 +0,0 @@
namespace ModelContextProtocol.Tests.Utils;
public class MockHttpHandler : HttpMessageHandler
{
public Func<HttpRequestMessage, Task<HttpResponseMessage>>? RequestHandler { get; set; }
protected async override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
{
if (RequestHandler == null)
throw new InvalidOperationException($"No {nameof(RequestHandler)} was set! Please set handler first and make request afterwards.");
cancellationToken.ThrowIfCancellationRequested();
var result = await RequestHandler.Invoke(request);
cancellationToken.ThrowIfCancellationRequested();
return result;
}
}

View file

@ -1,83 +0,0 @@
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Protocol.Types;
using System.Text.Json;
using System.Threading.Channels;
namespace ModelContextProtocol.Tests.Utils;
public class TestServerTransport : ITransport
{
private readonly Channel<JsonRpcMessage> _messageChannel;
public bool IsConnected { get; set; }
public ChannelReader<JsonRpcMessage> MessageReader => _messageChannel;
public List<JsonRpcMessage> SentMessages { get; } = [];
public Action<JsonRpcMessage>? OnMessageSent { get; set; }
public TestServerTransport()
{
_messageChannel = Channel.CreateUnbounded<JsonRpcMessage>(new UnboundedChannelOptions
{
SingleReader = true,
SingleWriter = true,
});
IsConnected = true;
}
public ValueTask DisposeAsync()
{
_messageChannel.Writer.TryComplete();
IsConnected = false;
return ValueTask.CompletedTask;
}
public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
{
SentMessages.Add(message);
if (message is JsonRpcRequest request)
{
if (request.Method == RequestMethods.RootsList)
await ListRoots(request, cancellationToken);
else if (request.Method == RequestMethods.SamplingCreateMessage)
await Sampling(request, cancellationToken);
else
await WriteMessageAsync(request, cancellationToken);
}
else if (message is JsonRpcNotification notification)
{
await WriteMessageAsync(notification, cancellationToken);
}
OnMessageSent?.Invoke(message);
}
private async Task ListRoots(JsonRpcRequest request, CancellationToken cancellationToken)
{
await WriteMessageAsync(new JsonRpcResponse
{
Id = request.Id,
Result = JsonSerializer.SerializeToNode(new ListRootsResult
{
Roots = []
}),
}, cancellationToken);
}
private async Task Sampling(JsonRpcRequest request, CancellationToken cancellationToken)
{
await WriteMessageAsync(new JsonRpcResponse
{
Id = request.Id,
Result = JsonSerializer.SerializeToNode(new CreateMessageResult { Content = new(), Model = "model", Role = Role.Assistant }),
}, cancellationToken);
}
private async Task WriteMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
{
await _messageChannel.Writer.WriteAsync(message, cancellationToken);
}
}

View file

@ -1,52 +0,0 @@
using System.Globalization;
using System.Text;
using Microsoft.Extensions.Logging;
namespace ModelContextProtocol.Tests.Utils;
public class XunitLoggerProvider(ITestOutputHelper output) : ILoggerProvider
{
public ILogger CreateLogger(string categoryName)
{
return new XunitLogger(output, categoryName);
}
public void Dispose()
{
}
private class XunitLogger(ITestOutputHelper output, string category) : ILogger
{
public void Log<TState>(
LogLevel logLevel, EventId eventId, TState state, Exception? exception, Func<TState, Exception?, string> formatter)
{
var sb = new StringBuilder();
var timestamp = DateTimeOffset.UtcNow.ToString("s", CultureInfo.InvariantCulture);
var prefix = $"| [{timestamp}] {category} {logLevel}: ";
var lines = formatter(state, exception);
sb.Append(prefix);
sb.Append(lines);
if (exception is not null)
{
sb.AppendLine();
sb.Append(exception.ToString());
}
output.WriteLine(sb.ToString());
}
public bool IsEnabled(LogLevel logLevel) => true;
public IDisposable BeginScope<TState>(TState state) where TState : notnull
=> new NoopDisposable();
private sealed class NoopDisposable : IDisposable
{
public void Dispose()
{
}
}
}
}