Work around Kestrel bug that could cause test hangs

Author: halter73

tests/ModelContextProtocol.AspNetCore.Tests/Utils/KestrelInMemoryConnection.cs

@@ -65,10 +65,21 @@ private class DuplexStream(IDuplexPipe duplexPipe, CancellationTokenSource conne
         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 async Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        {
+            // Normally, Kestrel will trigger RequestAborted when the connectionClosedCts fires causing it to gracefully close
+            // the connection. However, there's currently a race condition that can cause this to get missed. This at least
+            // unblocks HttpConnection.SendAsync when it disposes the underlying connection stream while awaiting the _readAheadTask
+            // as would happen with a real socket. https://github.com/dotnet/aspnetcore/pull/62385
+            using var linkedTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, connectionClosedCts.Token);
+            return await _readStream.ReadAsync(buffer, offset, count, linkedTokenSource.Token);
+        }
+
+        public override async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
+        {
+            using var linkedTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, connectionClosedCts.Token);
+            return await _readStream.ReadAsync(buffer, linkedTokenSource.Token);
+        }
 
         public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
             => _writeStream.WriteAsync(buffer, offset, count, cancellationToken);