Fix a lot of IntelliSense issues

src/PowerShellEditorServices/Server/PsesLanguageServer.cs

@@ -12,6 +12,7 @@
 using Microsoft.PowerShell.EditorServices.Services.PowerShell.Host;
 using Microsoft.PowerShell.EditorServices.Services.Template;
 using Newtonsoft.Json.Linq;
+using OmniSharp.Extensions.JsonRpc;
 using OmniSharp.Extensions.LanguageServer.Protocol.Models;
 using OmniSharp.Extensions.LanguageServer.Protocol.Server;
 using OmniSharp.Extensions.LanguageServer.Server;
@@ -100,7 +101,14 @@ public async Task StartAsync()
                     .WithHandler<PsesCodeLensHandlers>()
                     .WithHandler<PsesCodeActionHandler>()
                     .WithHandler<InvokeExtensionCommandHandler>()
-                    .WithHandler<PsesCompletionHandler>()
+                    // If PsesCompletionHandler is not marked as serial, then DidChangeTextDocument
+                    // notifications will end up cancelling completion. So quickly typing `Get-`
+                    // would result in no completions.
+                    //
+                    // This also lets completion requests interrupt time consuming background tasks
+                    // like the references code lens.
+                    .WithHandler<PsesCompletionHandler>(
+                        new JsonRpcHandlerOptions() { RequestProcessType = RequestProcessType.Serial })
                     .WithHandler<PsesHoverHandler>()
                     .WithHandler<PsesSignatureHelpHandler>()
                     .WithHandler<PsesDefinitionHandler>()

src/PowerShellEditorServices/Services/CodeLens/ICodeLensProvider.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.PowerShell.EditorServices.Services.TextDocument;
 using OmniSharp.Extensions.LanguageServer.Protocol.Models;
@@ -25,8 +26,9 @@ internal interface ICodeLensProvider
         /// <param name="scriptFile">
         /// The document for which CodeLenses should be provided.
         /// </param>
+        /// <param name="cancellationToken"></param>
         /// <returns>An array of CodeLenses.</returns>
-        CodeLens[] ProvideCodeLenses(ScriptFile scriptFile);
+        CodeLens[] ProvideCodeLenses(ScriptFile scriptFile, CancellationToken cancellationToken);
 
         /// <summary>
         /// Resolves a CodeLens that was created without a Command.
@@ -37,11 +39,13 @@ internal interface ICodeLensProvider
         /// <param name="scriptFile">
         /// The ScriptFile to resolve it in (sometimes unused).
         /// </param>
+        /// <param name="cancellationToken"></param>
         /// <returns>
         /// A Task which returns the resolved CodeLens when completed.
         /// </returns>
         Task<CodeLens> ResolveCodeLens(
             CodeLens codeLens,
-            ScriptFile scriptFile);
+            ScriptFile scriptFile,
+            CancellationToken cancellationToken);
     }
 }

src/PowerShellEditorServices/Services/CodeLens/PesterCodeLensProvider.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 using System;
 using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.PowerShell.EditorServices.Services;
 using Microsoft.PowerShell.EditorServices.Services.Symbols;
@@ -99,8 +100,9 @@ private static CodeLens[] GetPesterLens(PesterSymbolReference pesterSymbol, Scri
         /// Get all Pester CodeLenses for a given script file.
         /// </summary>
         /// <param name="scriptFile">The script file to get Pester CodeLenses for.</param>
+        /// <param name="cancellationToken"></param>
         /// <returns>All Pester CodeLenses for the given script file.</returns>
-        public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
+        public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile, CancellationToken cancellationToken)
         {
             // Don't return anything if codelens setting is disabled
             if (!_configurationService.CurrentSettings.Pester.CodeLens)
@@ -116,6 +118,7 @@ public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
                     continue;
                 }
 
