Apply RequestContext for ProcessMessage (#47049)

Author: live1206

eng/Packages.Data.props

@@ -252,7 +252,7 @@
   </ItemGroup>
 
   <ItemGroup Condition="'$(IsGeneratorLibrary)' == 'true'">
-    <PackageReference Update="Microsoft.Generator.CSharp.ClientModel" Version="1.0.0-alpha.20241030.4" />
+    <PackageReference Update="Microsoft.Generator.CSharp.ClientModel" Version="1.0.0-alpha.20241111.4" />
   </ItemGroup>
 
   <!--

eng/packages/http-client-csharp/generator/Azure.Generator/src/AzureClientPlugin.cs

@@ -24,6 +24,10 @@ public class AzureClientPlugin : ClientModelPlugin
     /// <inheritdoc/>
     public override AzureTypeFactory TypeFactory { get; }
 
+    private AzureOutputLibrary? _azureOutputLibrary;
+    /// <inheritdoc/>
+    public override AzureOutputLibrary OutputLibrary => _azureOutputLibrary ??= new();
+
     /// <summary>
     /// The Azure client plugin to generate the Azure client SDK.
     /// </summary>

eng/packages/http-client-csharp/generator/Azure.Generator/src/AzureOutputLibrary.cs

@@ -0,0 +1,16 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using Azure.Generator.Providers;
+using Microsoft.Generator.CSharp.ClientModel;
+using Microsoft.Generator.CSharp.Providers;
+
+namespace Azure.Generator
+{
+    /// <inheritdoc/>
+    public class AzureOutputLibrary : ScmOutputLibrary
+    {
+        /// <inheritdoc/>
+        protected override TypeProvider[] BuildTypeProviders() => [.. base.BuildTypeProviders(), new RequestContextExtensionsDefinition()];
+    }
+}

eng/packages/http-client-csharp/generator/Azure.Generator/src/Providers/Abstraction/HttpMessageProvider.cs

@@ -24,9 +24,6 @@ public HttpMessageProvider(ValueExpression original) : base(typeof(HttpMessage),
         public override ValueExpression BufferResponse()
             => Original.Property(nameof(HttpMessage.BufferResponse));
 
-        public override MethodBodyStatement[] ExtractResponse()
-            => [Original.Invoke(nameof(HttpMessage.ExtractResponseContent)).Terminate(), Return(Original.Property(nameof(HttpMessage.Response)))];
-
         public override HttpMessageApi FromExpression(ValueExpression original)
             => new HttpMessageProvider(original);
 

eng/packages/http-client-csharp/generator/Azure.Generator/src/Providers/Abstraction/HttpPipelineProvider.cs

@@ -6,7 +6,10 @@
 using Microsoft.Generator.CSharp.ClientModel.Providers;
 using Microsoft.Generator.CSharp.Expressions;
 using Microsoft.Generator.CSharp.Primitives;
+using Microsoft.Generator.CSharp.Providers;
+using Microsoft.Generator.CSharp.Snippets;
 using Microsoft.Generator.CSharp.Statements;
+using System.Threading;
 using static Microsoft.Generator.CSharp.Snippets.Snippet;
 
 namespace Azure.Generator.Providers.Abstraction
@@ -40,10 +43,30 @@ public override ValueExpression PerRetryPolicy(params ValueExpression[] argument
 
         public override ClientPipelineApi ToExpression() => this;
 
-        public override MethodBodyStatement Send(HttpMessageApi message, HttpRequestOptionsApi options)
-            => Original.Invoke(nameof(HttpPipeline.Send), [message, Default]).Terminate();
+        public override MethodBodyStatement[] ProcessMessage(HttpMessageApi message, HttpRequestOptionsApi options)
+            => BuildProcessMessage(message, options, false);
 
-        public override MethodBodyStatement SendAsync(HttpMessageApi message, HttpRequestOptionsApi options)
-            => Original.Invoke(nameof(HttpPipeline.SendAsync), [message, Default], true).Terminate();
+        public override MethodBodyStatement[] ProcessMessageAsync(HttpMessageApi message, HttpRequestOptionsApi options)
+            => BuildProcessMessage(message, options, true);
+
+        private MethodBodyStatement[] BuildProcessMessage(HttpMessageApi message, HttpRequestOptionsApi options, bool isAsync)
+        {
+            var userCancellationToken = new ParameterProvider("userCancellationToken", $"", new CSharpType(typeof(CancellationToken)));
+            var statusOption = new ParameterProvider("statusOption", $"", new CSharpType(typeof(ErrorOptions)));
+            return
+            [
+                new VariableTupleExpression(false, userCancellationToken, statusOption).Assign(options.Invoke("Parse")).Terminate(),
+                Original.Invoke(isAsync ? nameof(HttpPipeline.SendAsync) : nameof(HttpPipeline.Send), [message, userCancellationToken], isAsync).Terminate(),
+                MethodBodyStatement.EmptyLine,
+                new IfStatement(message.Response().IsError().And(new BinaryOperatorExpression("&", options.NullConditional().Property("ErrorOptions"), options.NoThrow()).NotEqual(options.NoThrow())))
+                {
+                    isAsync
+                    ? Throw(AzureClientPlugin.Instance.TypeFactory.ClientResponseApi.ToExpression().CreateAsync(message.Response()))
+                    : Throw(New.Instance(AzureClientPlugin.Instance.TypeFactory.ClientResponseApi.ClientResponseExceptionType, message.Response()))
+                },
+                MethodBodyStatement.EmptyLine,
+                Return(message.Response())
+            ];
+        }
     }
 }

eng/packages/http-client-csharp/generator/Azure.Generator/src/Providers/RequestContextExtensionsDefinition.cs

@@ -0,0 +1,47 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using Microsoft.Generator.CSharp.Expressions;
+using Microsoft.Generator.CSharp.Primitives;
+using Microsoft.Generator.CSharp.Providers;
+using Microsoft.Generator.CSharp.Snippets;
+using Microsoft.Generator.CSharp.Statements;
+using System.IO;
+using System.Threading;
+using static Microsoft.Generator.CSharp.Snippets.Snippet;
+
+namespace Azure.Generator.Providers
+{
+    internal class RequestContextExtensionsDefinition : TypeProvider
+    {
+        protected override TypeSignatureModifiers GetDeclarationModifiers() => TypeSignatureModifiers.Internal | TypeSignatureModifiers.Static;
+
+        protected override string BuildName() => "RequestContextExtensions";
+
+        protected override string BuildRelativeFilePath() => Path.Combine("src", "Generated", "Internal", $"{Name}.cs");
+
+        protected override MethodProvider[] BuildMethods() => [BuildParse()];
+
+        private MethodProvider BuildParse()
+        {
+            var requestContextParameter = new ParameterProvider("requestContext", $"", typeof(RequestContext));
+            var signature = new MethodSignature(
+                "Parse",
+                null,
+                MethodSignatureModifiers.Public | MethodSignatureModifiers.Static | MethodSignatureModifiers.Extension,
+                new CSharpType(typeof((CancellationToken, ErrorOptions))),
+                null,
+                [requestContextParameter]);
+
+            var method = new MethodProvider(signature, new MethodBodyStatement[]
+            {
+                new IfStatement(requestContextParameter.Equal(Null))
+                {
+                    Return(new TupleExpression(Static<CancellationToken>().Property(nameof(CancellationToken.None)), Static<ErrorOptions>().Property(nameof(ErrorOptions.Default))))
+                },
+                Return(new TupleExpression(requestContextParameter.Property(nameof(RequestContext.CancellationToken)), requestContextParameter.Property(nameof(RequestContext.ErrorOptions))))
+            }, this);
+            return method;
+        }
+    }
+}