Fix cancellation token handling in LimitArrayPoolWriteStream.WriteAsync methods (#121836)

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: MihaZupan <[email protected]>
Co-authored-by: Miha Zupan <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

src/libraries/System.Net.Http/src/System/Net/Http/HttpContent.cs

@@ -1103,12 +1103,22 @@ public override void Write(byte[] buffer, int offset, int count)
 
             public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
             {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    return Task.FromCanceled(cancellationToken);
+                }
+
                 Write(buffer, offset, count);
                 return Task.CompletedTask;
             }
 
             public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
             {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    return ValueTask.FromCanceled(cancellationToken);
+                }
+
                 Write(buffer.Span);
                 return default;
             }

src/libraries/System.Net.Http/tests/FunctionalTests/HttpClientTest.cs

@@ -1437,7 +1437,7 @@ private static async Task<T> WhenCanceled<T>(CancellationToken cancellationToken
             return default(T);
         }
 
-        private sealed class CustomResponseHandler : HttpMessageHandler
+        internal sealed class CustomResponseHandler : HttpMessageHandler
         {
             private readonly Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> _func;
 

src/libraries/System.Net.Http/tests/FunctionalTests/HttpContentTest.cs

@@ -596,6 +596,24 @@ await server.AcceptConnectionAsync(async connection =>
                 });
         }
 
+        [Fact]
+        public async Task GetAsync_WithCustomHandlerReturningStringContent_PreCanceledToken_ThrowsTaskCanceledException()
+        {
+            // Regression test for https://github.com/dotnet/runtime/issues/121810
+            // Ensures that when using a custom handler that returns pre-buffered content (StringContent),
+            // calling GetAsync with a pre-canceled token properly throws TaskCanceledException
+            using var handler = new HttpClientTest.CustomResponseHandler((req, ct) => Task.FromResult(new HttpResponseMessage()
+            {
+                Content = new StringContent("Hello from CustomHandler")
+            }));
+            using var httpClient = new HttpClient(handler);
+            var cts = new CancellationTokenSource();
+            cts.Cancel();
+
+            await Assert.ThrowsAsync<TaskCanceledException>(() =>
+                httpClient.GetAsync("https://example.com", cts.Token));
+        }
+
         [Theory]
         [InlineData(true)]
         [InlineData(false)]