+                cancellationToken.ThrowIfCancellationRequested();
                 if (_configurationService.CurrentSettings.Pester.UseLegacyCodeLens
                         && pesterSymbol.Command != PesterCommandType.Describe)
                 {
@@ -133,8 +136,9 @@ public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
         /// </summary>
         /// <param name="codeLens">The code lens to resolve.</param>
         /// <param name="scriptFile">The script file.</param>
+        /// <param name="cancellationToken"></param>
         /// <returns>The given CodeLens, wrapped in a task.</returns>
-        public Task<CodeLens> ResolveCodeLens(CodeLens codeLens, ScriptFile scriptFile) =>
+        public Task<CodeLens> ResolveCodeLens(CodeLens codeLens, ScriptFile scriptFile, CancellationToken cancellationToken) =>
             // This provider has no specific behavior for
             // resolving CodeLenses.
             Task.FromResult(codeLens);

src/PowerShellEditorServices/Services/CodeLens/ReferencesCodeLensProvider.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.PowerShell.EditorServices.Services;
 using Microsoft.PowerShell.EditorServices.Services.Symbols;
@@ -53,12 +54,14 @@ public ReferencesCodeLensProvider(WorkspaceService workspaceService, SymbolsServ
         /// Get all reference code lenses for a given script file.
         /// </summary>
         /// <param name="scriptFile">The PowerShell script file to get code lenses for.</param>
+        /// <param name="cancellationToken"></param>
         /// <returns>An array of CodeLenses describing all functions in the given script file.</returns>
-        public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
+        public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile, CancellationToken cancellationToken)
         {
             List<CodeLens> acc = new();
             foreach (SymbolReference sym in _symbolProvider.ProvideDocumentSymbols(scriptFile))
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 if (sym.SymbolType == SymbolType.Function)
                 {
                     acc.Add(new CodeLens
@@ -68,7 +71,7 @@ public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
                             Uri = scriptFile.DocumentUri,
                             ProviderId = nameof(ReferencesCodeLensProvider)
                         }, LspSerializer.Instance.JsonSerializer),
-                        Range = sym.ScriptRegion.ToRange()
+                        Range = sym.ScriptRegion.ToRange(),
                     });
                 }
             }
@@ -81,8 +84,12 @@ public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
         /// </summary>
         /// <param name="codeLens">The old code lens to get updated references for.</param>
         /// <param name="scriptFile"></param>
+        /// <param name="cancellationToken"></param>
         /// <returns>A new code lens object describing the same data as the old one but with updated references.</returns>
-        public async Task<CodeLens> ResolveCodeLens(CodeLens codeLens, ScriptFile scriptFile)
+        public async Task<CodeLens> ResolveCodeLens(
+            CodeLens codeLens,
+            ScriptFile scriptFile,
+            CancellationToken cancellationToken)
         {
             ScriptFile[] references = _workspaceService.ExpandScriptReferences(
                 scriptFile);
@@ -93,9 +100,10 @@ public async Task<CodeLens> ResolveCodeLens(CodeLens codeLens, ScriptFile script
                 codeLens.Range.Start.Character + 1);
 
             List<SymbolReference> referencesResult = await _symbolsService.FindReferencesOfSymbol(
-                foundSymbol,
-                references,
-                _workspaceService).ConfigureAwait(false);
+                    foundSymbol,
+                    references,
+                    _workspaceService,
+                    cancellationToken).ConfigureAwait(false);
 
             Location[] referenceLocations;
             if (referencesResult == null)
