Fix razor redirect

src/LanguageServer/Microsoft.CodeAnalysis.LanguageServer/HostWorkspace/HostDiagnosticAnalyzerProvider.cs

@@ -3,8 +3,9 @@
 // See the LICENSE file in the project root for more information.
 
 using System.Collections.Immutable;
-using System.Reflection;
+using System.Diagnostics;
 using Microsoft.CodeAnalysis.Diagnostics;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Workspaces.ProjectSystem;
 
 namespace Microsoft.CodeAnalysis.LanguageServer.HostWorkspace;
@@ -17,8 +18,28 @@ internal sealed class HostDiagnosticAnalyzerProvider(string? razorSourceGenerato
     {
         if (File.Exists(razorSourceGenerator))
         {
-            return [(razorSourceGenerator, ProjectSystemProject.RazorVsixExtensionId)];
+            // we also have to redirect the utilities and object pool assemblies
+            var razorDir = Path.GetDirectoryName(razorSourceGenerator) ?? "";
+            var razorUtilities = GetDependency(razorDir, RazorAnalyzerAssemblyResolver.RazorUtilsAssemblyName);
+            var objectPool = GetDependency(razorDir, RazorAnalyzerAssemblyResolver.ObjectPoolAssemblyName);
+
+            return
+            [
+                (razorSourceGenerator, ProjectSystemProject.RazorVsixExtensionId),
+                (razorUtilities, ProjectSystemProject.RazorVsixExtensionId),
+                (objectPool, ProjectSystemProject.RazorVsixExtensionId)
+            ];
         }
         return [];
+
+        static string GetDependency(string razorDir, string dependencyName)
+        {
+            var dependency = Path.Combine(razorDir, dependencyName + ".dll");
+            if (!File.Exists(dependency))
+            {
+                throw new FileNotFoundException($"Could not find razor dependency {dependency}");
+            }
+            return dependency;
+        }
     }
 }

src/Tools/ExternalAccess/Razor/Features/RazorAnalyzerAssemblyResolver.cs

@@ -49,63 +49,87 @@ internal sealed class RazorAnalyzerAssemblyResolver() : IAnalyzerAssemblyResolve
                 return null;
             }
 
-            var assembly = compilerLoadContext.Assemblies.FirstOrDefault(a => a.GetName().Name == assemblyName.Name);
-            if (assembly is not null)
-            {
-                return assembly;
-            }
+            // https://github.com/dotnet/roslyn/issues/76868
+            // load the complete closure of razor assemblies if we're asked to load any of them. Subsequent requests for the others will just return the ones loaded here
+            LoadAssemblyByFileName(compilerLoadContext, RazorCompilerAssemblyName, directory);
+            LoadAssemblyByFileName(compilerLoadContext, RazorUtilsAssemblyName, directory);
+            LoadAssemblyByFileName(compilerLoadContext, ObjectPoolAssemblyName, directory);
 
-            var assemblyFileName = $"{assemblyName.Name}.dll";
+            // return the actual assembly that we were asked to load.
+            return LoadAssembly(compilerLoadContext, assemblyName, directory);
 
-            // Depending on who wins the race to load these assemblies, the base directory will either be the tooling root (if Roslyn wins)
-            // or the ServiceHubCore subfolder (razor). In the root directory these are netstandard2.0 targeted, in ServiceHubCore they are 
-            // .net targeted. We need to always pick the same set of assemblies regardless of who causes us to load. Because this code only
-            // runs in a .net based host, it's safe to always choose the .net targeted ServiceHubCore versions.
-            if (!Path.GetFileName(directory.AsSpan().TrimEnd(Path.DirectorySeparatorChar)).Equals(ServiceHubCoreFolderName, StringComparison.OrdinalIgnoreCase))
+            static Assembly? LoadAssemblyByFileName(AssemblyLoadContext compilerLoadContext, string fileName, string directory)
             {
-                var serviceHubCoreDirectory = Path.Combine(directory, ServiceHubCoreFolderName);
-
-                // The logic above only applies to VS. In VS Code there is no service hub, so appending the folder would be silly.
-                if (Directory.Exists(serviceHubCoreDirectory))
+                // This is kind of odd that we find the assembly on disk, read it to get its assemblyName, then load it by assemblyName,
+                // which in turn attempts to find it on the disk, but ensures we go through the correct loading logic later on.
+                var onDiskName = Path.Combine(directory, $"{fileName}.dll");
+                if (File.Exists(onDiskName))
                 {
-                    directory = serviceHubCoreDirectory;
+                    return LoadAssembly(compilerLoadContext, AssemblyName.GetAssemblyName(onDiskName), directory);
                 }
+                return null;
             }
 
