Fix go to def for template spec modules

src/Bicep.LangServer.UnitTests/BicepExternalSourceRequestHandlerTests.cs

@@ -488,6 +488,26 @@ public void GetExternalSourceLinkUri_RequestedFilenameShouldBeBicepOrJson(Extern
             (decoded.RequestedFile ?? "main.json").Should().MatchRegex(".+\\.(bicep|json)$", "requested source file should end with .json or .bicep");
         }
 
+        [TestMethod]
+        public void GetTemplateSpecSourceLinkUri_WithTemplateSpecModuleReference_ReturnsEncodedUri()
+        {
+            var subscriptionId = Guid.NewGuid();
+            var resourceGroupName = "myRG";
+            var templateSpecName = "myTemplateSpec";
+            var version = "v1";
+            var referenceValue = $"{subscriptionId}/{resourceGroupName}/{templateSpecName}:{version}";
+
+            TemplateSpecModuleReference
+                .TryParse(null, referenceValue, BicepTestConstants.BuiltInConfigurationWithAllAnalyzersDisabled, new Uri("file:///no-parent-file-is-available.bicep"))
+                .IsSuccess(out var reference, out var errorBuilder)
+                .Should()
+                .BeTrue();
+
+            var result = BicepExternalSourceRequestHandler.GetTemplateSpeckSourceLinkUri(reference!);
+
+            result.Should().Be($"bicep-extsrc:ts%3A{subscriptionId}%2FmyRG%2FmyTemplateSpec%3Av1?ts%3A{subscriptionId}%2FmyRG%2FmyTemplateSpec%3Av1");
+        }
+
         private Uri GetExternalSourceLinkUri(ExternalSourceLinkTestData testData)
         {
             Uri? entrypointUri = testData.sourceEntrypoint is { } ? PathHelper.FilePathToFileUrl(testData.sourceEntrypoint) : null;
@@ -503,7 +523,7 @@ private Uri GetExternalSourceLinkUri(ExternalSourceLinkTestData testData)
                 new SourceArchiveBuilder().WithBicepFile(entrypointUri, "metadata description = 'bicep module'").Build()
                 : null;
 
-            return BicepExternalSourceRequestHandler.GetExternalSourceLinkUri(reference, sourceArchive);
+            return BicepExternalSourceRequestHandler.GetRegistryModuleSourceLinkUri(reference, sourceArchive);
         }
 
         private string TrimFirstCharacter(string s)

src/Bicep.LangServer/Handlers/BicepDefinitionHandler.cs

@@ -1,12 +1,14 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
+using System.Diagnostics;
 using Azure.Deployments.Core.Definitions.Schema;
 using Azure.Deployments.Core.Entities;
 using Azure.Deployments.Templates.Extensions;
 using Bicep.Core;
 using Bicep.Core.Emit;
 using Bicep.Core.Features;
 using Bicep.Core.FileSystem;
+using Bicep.Core.Modules;
 using Bicep.Core.Navigation;
 using Bicep.Core.Parsing;
 using Bicep.Core.Registry;
@@ -186,14 +188,24 @@ private LocationOrLocationLinks HandleModuleReference(CompilationContext context
 
         private Uri GetModuleSourceLinkUri(ISourceFile sourceFile, ArtifactReference reference)
         {
-            if (!this.CanClientAcceptRegistryContent() || !reference.IsExternal || reference is not OciArtifactReference ociReference)
+            if (!this.CanClientAcceptRegistryContent() || !reference.IsExternal)
             {
                 // the client doesn't support the bicep-extsrc scheme or we're dealing with a local module
                 // just use the file URI
                 return sourceFile.FileUri;
             }
 
-            return BicepExternalSourceRequestHandler.GetExternalSourceLinkUri(ociReference, moduleDispatcher?.TryGetModuleSources(reference).TryUnwrap());
+            if (reference is OciArtifactReference ociArtifactReference)
+            {
+                return BicepExternalSourceRequestHandler.GetRegistryModuleSourceLinkUri(ociArtifactReference, moduleDispatcher?.TryGetModuleSources(reference).TryUnwrap());
+            }
+
+            if (reference is TemplateSpecModuleReference templateSpecModuleReference)
+            {
+                return BicepExternalSourceRequestHandler.GetTemplateSpeckSourceLinkUri(templateSpecModuleReference);
+            }
+
+            throw new UnreachableException();
         }
 
         private LocationOrLocationLinks HandleWildcardImportDeclaration(CompilationContext context, DefinitionParams request, SymbolResolutionResult result, WildcardImportSymbol wildcardImport)

src/Bicep.LangServer/Handlers/BicepExternalSourceRequestHandler.cs

@@ -4,6 +4,7 @@
 using System.Diagnostics;
 using Bicep.Core.Diagnostics;
 using Bicep.Core.FileSystem;
+using Bicep.Core.Modules;
 using Bicep.Core.Registry;
 using Bicep.Core.Registry.Oci;
 using Bicep.Core.SourceCode;
@@ -113,11 +114,21 @@ public Task<BicepExternalSourceResponse> Handle(BicepExternalSourceParams reques
         /// <param name="reference">The module reference</param>
         /// <param name="sourceArchive">The source archive for the module, if sources are available</param>
         /// <returns>A bicep-extsrc: URI</returns>
-        public static Uri GetExternalSourceLinkUri(OciArtifactReference reference, SourceArchive? sourceArchive)
+        public static Uri GetRegistryModuleSourceLinkUri(OciArtifactReference reference, SourceArchive? sourceArchive)
         {
             return new ExternalSourceReference(reference, sourceArchive).ToUri();
         }
 
+        public static Uri GetTemplateSpeckSourceLinkUri(TemplateSpecModuleReference reference)
+        {
+            var uriBuilder = new UriBuilder($"{LangServerConstants.ExternalSourceFileScheme}:{Uri.EscapeDataString(reference.FullyQualifiedReference)}")
+            {
+                Query = Uri.EscapeDataString(reference.FullyQualifiedReference),
+            };
+
+            return uriBuilder.Uri;
+        }
+
         private BicepTelemetryEvent CreateSuccessTelemetry(SourceArchive? sourceArchive, string? requestedSourceFile)
         {
             return ExternalSourceRequestSuccess(

src/Bicep.LangServer/Handlers/ExternalSourceCodeLensProvider.cs

@@ -22,6 +22,11 @@ public static IEnumerable<CodeLens> GetCodeLenses(IModuleDispatcher moduleDispat
 
             if (request.TextDocument.Uri.Scheme == LangServerConstants.ExternalSourceFileScheme)
             {
+                if (request.TextDocument.Uri.Query.StartsWith("ts:", StringComparison.Ordinal))
+                {
+                    yield break;
+                }
+
                 string? message = null;
                 ExternalSourceReference? externalReference = null;
 

src/Bicep.LangServer/Handlers/ExternalSourceReference.cs

@@ -4,6 +4,7 @@
 using System.Diagnostics;
 using System.Text.RegularExpressions;
 using Bicep.Core.Diagnostics;
+using Bicep.Core.Modules;
 using Bicep.Core.Registry;
 using Bicep.Core.Registry.Oci;
 using Bicep.Core.SourceCode;