Fix generation (Azure#50488)

Author: live1206

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/Providers/OperationMethodProviders/ResourceOperationMethodProvider.cs

@@ -34,6 +34,10 @@ internal class ResourceOperationMethodProvider
         protected readonly MethodSignature _signature;
         protected readonly MethodBodyStatement[] _bodyStatements;
 
+        private readonly CSharpType? _responseGenericType;
+        private readonly bool _isGeneric;
+        private readonly bool _isLongRunningOperation;
+
         public ResourceOperationMethodProvider(
             ResourceClientProvider resourceClientProvider,
             InputServiceMethod method,
@@ -44,6 +48,9 @@ public ResourceOperationMethodProvider(
             _serviceMethod = method;
             _convenienceMethod = convenienceMethod;
             _isAsync = isAsync;
+            _responseGenericType = _serviceMethod.GetResponseBodyType();
+            _isGeneric = _responseGenericType != null;
+            _isLongRunningOperation = _serviceMethod.IsLongRunningOperation();
             _clientDiagnosticsField = resourceClientProvider.GetClientDiagnosticsField();
             _signature = CreateSignature();
             _bodyStatements = BuildBodyStatements();
@@ -87,10 +94,9 @@ protected virtual MethodSignature CreateSignature()
 
         private TryExpression BuildTryExpression()
         {
-            var cancellationTokenParameter = _convenienceMethod.Signature.Parameters.Single(p => p.Type.Equals(typeof(CancellationToken)));
+            var cancellationTokenParameter = _convenienceMethod.Signature.Parameters.FirstOrDefault(p => p.Type.Equals(typeof(CancellationToken))) ?? KnownParameters.CancellationTokenParameter;
 
             var requestMethod = _resourceClientProvider.GetClientProvider().GetRequestMethodByOperation(_serviceMethod.Operation);
-
             var tryStatements = new List<MethodBodyStatement>
             {
                 ResourceMethodSnippets.CreateRequestContext(cancellationTokenParameter, out var contextVariable)
@@ -103,7 +109,7 @@ private TryExpression BuildTryExpression()
 
             tryStatements.AddRange(BuildClientPipelineProcessing(messageVariable, contextVariable, out var responseVariable));
 
-            if (_serviceMethod.IsLongRunningOperation())
+            if (_isLongRunningOperation)
             {
                 tryStatements.AddRange(
                 BuildLroHandling(
@@ -123,22 +129,17 @@ private IReadOnlyList<MethodBodyStatement> BuildClientPipelineProcessing(
             VariableExpression contextVariable,
             out VariableExpression responseVariable)
         {
-            var responseType = _convenienceMethod.Signature.ReturnType!.UnWrapAsync();
-
-            // Check if the return type is a generic Response<T> or just Response
-            if (!responseType.Equals(typeof(Response)))
+            if (_isGeneric && !_isLongRunningOperation)
             {
-                // For methods returning Response<T> (e.g., Response<MyResource>), use generic response processing
                 return ResourceMethodSnippets.CreateGenericResponsePipelineProcessing(
                     messageVariable,
                     contextVariable,
-                    responseType,
+                    _responseGenericType!,
                     _isAsync,
                     out responseVariable);
             }
             else
             {
-                // For methods returning just Response (no generic type), use non-generic response processing
                 return ResourceMethodSnippets.CreateNonGenericResponsePipelineProcessing(
                     messageVariable,
                     contextVariable,
@@ -155,9 +156,8 @@ private IReadOnlyList<MethodBodyStatement> BuildLroHandling(
             var statements = new List<MethodBodyStatement>();
 
             var finalStateVia = _serviceMethod.GetOperationFinalStateVia();
-            bool isGeneric = _serviceMethod.GetResponseBodyType() != null;
 
-            var armOperationType = isGeneric
+            var armOperationType = _isGeneric
                 ? ManagementClientGenerator.Instance.OutputLibrary.GenericArmOperation.Type
                     .MakeGenericType([_resourceClientProvider.ResourceClientCSharpType])
                 : ManagementClientGenerator.Instance.OutputLibrary.ArmOperation.Type;
@@ -166,11 +166,11 @@ private IReadOnlyList<MethodBodyStatement> BuildLroHandling(
                 _clientDiagnosticsField,
                 This.Property("Pipeline"),
                 messageVariable.Property("Request"),
-                isGeneric ? responseVariable.Invoke("GetRawResponse") : responseVariable,
+                responseVariable,
                 Static(typeof(OperationFinalStateVia)).Property(finalStateVia.ToString())
             ];
 
-            var operationInstanceArguments = isGeneric
+            var operationInstanceArguments = _isGeneric
                 ? [
                     New.Instance(_resourceClientProvider.Source.Type, This.Property("Client")),
                     .. armOperationArguments
@@ -184,13 +184,13 @@ .. armOperationArguments
                 out var operationVariable);
             statements.Add(operationDeclaration);
 
-            var waitMethod = isGeneric
+            var waitMethod = _isGeneric
                 ? (_isAsync ? "WaitForCompletionAsync" : "WaitForCompletion")
                 : (_isAsync ? "WaitForCompletionResponseAsync" : "WaitForCompletionResponse");
 
             var waitInvocation = _isAsync
-                ? operationVariable.Invoke(waitMethod, [cancellationTokenParameter], null, _isAsync).Terminate()
-                : operationVariable.Invoke(waitMethod, cancellationTokenParameter).Terminate();
+                           ? operationVariable.Invoke(waitMethod, [cancellationTokenParameter], null, _isAsync).Terminate()
+                           : operationVariable.Invoke(waitMethod, cancellationTokenParameter).Terminate();
 
             var waitIfCompletedStatement = new IfStatement(
                 KnownAzureParameters.WaitUntil.Equal(
@@ -243,7 +243,11 @@ protected IReadOnlyList<ParameterProvider> GetOperationMethodParameters()
 
             foreach (var parameter in _convenienceMethod.Signature.Parameters)
             {
-                if (!_resourceClientProvider.ImplicitParameterNames.Contains(parameter.Name))
+                if (parameter.Type.Equals(typeof(RequestContext)))
+                {
+                    result.Add(KnownParameters.CancellationTokenParameter);
+                }
+                else if (!_resourceClientProvider.ImplicitParameterNames.Contains(parameter.Name))
                 {
                     result.Add(parameter);
                 }

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/Providers/ResourceClientProvider.cs

@@ -247,23 +247,23 @@ protected override MethodProvider[] BuildMethods()
                     continue;
                 }
 
-                // Check if this is an update operation (PUT method for non-singleton resource)
-                var isUpdateOperation = method.Operation.HttpMethod == HttpMethod.Put.ToString() && !IsSingleton;
+                // Check if this is an update operation (PUT or Patch method for non-singleton resource)
+                var isUpdateOperation = (method.Operation.HttpMethod == HttpMethod.Put.ToString() || method.Operation.HttpMethod == HttpMethod.Patch.ToString()) && !IsSingleton;
 
                 if (isUpdateOperation)
                 {
                     var updateMethodProvider = new UpdateOperationMethodProvider(this, method, convenienceMethod, false);
                     operationMethods.Add(updateMethodProvider);
 
-                    var asyncConvenienceMethod = GetCorrespondingConvenienceMethod(method.Operation, true);
+                    var asyncConvenienceMethod = _restClientProvider.GetConvenienceMethodByOperation(method.Operation, true);
                     var updateAsyncMethodProvider = new UpdateOperationMethodProvider(this, method, asyncConvenienceMethod, true);
                     operationMethods.Add(updateAsyncMethodProvider);
                 }
                 else
                 {
-                    operationMethods.Add(BuildOperationMethod(method, convenienceMethod, false));
-                    var asyncConvenienceMethod = GetCorrespondingConvenienceMethod(method.Operation, true);
-                    operationMethods.Add(BuildOperationMethod(method, asyncConvenienceMethod, true));
+                    operationMethods.Add(new ResourceOperationMethodProvider(this, method, convenienceMethod, false));
+                    var asyncConvenienceMethod = _restClientProvider.GetConvenienceMethodByOperation(method.Operation, true);
+                    operationMethods.Add(new ResourceOperationMethodProvider(this, method, asyncConvenienceMethod, true));
                 }
             }
 
@@ -280,20 +280,6 @@ protected override MethodProvider[] BuildMethods()
             ];
         }
 
-        protected MethodProvider BuildOperationMethod(InputServiceMethod method, MethodProvider convenienceMethod, bool isAsync)
-        {
-            return BuildOperationMethodCore(method, convenienceMethod, isAsync);
-        }
-
-        protected MethodProvider BuildOperationMethodCore(InputServiceMethod method, MethodProvider convenienceMethod, bool isAsync)
-        {
-            return new ResourceOperationMethodProvider(this, method, convenienceMethod, isAsync);
-        }
-
-        // TODO: get clean name of operation Name
-        protected MethodProvider GetCorrespondingConvenienceMethod(InputOperation operation, bool isAsync)
-            => _restClientProvider.CanonicalView.Methods.Single(m => m.Signature.Name.Equals(isAsync ? $"{operation.Name}Async" : operation.Name, StringComparison.OrdinalIgnoreCase) && m.Signature.Parameters.Any(p => p.Type.Equals(typeof(CancellationToken))));
-
         public ScopedApi<bool> TryGetApiVersion(out ScopedApi<string> apiVersion)
         {
             var apiVersionDeclaration = new VariableExpression(typeof(string), $"{SpecName.ToLower()}ApiVersion");
@@ -335,7 +321,7 @@ public ValueExpression[] PopulateArguments(
                     if (convenienceMethod != null)
                     {
                         var resource = convenienceMethod.Signature.Parameters
-                            .Single(p => p.Type.Equals(ResourceData.Type));
+                            .Single(p => p.Type.Equals(ResourceData.Type) || p.Type.Equals(typeof(RequestContent)));
                         arguments.Add(resource);
                     }
                     else

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/Providers/ResourceCollectionClientProvider.cs

@@ -151,7 +151,7 @@ private List<MethodProvider> BuildCreateOrUpdateMethods()
             foreach (var isAsync in new List<bool> { true, false})
             {
                 var convenienceMethod = _restClientProvider.GetConvenienceMethodByOperation(_create!.Operation, isAsync);
-                result.Add(BuildOperationMethod(_create, convenienceMethod, isAsync));
+                result.Add(new ResourceOperationMethodProvider(this, _create, convenienceMethod, isAsync));
             }
 
             return result;
@@ -174,7 +174,7 @@ private List<MethodProvider> BuildGetMethods()
             foreach (var isAsync in new List<bool> { true, false})
             {
                 var convenienceMethod = _restClientProvider.GetConvenienceMethodByOperation(_get!.Operation, isAsync);
-                result.Add(BuildOperationMethod(_get, convenienceMethod, isAsync));
+                result.Add(new ResourceOperationMethodProvider(this, _get, convenienceMethod, isAsync));
             }
 
             return result;

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/Providers/TagMethodProviders/BaseTagMethodProvider.cs

@@ -104,11 +104,10 @@ protected static List<MethodBodyStatement> CreateRequestContextAndProcessMessage
 
             statements.Add(ResourceMethodSnippets.CreateHttpMessage(resourceClientProvider, "CreateGetRequest", arguments, out var messageVariable));
 
-            var responseType = new CSharpType(typeof(Response<>), resourceClientProvider.ResourceData.Type);
             statements.AddRange(ResourceMethodSnippets.CreateGenericResponsePipelineProcessing(
                 messageVariable,
                 contextVariable,
-                responseType,
+                resourceClientProvider.ResourceData.Type,
                 isAsync,
                 out responseVariable));
 

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/Snippets/ResourceMethodSnippets.cs

@@ -89,7 +89,7 @@ public static MethodBodyStatement CreateHttpMessage(
         public static IReadOnlyList<MethodBodyStatement> CreateGenericResponsePipelineProcessing(
             VariableExpression messageVariable,
             VariableExpression contextVariable,
-            CSharpType responseType,
+            CSharpType responseGenericType,
             bool isAsync,
             out VariableExpression responseVariable)
         {
@@ -107,10 +107,10 @@ public static IReadOnlyList<MethodBodyStatement> CreateGenericResponsePipelinePr
             // Response<T> response = Response.FromValue((T)result, result);
             var responseDeclaration = Declare(
                 "response",
-                responseType,
+                new CSharpType(typeof(Response<>), responseGenericType),
                 Static(typeof(Response)).Invoke(
                     nameof(Response.FromValue),
-                    [resultVariable.CastTo(responseType.Arguments[0]), resultVariable]),
+                    [resultVariable.CastTo(responseGenericType), resultVariable]),
                 out responseVariable);
             statements.Add(responseDeclaration);
 

eng/packages/http-client-csharp-mgmt/generator/TestProjects/Local/Mgmt-TypeSpec/foo.tsp

@@ -40,5 +40,7 @@ interface Foos {
 
     delete is ArmResourceDeleteWithoutOkAsync<Foo>;
 
+    update is ArmCustomPatchAsync<Foo>;
+
     list is ArmResourceListByParent<Foo>;
 }