[http-client-csharp] Fix partial method customization rebinding

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator.ClientModel/src/Providers/ScmMethodProviderCollection.cs

@@ -157,6 +157,8 @@ private ScmMethodProvider BuildConvenienceMethod(MethodProvider protocolMethod,
             // Detect a partial method declaration in the client's custom code matching this convenience method.
             MethodSignature? customSignature = null;
             PartialMethodCustomization.TryFindCustomSignature(client, methodName, signatureParameters, out customSignature);
+            var generatedModifiers = GetConvenienceMethodModifiers(protocolMethod.Signature.Modifiers, signatureParameters);
+            var generatedReturnType = GetResponseType(ServiceMethod.Operation.Responses, true, isAsync, out _);
 
             // Parameters used to construct the method body. When customizing, we clone the generator's
             // parameter providers with the customer's names but keep all generator metadata so that the
@@ -183,6 +185,7 @@ private ScmMethodProvider BuildConvenienceMethod(MethodProvider protocolMethod,
                 methodSignature = PartialMethodCustomization.BuildPartialSignature(
                     customSignature,
                     renamedSignatureParameters,
+                    generatedReturnType,
                     additionalModifiers: isAsync && _pagingServiceMethod == null
                         ? MethodSignatureModifiers.Async
                         : MethodSignatureModifiers.None);
@@ -193,8 +196,8 @@ private ScmMethodProvider BuildConvenienceMethod(MethodProvider protocolMethod,
                 methodSignature = new MethodSignature(
                     methodName,
                     DocHelpers.GetFormattableDescription(ServiceMethod.Operation.Summary, ServiceMethod.Operation.Doc) ?? FormattableStringHelpers.FromString(ServiceMethod.Operation.Name),
-                    GetConvenienceMethodModifiers(protocolMethod.Signature.Modifiers, signatureParameters),
-                    GetResponseType(ServiceMethod.Operation.Responses, true, isAsync, out _),
+                    generatedModifiers,
+                    generatedReturnType,
                     null,
                     signatureParameters,
                     Attributes: BuildConvenienceMethodAttributes());
@@ -1100,6 +1103,7 @@ private ScmMethodProvider BuildProtocolMethod(MethodProvider createRequestMethod
             // generated body using the customized parameter references.
             MethodSignature? customSignature = null;
             PartialMethodCustomization.TryFindCustomSignature(client, methodName, parameters, out customSignature);
+            var generatedReturnType = GetResponseType(ServiceMethod.Operation.Responses, false, isAsync, out _);
 
             MethodSignature methodSignature;
             ParameterProvider[] bodyParameters;
@@ -1115,6 +1119,7 @@ private ScmMethodProvider BuildProtocolMethod(MethodProvider createRequestMethod
                 methodSignature = PartialMethodCustomization.BuildPartialSignature(
                     customSignature,
                     requiredCustomParameters,
+                    generatedReturnType,
                     additionalModifiers: isAsync && _pagingServiceMethod == null
                         ? MethodSignatureModifiers.Async
                         : MethodSignatureModifiers.None);
@@ -1130,7 +1135,7 @@ private ScmMethodProvider BuildProtocolMethod(MethodProvider createRequestMethod
                     methodName,
                     DocHelpers.GetFormattableDescription(ServiceMethod.Operation.Summary, ServiceMethod.Operation.Doc) ?? FormattableStringHelpers.FromString(ServiceMethod.Operation.Name),
                     methodModifiers,
-                    GetResponseType(ServiceMethod.Operation.Responses, false, isAsync, out _),
+                    generatedReturnType,
                     $"The response returned from the service.",
                     parameters);
                 bodyParameters = parameters;

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/src/Providers/PartialMethodCustomization.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using Microsoft.TypeSpec.Generator.Expressions;
 using Microsoft.TypeSpec.Generator.Primitives;
 
 namespace Microsoft.TypeSpec.Generator.Providers
@@ -157,6 +158,34 @@ public static IReadOnlyList<ParameterProvider> RenameAndCloneParameters(
             return renamed;
         }
 
+        /// <summary>
+        /// Renames the generator parameters in place using the customer's partial declaration
+        /// parameter names while preserving the existing parameter instances.
+        /// </summary>
+        /// <remarks>
+        /// Use this when a method body has already been constructed with the generator
+        /// parameters. Body expressions may hold cached <see cref="Expressions.VariableExpression"/>
+        /// instances from those parameters, so replacing the parameters with clones would update
+        /// the signature but leave the body referencing the old declarations.
+        /// </remarks>
+        public static IReadOnlyList<ParameterProvider> RenameParametersInPlace(
+            IReadOnlyList<ParameterProvider> generatorParameters,
+            IReadOnlyList<ParameterProvider> customParameters,
+            bool removeDefaults)
+        {
+            for (int i = 0; i < generatorParameters.Count && i < customParameters.Count; i++)
+            {
+                var generatedParameter = generatorParameters[i];
+                generatedParameter.Update(name: customParameters[i].Name);
+                if (removeDefaults)
+                {
+                    generatedParameter.DefaultValue = null;
+                }
+            }
+
+            return generatorParameters;
+        }
+
         /// <summary>
         /// Builds a <see cref="MethodSignature"/> for a partial method implementation using
         /// <paramref name="customSignature"/> (modifiers, name, return type, attributes, and

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/src/Providers/TypeProvider.cs

@@ -435,15 +435,16 @@ private static MethodProvider CreatePartialMethodFromCustomSignature(MethodSigna
             // Partial method implementations require all parameters to be required (no default values).
             // The generator's parameters carry the metadata and the declarations referenced by the
             // method body and XML docs; the custom signature only supplies the parameter names.
-            var requiredParameters = PartialMethodCustomization.RenameAndCloneParameters(
+            var requiredParameters = PartialMethodCustomization.RenameParametersInPlace(
                 generatedMethod.Signature.Parameters,
                 customSignature.Parameters,
                 removeDefaults: true);
 
             var partialSignature = PartialMethodCustomization.BuildPartialSignature(
                 customSignature,
                 requiredParameters,
-                generatedMethod.Signature.ReturnType);
+                generatedMethod.Signature.ReturnType,
+                generatedMethod.Signature.Modifiers & MethodSignatureModifiers.Async);
 
             MethodProvider partialMethod = generatedMethod.BodyExpression != null
                 ? new MethodProvider(partialSignature, generatedMethod.BodyExpression, generatedMethod.EnclosingType, generatedMethod.XmlDocs, generatedMethod.Suppressions)

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/src/Writers/CodeWriter.cs

@@ -830,15 +830,15 @@ public IDisposable WriteMethodDeclarationNoScope(MethodSignatureBase methodBase,
                 .AppendRawIf("private ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Private))
                 .AppendRawIf("protected ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Protected))
                 .AppendRawIf("internal ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Internal))
-                .AppendRawIf("static ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Static))
-                .AppendRawIf("partial ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Partial));
+                .AppendRawIf("static ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Static));
 
             if (methodBase is MethodSignature method)
             {
                 AppendRawIf("virtual ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Virtual))
                     .AppendRawIf("abstract ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Abstract))
                     .AppendRawIf("override ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Override))
                     .AppendRawIf("new ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.New))
+                    .AppendRawIf("partial ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Partial))
                     .AppendRawIf("async ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Async));
 
                 var isImplicit = methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Implicit);
@@ -854,7 +854,6 @@ public IDisposable WriteMethodDeclarationNoScope(MethodSignatureBase methodBase,
                         AppendRaw("void ");
                     }
                 }
-
                 AppendRawIf("implicit ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Implicit))
                     .AppendRawIf("explicit ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Explicit))
                     .AppendRawIf("operator ", methodBase.Modifiers.HasFlag(MethodSignatureModifiers.Operator));

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/test/OutputLibraryVisitorTests.cs

@@ -272,6 +272,38 @@ [new ParameterProvider("param1", $"", typeof(int))]),
             Assert.AreEqual(typeof(float), typeProvider.Methods[0].Signature.Parameters[0].Type.FrameworkType);
         }
 
+        [Test]
+        public async Task PartialMethodCustomizationRenamesGeneratedBodyReferences()
+        {
+            var typeProvider = new MutableMethodsTypeProvider();
+            var parameter = new ParameterProvider("param1", $"", typeof(int), defaultValue: Snippet.Default);
+            var methodProvider = new MethodProvider(
+                new MethodSignature("TestMethod", $"", MethodSignatureModifiers.Public | MethodSignatureModifiers.Virtual, typeof(int), $"", [parameter]),
+                Snippet.Return(parameter),
+                typeProvider);
+            typeProvider.MethodProviders = [methodProvider];
+
+            var generator = await MockHelpers.LoadMockGeneratorAsync(
+                createOutputLibrary: () => new TestOutputLibrary(typeProvider),
+                compilation: async () => await Helpers.GetCompilationFromDirectoryAsync());
+
+            typeProvider.Update(reset: true);
+            foreach (var type in generator.Object.OutputLibrary.TypeProviders)
+            {
+                type.EnsureBuilt();
+            }
+
+            typeProvider.Update(methods: typeProvider.FilterCustomizedMethods(typeProvider.Methods));
+
+            var partialMethod = typeProvider.Methods.Single();
+            Assert.IsTrue(partialMethod.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Partial));
+            Assert.IsTrue(partialMethod.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Virtual));
+            Assert.AreEqual("renamed", partialMethod.Signature.Parameters[0].Name);
+            Assert.IsNull(partialMethod.Signature.Parameters[0].DefaultValue);
+            Assert.That(partialMethod.BodyStatements?.ToDisplayString(), Does.Contain("return renamed;"));
+            Assert.That(partialMethod.BodyStatements?.ToDisplayString(), Does.Not.Contain("param1"));
+        }
+
         [Test]
         public async Task MatchingConstructorSignatureIsFilteredAfterVisitorMutation()
         {

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/test/TestData/OutputLibraryVisitorTests/PartialMethodCustomizationRenamesGeneratedBodyReferences/TestName.cs

@@ -0,0 +1,7 @@
+namespace Test
+{
+    public partial class TestName
+    {
+        public virtual partial int TestMethod(int renamed = 0);
+    }
+}

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/test/Writers/CodeWriterTests.cs

@@ -681,6 +681,23 @@ public void CodeWriter_WriteMethodDeclaration_WithPartialModifier()
             Assert.AreEqual("private partial void DefineAdditionalProperties()", result);
         }
 
+        [Test]
+        public void CodeWriter_WriteMethodDeclaration_WithVirtualPartialModifier()
+        {
+            var methodSignature = new MethodSignature(
+                "CreateOrUpdate",
+                $"Creates or updates a resource.",
+                MethodSignatureModifiers.Public | MethodSignatureModifiers.Virtual | MethodSignatureModifiers.Partial,
+                typeof(int),
+                null,
+                []);
+            using var codeWriter = new CodeWriter();
+            codeWriter.WriteMethodDeclarationNoScope(methodSignature);
+
+            var result = codeWriter.ToString(false);
+            Assert.AreEqual("public virtual partial int CreateOrUpdate()", result);
+        }
+
         [Test]
         public void CodeWriter_WriteMethod_PartialMethodWithBody()
         {