Add ExtractResponse method to Azure.Core's HttpMessage in Azure.Core/ClientModel integration branch

sdk/core/Azure.Core/api/Azure.Core.net461.cs

@@ -489,6 +489,7 @@ public HttpMessage(Azure.Core.Request request, Azure.Core.ResponseClassifier res
         public new Azure.Core.Request Request { get { throw null; } }
         public new Azure.Response Response { get { throw null; } set { } }
         public new Azure.Core.ResponseClassifier ResponseClassifier { get { throw null; } set { } }
+        public new Azure.Response? ExtractResponse() { throw null; }
         public System.IO.Stream? ExtractResponseContent() { throw null; }
         public void SetProperty(string name, object value) { }
         public bool TryGetProperty(string name, out object? value) { throw null; }

sdk/core/Azure.Core/api/Azure.Core.net472.cs

@@ -489,6 +489,7 @@ public HttpMessage(Azure.Core.Request request, Azure.Core.ResponseClassifier res
         public new Azure.Core.Request Request { get { throw null; } }
         public new Azure.Response Response { get { throw null; } set { } }
         public new Azure.Core.ResponseClassifier ResponseClassifier { get { throw null; } set { } }
+        public new Azure.Response? ExtractResponse() { throw null; }
         public System.IO.Stream? ExtractResponseContent() { throw null; }
         public void SetProperty(string name, object value) { }
         public bool TryGetProperty(string name, out object? value) { throw null; }

sdk/core/Azure.Core/api/Azure.Core.net6.0.cs

@@ -489,6 +489,7 @@ public HttpMessage(Azure.Core.Request request, Azure.Core.ResponseClassifier res
         public new Azure.Core.Request Request { get { throw null; } }
         public new Azure.Response Response { get { throw null; } set { } }
         public new Azure.Core.ResponseClassifier ResponseClassifier { get { throw null; } set { } }
+        public new Azure.Response? ExtractResponse() { throw null; }
         public System.IO.Stream? ExtractResponseContent() { throw null; }
         public void SetProperty(string name, object value) { }
         public bool TryGetProperty(string name, out object? value) { throw null; }

sdk/core/Azure.Core/api/Azure.Core.netstandard2.0.cs

@@ -489,6 +489,7 @@ public HttpMessage(Azure.Core.Request request, Azure.Core.ResponseClassifier res
         public new Azure.Core.Request Request { get { throw null; } }
         public new Azure.Response Response { get { throw null; } set { } }
         public new Azure.Core.ResponseClassifier ResponseClassifier { get { throw null; } set { } }
+        public new Azure.Response? ExtractResponse() { throw null; }
         public System.IO.Stream? ExtractResponseContent() { throw null; }
         public void SetProperty(string name, object value) { }
         public bool TryGetProperty(string name, out object? value) { throw null; }

sdk/core/Azure.Core/src/HttpMessage.cs

@@ -44,7 +44,10 @@ public HttpMessage(Request request, ResponseClassifier responseClassifier)
                 if (base.Response is null)
                 {
 #pragma warning disable CA1065 // Do not raise exceptions in unexpected locations
-                    throw new InvalidOperationException("Response was not set, make sure SendAsync was called");
+                    throw new InvalidOperationException($"{nameof(Response)} is not set on this message. "
+                        + "This is may be because the message was not sent via pipeline.Send, "
+                        + "the pipeline transport did not populate the response, or because "
+                        + $"{nameof(ExtractResponse)} was called.");
 #pragma warning restore CA1065 // Do not raise exceptions in unexpected locations
                 }
                 return (Response)base.Response;
@@ -147,11 +150,11 @@ private class MessagePropertyKey { }
         #endregion
 
         /// <summary>
-        /// Returns the response content stream and releases it ownership to the caller.
-        ///
-        /// After calling this method, any attempt to use the
-        /// <see cref="PipelineResponse.ContentStream"/> or <see cref="PipelineResponse.Content"/>
-        /// properties on <see cref="Response"/> will result in an exception being thrown.
+        /// Returns the response content stream and releases its ownership to the caller.
+        /// After this method has been called, any use of the
+        /// <see cref="PipelineResponse.ContentStream"/> or <see cref="Response.Content"/>
+        /// properties on this message will result in an <see cref="InvalidOperationException"/>
+        /// being thrown.
         /// </summary>
         /// <returns>The content stream, or <code>null</code> if <see cref="Response"/>
         /// did not have content set.</returns>
@@ -174,6 +177,15 @@ private class MessagePropertyKey { }
             }
         }
 
+        /// <summary>
+        /// Returns the value of the <see cref="Response"/> property and
+        /// transfers dispose ownership of the response to the caller. After
+        /// calling this method, the <see cref="Response"/> property will be
+        /// null and the caller will be responsible for disposing the returned
+        /// value, which may hold a live network stream.
+        /// </summary>
+        public new Response? ExtractResponse() => (Response?)base.ExtractResponse();
+
         private class ResponseShouldNotBeUsedStream : Stream
         {
             public Stream Original { get; }

sdk/core/Azure.Core/tests/HttpMessageTests.cs

@@ -1,6 +1,9 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using System.Threading.Tasks;
+using System.Threading;
+using System;
 using Azure.Core.Pipeline;
 using Azure.Core.TestFramework;
 using NUnit.Framework;
@@ -380,7 +383,52 @@ public void AppliesHandlerWithLastSetWinsSemantics()
             Assert.IsTrue(message.ResponseClassifier.IsErrorResponse(message));
         }
 
+        [Test]
+        [TestCase(true)]
+        [TestCase(false)]
+        public async Task ResponseStreamAccessibleAfterMessageDisposed(bool buffer)
+        {
+            byte[] serverBytes = new byte[1000];
+            new Random().NextBytes(serverBytes);
+
+            HttpPipeline pipeline = HttpPipelineBuilder.Build(new MockClientOptions { Retry = { NetworkTimeout = Timeout.InfiniteTimeSpan } });
+
+            using TestServer testServer = new(async context =>
+            {
+                await context.Response.Body.WriteAsync(serverBytes, 0, serverBytes.Length).ConfigureAwait(false);
+            });
+
+            Response response;
+            using (HttpMessage message = pipeline.CreateMessage())
+            {
+                message.Request.Uri.Reset(testServer.Address);
+                message.BufferResponse = buffer;
+
+                await pipeline.SendAsync(message, default).ConfigureAwait(false);
+
+                response = message.ExtractResponse();
+
+                Assert.IsFalse(message.HasResponse);
+            }
+
+            Assert.NotNull(response.ContentStream);
+
+            byte[] clientBytes = new byte[serverBytes.Length];
+            int readLength = 0;
+            while (readLength < serverBytes.Length)
+            {
+                readLength += await response.ContentStream.ReadAsync(clientBytes, 0, serverBytes.Length);
+            }
+
+            Assert.AreEqual(serverBytes.Length, readLength);
+            CollectionAssert.AreEqual(serverBytes, clientBytes);
+        }
+
         #region Helpers
+        internal class MockClientOptions : ClientOptions
+        {
+        }
+
         private class StatusCodeHandler : ResponseClassificationHandler
         {
             private readonly int _statusCode;

sdk/core/System.ClientModel/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Added `ExtractResponse` method to `PipelineMessage` to enable returning an undisposed `PipelineResponse` from protocol methods.
+
 ### Breaking Changes
 
 - Change `HttpClientPipelineTransport.Shared` from a field to a property.

sdk/core/System.ClientModel/api/System.ClientModel.net6.0.cs

@@ -149,6 +149,7 @@ protected internal PipelineMessage(System.ClientModel.Primitives.PipelineRequest
         public void Apply(System.ClientModel.Primitives.RequestOptions options) { }
         public void Dispose() { }
         protected virtual void Dispose(bool disposing) { }
+        public System.ClientModel.Primitives.PipelineResponse? ExtractResponse() { throw null; }
         public void SetProperty(System.Type type, object value) { }
         public bool TryGetProperty(System.Type type, out object? value) { throw null; }
     }

sdk/core/System.ClientModel/api/System.ClientModel.netstandard2.0.cs

@@ -148,6 +148,7 @@ protected internal PipelineMessage(System.ClientModel.Primitives.PipelineRequest
         public void Apply(System.ClientModel.Primitives.RequestOptions options) { }
         public void Dispose() { }
         protected virtual void Dispose(bool disposing) { }
+        public System.ClientModel.Primitives.PipelineResponse? ExtractResponse() { throw null; }
         public void SetProperty(System.Type type, object value) { }
         public bool TryGetProperty(System.Type type, out object? value) { throw null; }
     }

sdk/core/System.ClientModel/src/Message/PipelineMessage.cs

@@ -32,6 +32,21 @@ public PipelineResponse? Response
         protected internal set;
     }
 
+    public PipelineResponse? ExtractResponse()
+    {
+        PipelineResponse? response = Response;
+        Response = null;
+        return response;
+    }
+
+    internal void AssertResponse()
+    {
+        if (Response is null)
+        {
+            throw new InvalidOperationException($"'{nameof(Response)}' property is not set on message.");
+        }
+    }
+
     #region Pipeline invocation options
 
     public CancellationToken CancellationToken
@@ -117,41 +132,30 @@ PerTryPolicies is not null ||
 
     #endregion
 
-    internal void AssertResponse()
+    #region IDisposable
+
+    public void Dispose()
     {
-        if (Response is null)
-        {
-            throw new InvalidOperationException("Response is not set on message.");
-        }
+        Dispose(true);
+        GC.SuppressFinalize(this);
     }
 
-    #region IDisposable
-
     protected virtual void Dispose(bool disposing)
     {
         if (disposing && !_disposed)
         {
-            var request = Request;
+            PipelineResponse? response = Response;
+            response?.Dispose();
+            Response = null;
+
+            PipelineRequest request = Request;
             request?.Dispose();
 
             _propertyBag.Dispose();
 
-            var response = Response;
-            if (response != null)
-            {
-                response.Dispose();
-                Response = null;
-            }
-
             _disposed = true;
         }
     }
 
-    public void Dispose()
-    {
-        Dispose(true);
-        GC.SuppressFinalize(this);
-    }
-
     #endregion
 }

sdk/core/System.ClientModel/tests/Message/PipelineMessageTests.cs

@@ -1,14 +1,22 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-using NUnit.Framework;
 using System.ClientModel.Primitives;
 using System.Threading;
+using System.Threading.Tasks;
+using Azure.Core.TestFramework;
+using ClientModel.Tests.Mocks;
+using NUnit.Framework;
+using SyncAsyncTestBase = ClientModel.Tests.SyncAsyncTestBase;
 
 namespace System.ClientModel.Tests.Message;
 
-public class PipelineMessageTests
+public class PipelineMessageTests : SyncAsyncTestBase
 {
+    public PipelineMessageTests(bool isAsync) : base(isAsync)
+    {
+    }
+
     [Test]
     public void ApplyAddsRequestHeaders()
     {
@@ -111,6 +119,48 @@ public void TryGetTypeKeyedPropertyReturnsCorrectValues()
         Assert.AreEqual(4444, ((T4)value!).Value);
     }
 
+    [Test]
+    [TestCase(true)]
+    [TestCase(false)]
+    public async Task ResponseStreamAccessibleAfterMessageDisposed(bool buffer)
+    {
+        byte[] serverBytes = new byte[1000];
+        new Random().NextBytes(serverBytes);
+
+        ClientPipelineOptions options = new() { NetworkTimeout = Timeout.InfiniteTimeSpan };
+        ClientPipeline pipeline = ClientPipeline.Create(options);
+
+        using TestServer testServer = new(async context =>
+        {
+            await context.Response.Body.WriteAsync(serverBytes, 0, serverBytes.Length).ConfigureAwait(false);
+        });
+
+        PipelineResponse? response;
+        using (PipelineMessage message = pipeline.CreateMessage())
+        {
+            message.Request.Uri = testServer.Address;
+            message.BufferResponse = buffer;
+
+            await pipeline.SendSyncOrAsync(message, IsAsync);
+
+            response = message.ExtractResponse();
+
+            Assert.IsNull(message.Response);
+        }
+
+        Assert.NotNull(response!.ContentStream);
+
+        byte[] clientBytes = new byte[serverBytes.Length];
+        int readLength = 0;
+        while (readLength < serverBytes.Length)
+        {
+            readLength += await response.ContentStream!.ReadAsync(clientBytes, 0, serverBytes.Length);
+        }
+
+        Assert.AreEqual(serverBytes.Length, readLength);
+        CollectionAssert.AreEqual(serverBytes, clientBytes);
+    }
+
     #region Helpers
     private struct T1
     {