Fix the client diagnostics value for pageable methods (Azure#51961)

Co-authored-by: Copilot <[email protected]>

Author: ArcturusZhang

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

@@ -30,6 +30,10 @@ public class ManagementOutputLibrary : AzureOutputLibrary
         private ProviderConstantsProvider? _providerConstants;
         internal ProviderConstantsProvider ProviderConstants => _providerConstants ??= new ProviderConstantsProvider();
 
+        // TODO -- this is really a bad practice that this map is not built in one place, but we are building it while generating stuff and in the meantime we might read it.
+        // but currently this is the best we could do right now.
+        internal Dictionary<TypeProvider, string> PageableMethodScopes { get; } = new();
+
         private IReadOnlyList<ResourceClientProvider>? _resourceClients;
         private IReadOnlyList<ResourceCollectionClientProvider>? _resourceCollections;
         private IReadOnlyList<MockableResourceProvider>? _mockableResourceProviders;

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

@@ -102,7 +102,11 @@ protected MethodBodyStatement[] BuildBodyStatements()
         {
             var statements = new List<MethodBodyStatement>();
 
-            var collectionResultOfT = ((ScmMethodProvider)_convenienceMethod).CollectionDefinition!.Type;
+            var collectionResult = ((ScmMethodProvider)_convenienceMethod).CollectionDefinition!;
+            var diagnosticScope = ResourceHelpers.GetDiagnosticScope(_enclosingType, _methodName, _isAsync);
+            ManagementClientGenerator.Instance.OutputLibrary.PageableMethodScopes.Add(collectionResult, diagnosticScope);
+
+            var collectionResultOfT = collectionResult.Type;
             statements.Add(ResourceMethodSnippets.CreateRequestContext(KnownParameters.CancellationTokenParameter, out var contextVariable));
 
             var requestMethod = _restClientInfo.RestClientProvider.GetRequestMethodByOperation(_method.Operation);

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/Utilities/ResourceHelpers.cs

@@ -6,6 +6,7 @@
 using Azure.ResourceManager.Resources;
 using Microsoft.TypeSpec.Generator.Input.Extensions;
 using Microsoft.TypeSpec.Generator.Primitives;
+using Microsoft.TypeSpec.Generator.Providers;
 using System;
 using System.Runtime.CompilerServices;
 
@@ -60,6 +61,12 @@ public static string GetRestClientPropertyName(string clientName)
             };
         }
 
+        public static string GetDiagnosticScope(TypeProvider enclosingType, string methodName, bool isAsync)
+        {
+            var rawMethodName = isAsync && methodName.EndsWith("Async") ? methodName[..^5] : methodName; // trim "Async" if the method is async method
+            return $"{enclosingType.Type.Name}.{rawMethodName}";
+        }
+
         /// <summary>
         /// Determines if the given resource operation kind should also be treated as an LRO (Long-Running Operation).
         /// When this happens, an operation is not a true long-running operation at the REST API level,

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/Visitors/PaginationVisitor.cs