@@ -107,6 +115,10 @@ public async Task<CodeLens> ResolveCodeLens(CodeLens codeLens, ScriptFile script
                 List<Location> acc = new();
                 foreach (SymbolReference foundReference in referencesResult)
                 {
+                    // This async method is pretty dense with synchronous code
+                    // so it's helpful to add some yields.
+                    await Task.Yield();
+                    cancellationToken.ThrowIfCancellationRequested();
                     if (IsReferenceDefinition(foundSymbol, foundReference))
                     {
                         continue;
@@ -140,9 +152,9 @@ public async Task<CodeLens> ResolveCodeLens(CodeLens codeLens, ScriptFile script
                     Title = GetReferenceCountHeader(referenceLocations.Length),
                     Arguments = JArray.FromObject(new object[]
                     {
-                        scriptFile.DocumentUri,
-                        codeLens.Range.Start,
-                        referenceLocations
+                    scriptFile.DocumentUri,
+                    codeLens.Range.Start,
+                    referenceLocations
                     },
                     LspSerializer.Instance.JsonSerializer)
                 }

src/PowerShellEditorServices/Services/PowerShell/Execution/SynchronousTask.cs

@@ -35,7 +35,7 @@ protected SynchronousTask(
             CancellationToken cancellationToken)
         {
             Logger = logger;
-            _taskCompletionSource = new TaskCompletionSource<TResult>();
+            _taskCompletionSource = new TaskCompletionSource<TResult>(TaskCreationOptions.RunContinuationsAsynchronously);
             _taskRequesterCancellationToken = cancellationToken;
             _executionCanceled = false;
         }
@@ -76,6 +76,7 @@ public void ExecuteSynchronously(CancellationToken executorCancellationToken)
         {
             if (IsCanceled)
             {
+                SetCanceled();
                 return;
             }
 

src/PowerShellEditorServices/Services/PowerShell/Utility/CommandHelpers.cs

@@ -16,6 +16,10 @@ namespace Microsoft.PowerShell.EditorServices.Services.PowerShell.Utility
     /// </summary>
     internal static class CommandHelpers
     {
+        public record struct AliasMap(
+            Dictionary<string, List<string>> CmdletToAliases,
+            Dictionary<string, string> AliasToCmdlets);
+
         private static readonly HashSet<string> s_nounExclusionList = new()
         {
             // PowerShellGet v2 nouns
@@ -180,19 +184,21 @@ not CommandTypes.Function and
         /// Gets all aliases found in the runspace
         /// </summary>
         /// <param name="executionService"></param>
-        public static async Task<(Dictionary<string, List<string>>, Dictionary<string, string>)> GetAliasesAsync(IInternalPowerShellExecutionService executionService)
+        /// <param name="cancellationToken"></param>
+        public static async Task<AliasMap> GetAliasesAsync(
+            IInternalPowerShellExecutionService executionService,
+            CancellationToken cancellationToken = default)
         {
             Validate.IsNotNull(nameof(executionService), executionService);
 
-            IEnumerable<CommandInfo> aliases = await executionService.ExecuteDelegateAsync(
-                nameof(GetAliasesAsync),
-                executionOptions: null,
-                (pwsh, _) =>
-                {
-                    CommandInvocationIntrinsics invokeCommand = pwsh.Runspace.SessionStateProxy.InvokeCommand;
-                    return invokeCommand.GetCommands("*", CommandTypes.Alias, nameIsPattern: true);
-                },
-                CancellationToken.None).ConfigureAwait(false);
+            // Need to execute a PSCommand here as Runspace.SessionStateProxy cannot be used from
+            // our PSRL on idle handler.
+            IReadOnlyList<CommandInfo> aliases = await executionService.ExecutePSCommandAsync<CommandInfo>(
+                new PSCommand()
+                    .AddCommand("Microsoft.PowerShell.Core\\Get-Command")
+                    .AddParameter("ListImported", true)
+                    .AddParameter("CommandType", CommandTypes.Alias),
+                cancellationToken).ConfigureAwait(false);
 
             foreach (AliasInfo aliasInfo in aliases)
             {
@@ -206,7 +212,8 @@ not CommandTypes.Function and
                 s_aliasToCmdletCache.TryAdd(aliasInfo.Name, aliasInfo.Definition);
             }
 
-            return (new Dictionary<string, List<string>>(s_cmdletToAliasCache),
+            return new AliasMap(
+                new Dictionary<string, List<string>>(s_cmdletToAliasCache),
                 new Dictionary<string, string>(s_aliasToCmdletCache));
         }
     }

src/PowerShellEditorServices/Services/Symbols/SymbolsService.cs

@@ -11,6 +11,7 @@
 using System.Management.Automation.Language;
 using System.Runtime.InteropServices;
 using System.Text.RegularExpressions;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
 using Microsoft.PowerShell.EditorServices.CodeLenses;
@@ -160,18 +161,25 @@ public static SymbolReference FindSymbolAtLocation(
         /// <param name="foundSymbol">The symbol to find all references for</param>
         /// <param name="referencedFiles">An array of scriptFiles too search for references in</param>
         /// <param name="workspace">The workspace that will be searched for symbols</param>
+        /// <param name="cancellationToken"></param>
         /// <returns>FindReferencesResult</returns>
         public async Task<List<SymbolReference>> FindReferencesOfSymbol(
             SymbolReference foundSymbol,
             ScriptFile[] referencedFiles,
-            WorkspaceService workspace)
+            WorkspaceService workspace,
+            CancellationToken cancellationToken = default)
         {
             if (foundSymbol == null)
             {
                 return null;
             }
 
-            (Dictionary<string, List<string>> cmdletToAliases, Dictionary<string, string> aliasToCmdlets) = await CommandHelpers.GetAliasesAsync(_executionService).ConfigureAwait(false);
+            CommandHelpers.AliasMap aliases = await CommandHelpers.GetAliasesAsync(
+                _executionService,
+                cancellationToken).ConfigureAwait(false);
+
+            Dictionary<string, List<string>> cmdletToAliases = aliases.CmdletToAliases;
+            Dictionary<string, string> aliasToCmdlets = aliases.AliasToCmdlets;
 
             // We want to look for references first in referenced files, hence we use ordered dictionary
             // TODO: File system case-sensitivity is based on filesystem not OS, but OS is a much cheaper heuristic
@@ -188,6 +196,10 @@ public async Task<List<SymbolReference>> FindReferencesOfSymbol(
             {
                 if (!fileMap.Contains(filePath))
                 {
+                    // This async method is pretty dense with synchronous code
+                    // so it's helpful to add some yields.
+                    await Task.Yield();
+                    cancellationToken.ThrowIfCancellationRequested();
                     if (!workspace.TryGetFile(filePath, out ScriptFile scriptFile))
                     {
                         // If we can't access the file for some reason, just ignore it
@@ -223,6 +235,9 @@ public async Task<List<SymbolReference>> FindReferencesOfSymbol(
                     reference.FilePath = file.FilePath;
                     symbolReferences.Add(reference);
                 }
+
+                await Task.Yield();
+                cancellationToken.ThrowIfCancellationRequested();
             }
 
             return symbolReferences;

src/PowerShellEditorServices/Services/Symbols/Vistors/AstOperations.cs

@@ -14,6 +14,7 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.PowerShell.EditorServices.Services.PowerShell;
 using Microsoft.PowerShell.EditorServices.Services.PowerShell.Execution;
+using Microsoft.PowerShell.EditorServices.Services.PowerShell.Host;
 
 namespace Microsoft.PowerShell.EditorServices.Services.Symbols
 {
@@ -87,6 +88,34 @@ await executionService.ExecuteDelegateAsync(
                 (pwsh, _) =>
                 {
                     stopwatch.Start();
+
+                    // If the current runspace is not out of process, then we call TabExpansion2 so
+                    // that we have the ability to issue pipeline stop requests on cancellation.
+                    if (executionService is PsesInternalHost psesInternalHost
+                        && !psesInternalHost.Runspace.RunspaceIsRemote)
+                    {
+                        IReadOnlyList<CommandCompletion> completionResults = new SynchronousPowerShellTask<CommandCompletion>(
+                            logger,
+                            psesInternalHost,
+                            new PSCommand()
+                                .AddCommand("TabExpansion2")
+                                    .AddParameter("ast", scriptAst)
+                                    .AddParameter("tokens", currentTokens)
+                                    .AddParameter("positionOfCursor", cursorPosition),
+                            executionOptions: null,
+                            cancellationToken)
+                            .ExecuteAndGetResult(cancellationToken);
+
+                        if (completionResults is { Count: > 0 })
+                        {
+                            commandCompletion = completionResults[0];
+                        }
+
+                        return;
+                    }
+
+                    // If the current runspace is out of process, we can't call TabExpansion2
+                    // because the output will be serialized.
                     commandCompletion = CommandCompletion.CompleteInput(
                         scriptAst,
                         currentTokens,

src/PowerShellEditorServices/Services/TextDocument/Handlers/CodeLensHandlers.cs

@@ -41,11 +41,11 @@ public PsesCodeLensHandlers(ILoggerFactory factory, SymbolsService symbolsServic
         public override Task<CodeLensContainer> Handle(CodeLensParams request, CancellationToken cancellationToken)
         {
             ScriptFile scriptFile = _workspaceService.GetFile(request.TextDocument.Uri);
-            CodeLens[] codeLensResults = ProvideCodeLenses(scriptFile);
+            CodeLens[] codeLensResults = ProvideCodeLenses(scriptFile, cancellationToken);
             return Task.FromResult(new CodeLensContainer(codeLensResults));
         }
 
-        public override Task<CodeLens> Handle(CodeLens request, CancellationToken cancellationToken)
+        public override async Task<CodeLens> Handle(CodeLens request, CancellationToken cancellationToken)
         {
             // TODO: Catch deserialization exception on bad object
             CodeLensData codeLensData = request.Data.ToObject<CodeLensData>();
@@ -55,18 +55,19 @@ public override Task<CodeLens> Handle(CodeLens request, CancellationToken cancel
                 .FirstOrDefault(provider => provider.ProviderId.Equals(codeLensData.ProviderId, StringComparison.Ordinal));
 
             ScriptFile scriptFile = _workspaceService.GetFile(codeLensData.Uri);
-
-            return originalProvider.ResolveCodeLens(request, scriptFile);
+            return await originalProvider.ResolveCodeLens(request, scriptFile, cancellationToken)
+                .ConfigureAwait(false);
         }
 
         /// <summary>
         /// Get all the CodeLenses for a given script file.
         /// </summary>
         /// <param name="scriptFile">The PowerShell script file to get CodeLenses for.</param>
+        /// <param name="cancellationToken"></param>
         /// <returns>All generated CodeLenses for the given script file.</returns>
-        private CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
+        private CodeLens[] ProvideCodeLenses(ScriptFile scriptFile, CancellationToken cancellationToken)
         {
-            return InvokeProviders(provider => provider.ProvideCodeLenses(scriptFile))
+            return InvokeProviders(provider => provider.ProvideCodeLenses(scriptFile, cancellationToken))
                 .SelectMany(codeLens => codeLens)
                 .ToArray();
         }