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

tests/Common/Utils/DelegatingTestOutputHelper.cs

@@ -1,6 +1,6 @@
 namespace ModelContextProtocol.Tests.Utils;
 
-public class DelegatingTestOutputHelper : ITestOutputHelper
+public class DelegatingTestOutputHelper() : ITestOutputHelper
 {
     public ITestOutputHelper? CurrentTestOutputHelper { get; set; }
 

tests/Common/Utils/MockLoggerProvider.cs

@@ -0,0 +1,32 @@
+using Microsoft.Extensions.Logging;
+using System.Collections.Concurrent;
+
+namespace ModelContextProtocol.Tests.Utils;
+
+public class MockLoggerProvider() : ILoggerProvider
+{
+    public ConcurrentQueue<(string Category, LogLevel LogLevel, string Message, Exception? Exception)> LogMessages { get; } = [];
+
+    public ILogger CreateLogger(string categoryName)
+    {
+        return new MockLogger(this, categoryName);
+    }
+
+    public void Dispose()
+    {
+    }
+
+    private class MockLogger(MockLoggerProvider mockProvider, string category) : ILogger
+    {
+        public void Log<TState>(
+            LogLevel logLevel, EventId eventId, TState state, Exception? exception, Func<TState, Exception?, string> formatter)
+        {
+            mockProvider.LogMessages.Enqueue((category, logLevel, formatter(state, exception), exception));
+        }
+
+        public bool IsEnabled(LogLevel logLevel) => true;
+
+        // The MockLoggerProvider is a convenient NoopDisposable
+        public IDisposable BeginScope<TState>(TState state) where TState : notnull => mockProvider;
+    }
+}

tests/Common/Utils/TestServerTransport.cs

@@ -1,6 +1,7 @@
 using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Protocol.Types;
+using ModelContextProtocol.Utils.Json;
 using System.Text.Json;
 using System.Threading.Channels;
 
@@ -32,7 +33,7 @@ public class TestServerTransport : ITransport
     {
         _messageChannel.Writer.TryComplete();
         IsConnected = false;
-        return ValueTask.CompletedTask;
+        return default;
     }
 
     public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
@@ -63,7 +64,7 @@ public class TestServerTransport : ITransport
             Result = JsonSerializer.SerializeToNode(new ListRootsResult
             {
                 Roots = []
-            }),
+            }, McpJsonUtilities.DefaultOptions),
         }, cancellationToken);
     }
 
@@ -72,7 +73,7 @@ public class TestServerTransport : ITransport
         await WriteMessageAsync(new JsonRpcResponse
         {
             Id = request.Id,
-            Result = JsonSerializer.SerializeToNode(new CreateMessageResult { Content = new(), Model = "model", Role = Role.Assistant }),
+            Result = JsonSerializer.SerializeToNode(new CreateMessageResult { Content = new(), Model = "model", Role = Role.User }, McpJsonUtilities.DefaultOptions),
         }, cancellationToken);
     }
 

tests/ModelContextProtocol.AspNetCore.Tests/ModelContextProtocol.AspNetCore.Tests.csproj

@@ -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" />

tests/ModelContextProtocol.AspNetCore.Tests/StreamableHttpTests.cs

@@ -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);
     }
 

tests/ModelContextProtocol.Tests/Client/McpClientFactoryTests.cs

@@ -45,7 +45,10 @@ public class McpClientFactoryTests
         Task t = McpClientFactory.CreateAsync(
             new StreamClientTransport(new Pipe().Writer.AsStream(), new Pipe().Reader.AsStream()),
             cancellationToken: cts.Token);
-        Assert.False(t.IsCompleted);
+        if (!preCanceled)
+        {
+            Assert.False(t.IsCompleted);
+        }
 
         if (!preCanceled)
         {

tests/ModelContextProtocol.Tests/ClientServerTestBase.cs

@@ -62,7 +62,7 @@ public abstract class ClientServerTestBase : LoggedTest, IAsyncDisposable
         Dispose();
     }
 
-    protected async Task<IMcpClient> CreateMcpClientForServer(McpClientOptions? options = null)
+    protected async Task<IMcpClient> CreateMcpClientForServer()
     {
         return await McpClientFactory.CreateAsync(
             new StreamClientTransport(

tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj

@@ -24,6 +24,10 @@
     <DisableTestingPlatformServerCapability>true</DisableTestingPlatformServerCapability>
   </PropertyGroup>
 
+  <ItemGroup>
+    <Compile Include="..\Common\**\*.cs" />
+  </ItemGroup>
+
   <ItemGroup>
     <PackageReference Include="coverlet.collector">
       <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
@@ -55,13 +59,6 @@
     <ProjectReference Include="..\..\samples\TestServerWithHosting\TestServerWithHosting.csproj" />
   </ItemGroup>
 
-  <!-- Dependencies only needed by net472 -->
-  <ItemGroup Condition="'$(TargetFramework)' == 'net472'">
-    <Compile Include="..\..\src\Common\Throw.cs" Link="Utils\Throw.cs" />
-    <Compile Include="..\..\src\Common\Polyfills\**\*.cs" />
-    <PackageReference Include="System.Net.Http" />
-  </ItemGroup>
-
   <ItemGroup>
     <Content Condition="$([MSBuild]::IsOSPlatform('Windows'))" Include="$([System.IO.Path]::GetFullPath('$(ArtifactsBinDir)ModelContextProtocol.TestServer\$(Configuration)'))\$(TargetFramework)\TestServer.exe">
       <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>