-            var assemblyPath = Path.Combine(directory, assemblyFileName);
-            if (File.Exists(assemblyPath))
+            static Assembly? LoadAssembly(AssemblyLoadContext compilerLoadContext, AssemblyName assemblyName, string directory)
             {
-                // https://github.com/dotnet/roslyn/issues/76868
-                //
-                // There is a subtle race condition in this logic as another thread could load the assembly in between 
-                // the above calls and this one. Short term will just catch and grab the loaded assembly but longer 
-                // term need to think about creating a dedicated AssemblyLoadContext for the razor assemblies 
-                // which avoids this race condition.
-                try
+                var assembly = compilerLoadContext.Assemblies.FirstOrDefault(a => a.GetName().Name == assemblyName.Name);
+                if (assembly is not null)
                 {
-                    assembly = compilerLoadContext.LoadFromAssemblyPath(assemblyPath);
+                    return assembly;
                 }
-                catch
+
+                var assemblyFileName = $"{assemblyName.Name}.dll";
+
+                // Depending on who wins the race to load these assemblies, the base directory will either be the tooling root (if Roslyn wins)
+                // or the ServiceHubCore subfolder (razor). In the root directory these are netstandard2.0 targeted, in ServiceHubCore they are 
+                // .net targeted. We need to always pick the same set of assemblies regardless of who causes us to load. Because this code only
+                // runs in a .net based host, it's safe to always choose the .net targeted ServiceHubCore versions.
+                if (!Path.GetFileName(directory.AsSpan().TrimEnd(Path.DirectorySeparatorChar)).Equals(ServiceHubCoreFolderName, StringComparison.OrdinalIgnoreCase))
                 {
-                    assembly = compilerLoadContext.Assemblies.Single(a => a.GetName().Name == assemblyName.Name);
+                    var serviceHubCoreDirectory = Path.Combine(directory, ServiceHubCoreFolderName);
+
+                    // The logic above only applies to VS. In VS Code there is no service hub, so appending the folder would be silly.
+                    if (Directory.Exists(serviceHubCoreDirectory))
+                    {
+                        directory = serviceHubCoreDirectory;
+                    }
                 }
-            }
-            else
-            {
-                // There are assemblies in the razor sdk generator directory that do not exist in the VS installation. That
-                // means when the paths are redirected, it's possible that the assembly is not found. In that case, we should
-                // load the assembly from the VS installation by querying through the compiler context.
-                try
+
+                var assemblyPath = Path.Combine(directory, assemblyFileName);
+                if (File.Exists(assemblyPath))
                 {
-                    assembly = compilerLoadContext.LoadFromAssemblyName(assemblyName);
+                    // https://github.com/dotnet/roslyn/issues/76868
+                    //
+                    // There is a subtle race condition in this logic as another thread could load the assembly in between 
+                    // the above calls and this one. Short term will just catch and grab the loaded assembly but longer 
+                    // term need to think about creating a dedicated AssemblyLoadContext for the razor assemblies 
+                    // which avoids this race condition.
+                    try
+                    {
+                        assembly = compilerLoadContext.LoadFromAssemblyPath(assemblyPath);
+                    }
+                    catch
+                    {
+                        assembly = compilerLoadContext.Assemblies.Single(a => a.GetName().Name == assemblyName.Name);
+                    }
                 }
-                catch (FileNotFoundException)
+                else
                 {
-                    assembly = null;
+                    // There are assemblies in the razor sdk generator directory that do not exist in the VS installation. That
+                    // means when the paths are redirected, it's possible that the assembly is not found. In that case, we should
+                    // load the assembly from the VS installation by querying through the compiler context.
+                    try
+                    {
+                        assembly = compilerLoadContext.LoadFromAssemblyName(assemblyName);
+                    }
+                    catch (FileNotFoundException)
+                    {
+                        assembly = null;
+                    }
                 }
-            }
 
-            return assembly;
+                return assembly;
+            }
         }
     }
 }