@@ -13,32 +13,122 @@ namespace Azure.Generator.Management.Visitors;
 
 internal class PaginationVisitor : ScmLibraryVisitor
 {
+    private bool _isCollectionResultType;
+
+    /// <summary>
+    /// This method is the entrance of the visitor.
+    /// We check the name of this type to ensure this visitor is only applied to the collection result types.
+    /// </summary>
+    /// <param name="type"></param>
+    /// <returns></returns>
+    protected override TypeProvider? VisitType(TypeProvider type)
+    {
+        _isCollectionResultType = type.Name.AsSpan().Contains("CollectionResult".AsSpan(), StringComparison.Ordinal);
+        return base.VisitType(type);
+    }
+
     protected override MethodBodyStatement? VisitStatements(MethodBodyStatements statements, MethodProvider method)
     {
-        if (method.EnclosingType.Name.AsSpan().Contains("CollectionResult".AsSpan(), StringComparison.Ordinal) && method.Signature.Name.Equals("AsPages"))
+        if (_isCollectionResultType && IsAsPagesMethod(method))
         {
-            var doWhileStatement = statements.OfType<DoWhileStatement>().FirstOrDefault();
-            if (doWhileStatement is not null)
-            {
-                var body = doWhileStatement.Body;
+            DoVisitAsPagesMethodStatements(statements, method);
+        }
+        return base.VisitStatements(statements, method);
+    }
+
+    private static bool IsAsPagesMethod(MethodProvider method) => method.Signature.Name.Equals("AsPages");
+    private static bool IsGetNextResponseMethod(MethodProvider method) => method.Signature.Name.Equals("GetNextResponse");
 
-                // get the response to model casting expression
-                var responseToModelStatement = body.OfType<ExpressionStatement>().Skip(1).FirstOrDefault();
-                if (responseToModelStatement is not null)
+    private void DoVisitAsPagesMethodStatements(MethodBodyStatements statements, MethodProvider method)
+    {
+        var doWhileStatement = statements.OfType<DoWhileStatement>().FirstOrDefault();
+        if (doWhileStatement is not null)
+        {
+            // we manually go over the body statements because currently the framework does not do this.
+            // TODO -- we do not have to do this once https://github.com/microsoft/typespec/issues/8177 is fixed.
+            foreach (var statement in doWhileStatement.Body)
+            {
+                if (statement is ExpressionStatement { Expression: AssignmentExpression assignment } expressionStatement)
                 {
-                    responseToModelStatement.Update(ConstructFromResponseExpression(responseToModelStatement));
+                    var updatedExpression = DoVisitAssignmentExpressionForAsPagesMethod(assignment, method);
+                    if (updatedExpression is not null)
+                    {
+                        // update the expression in the statement.
+                        expressionStatement.Update(updatedExpression);
+                    }
                 }
             }
         }
-        return base.VisitStatements(statements, method);
     }
 
-    private static AssignmentExpression ConstructFromResponseExpression(ExpressionStatement responseToModelStatement)
+    protected override ValueExpression? VisitAssignmentExpression(AssignmentExpression expression, MethodProvider method)
+    {
+        if (_isCollectionResultType && IsAsPagesMethod(method))
+        {
+            var newExpression = DoVisitAssignmentExpressionForAsPagesMethod(expression, method);
+            if (newExpression is not null)
+            {
+                // we have updated the expression, so we return it.
+                return newExpression;
+            }
+        }
+
+        if (_isCollectionResultType && IsGetNextResponseMethod(method))
+        {
+            var newExpression = DoVisitAssignmentExpressionForGetNextResponseMethod(expression, method);
+            if (newExpression is not null)
+            {
+                // we have updated the expression, so we return it.
+                return newExpression;
+            }
+        }
+        return base.VisitAssignmentExpression(expression, method);
+    }
+
+    private ValueExpression? DoVisitAssignmentExpressionForAsPagesMethod(AssignmentExpression expression, MethodProvider method)
     {
-        var assignmentExpression = responseToModelStatement.Expression as AssignmentExpression;
-        var castExpression = assignmentExpression?.Value as CastExpression;
-        var value = Static(castExpression?.Type!).Invoke(SerializationVisitor.FromResponseMethodName, [castExpression?.Inner!]);
-        var variable = assignmentExpression!.Variable;
-        return variable.Assign(value);
+        if (expression.Value is CastExpression castExpression && IsResponseToModelCastExpression(castExpression))
+        {
+            var value = Static(castExpression.Type!).Invoke(SerializationVisitor.FromResponseMethodName, [castExpression.Inner!]);
+            var variable = expression.Variable;
+            return variable.Assign(value);
+        }
+        // do nothing if nothing is changed.
+        return null;
+
+        static bool IsResponseToModelCastExpression(CastExpression castExpression)
+        {
+            if (castExpression.Inner is VariableExpression variableExpression &&
+                variableExpression.Type is { IsFrameworkType: true, FrameworkType: { } frameworkType } &&
+                frameworkType == typeof(Response))
+            {
+                return true;
+            }
+            return false;
+        }
+    }
+
+    private ValueExpression? DoVisitAssignmentExpressionForGetNextResponseMethod(AssignmentExpression expression, MethodProvider method)
+    {
+        if (expression is
+            {
+                Variable: DeclarationExpression { Variable: { } variable, IsUsing: true } declaration,
+                Value: InvokeMethodExpression invokeMethodExpression
+            }
+            && variable.Declaration.RequestedName == "scope")
+        {
+            // first we fetch the diagnostics scope from outputlibrary
+            if (ManagementClientGenerator.Instance.OutputLibrary.PageableMethodScopes.TryGetValue(method.EnclosingType, out var diagnosticsScopeValue))
+            {
+                // change its first argument to be the new diagnostics scope value.
+                invokeMethodExpression.Update(
+                    arguments: [Literal(diagnosticsScopeValue)]);
+
+                return expression;
+            }
+            return null;
+        }
+        // do nothing if nothing is changed.
+        return null;
     }
 }