Merge pull request #1936 from PowerShell/andschwa/onidle

Fix regression around `OnIdle` handler

Author: andyleejordan

Diff for: src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs

@@ -503,18 +503,8 @@ private void Run()
                 (PowerShell pwsh, RunspaceInfo localRunspaceInfo, EngineIntrinsics engineIntrinsics) = CreateInitialPowerShellSession();
                 _mainRunspaceEngineIntrinsics = engineIntrinsics;
                 _localComputerName = localRunspaceInfo.SessionDetails.ComputerName;
-
-                // NOTE: In order to support running events registered to PowerShell's OnIdle
-                // handler, we have to have our top-level PowerShell instance be nested (otherwise
-                // we get a PSInvalidOperationException because pipelines cannot be run
-                // concurrently). Specifically this bug cropped up when a profile loaded code which
-                // registered (and subsequently ran) on the OnIdle handler since it was hitting the
-                // non-nested PowerShell instance. So now we just start with a nested instance.
-                // While the PowerShell object is nested, as a frame type, this is our top-level
-                // frame and therefore NOT nested in that sense.
-                PowerShell nestedPwsh = CreateNestedPowerShell(localRunspaceInfo);
-                _runspaceStack.Push(new RunspaceFrame(nestedPwsh.Runspace, localRunspaceInfo));
-                PushPowerShellAndRunLoop(nestedPwsh, PowerShellFrameType.Normal | PowerShellFrameType.Repl, localRunspaceInfo);
+                _runspaceStack.Push(new RunspaceFrame(pwsh.Runspace, localRunspaceInfo));
+                PushPowerShellAndRunLoop(pwsh, PowerShellFrameType.Normal | PowerShellFrameType.Repl, localRunspaceInfo);
             }
             catch (Exception e)
             {
@@ -1015,7 +1005,9 @@ private static PowerShell CreateNestedPowerShell(RunspaceInfo currentRunspace)
             // PowerShell.CreateNestedPowerShell() sets IsNested but not IsChild
             // This means it throws due to the parent pipeline not running...
             // So we must use the RunspaceMode.CurrentRunspace option on PowerShell.Create() instead
-            return PowerShell.Create(RunspaceMode.CurrentRunspace);
+            PowerShell pwsh = PowerShell.Create(RunspaceMode.CurrentRunspace);
+            pwsh.Runspace.ThreadOptions = PSThreadOptions.UseCurrentThread;
+            return pwsh;
         }
 
         private static PowerShell CreatePowerShellForRunspace(Runspace runspace)
@@ -1085,7 +1077,7 @@ private Runspace CreateInitialRunspace(InitialSessionState initialSessionState)
         }
 
         // NOTE: This token is received from PSReadLine, and it _is_ the ReadKey cancellation token!
-        private void OnPowerShellIdle(CancellationToken idleCancellationToken)
+        internal void OnPowerShellIdle(CancellationToken idleCancellationToken)
         {
             IReadOnlyList<PSEventSubscriber> eventSubscribers = _mainRunspaceEngineIntrinsics.Events.Subscribers;
 

Diff for: test/PowerShellEditorServices.Test.Shared/Profile/Test.PowerShellEditorServices_profile.ps1

@@ -5,3 +5,5 @@
 function Assert-ProfileLoaded {
 	return $true
 }
+
+Register-EngineEvent -SourceIdentifier PowerShell.OnIdle -MaxTriggerCount 1 -Action { $global:handledInProfile = $true }

Diff for: test/PowerShellEditorServices.Test/PsesHostFactory.cs

@@ -28,7 +28,7 @@ internal static class PsesHostFactory
 
         public static readonly string BundledModulePath = Path.GetFullPath(TestUtilities.NormalizePath("../../../../../module"));
 
-        public static PsesInternalHost Create(ILoggerFactory loggerFactory)
+        public static PsesInternalHost Create(ILoggerFactory loggerFactory, bool loadProfiles = false)
         {
             // We intentionally use `CreateDefault2()` as it loads `Microsoft.PowerShell.Core` only,
             // which is a more minimal and therefore safer state.
@@ -62,7 +62,7 @@ public static PsesInternalHost Create(ILoggerFactory loggerFactory)
             PsesInternalHost psesHost = new(loggerFactory, null, testHostDetails);
 
             // NOTE: Because this is used by constructors it can't use await.
-            if (psesHost.TryStartAsync(new HostStartOptions { LoadProfiles = false }, CancellationToken.None).GetAwaiter().GetResult())
+            if (psesHost.TryStartAsync(new HostStartOptions { LoadProfiles = loadProfiles }, CancellationToken.None).GetAwaiter().GetResult())
             {
                 return psesHost;
             }

Diff for: test/PowerShellEditorServices.Test/Session/PsesInternalHostTests.cs

@@ -111,39 +111,6 @@ public async Task CanCancelExecutionWithMethod()
             Assert.True(executeTask.IsCanceled);
         }
 
-        [Fact]
-        public async Task CanResolveAndLoadProfilesForHostId()
-        {
-            // Load the profiles for the test host name
-            await psesHost.LoadHostProfilesAsync(CancellationToken.None).ConfigureAwait(true);
-
-            // Ensure that the $PROFILE variable is a string with the value of CurrentUserCurrentHost.
-            IReadOnlyList<string> profileVariable = await psesHost.ExecutePSCommandAsync<string>(
-                new PSCommand().AddScript("$PROFILE"),
-                CancellationToken.None).ConfigureAwait(true);
-
-            Assert.Collection(profileVariable,
-                (p) => Assert.Equal(PsesHostFactory.TestProfilePaths.CurrentUserCurrentHost, p));
-
-            // Ensure that all the profile paths are set in the correct note properties.
-            IReadOnlyList<string> profileProperties = await psesHost.ExecutePSCommandAsync<string>(
-                new PSCommand().AddScript("$PROFILE | Get-Member -Type NoteProperty"),
-                CancellationToken.None).ConfigureAwait(true);
-
-            Assert.Collection(profileProperties,
-                (p) => Assert.Equal($"string AllUsersAllHosts={PsesHostFactory.TestProfilePaths.AllUsersAllHosts}", p, ignoreCase: true),
-                (p) => Assert.Equal($"string AllUsersCurrentHost={PsesHostFactory.TestProfilePaths.AllUsersCurrentHost}", p, ignoreCase: true),
-                (p) => Assert.Equal($"string CurrentUserAllHosts={PsesHostFactory.TestProfilePaths.CurrentUserAllHosts}", p, ignoreCase: true),
-                (p) => Assert.Equal($"string CurrentUserCurrentHost={PsesHostFactory.TestProfilePaths.CurrentUserCurrentHost}", p, ignoreCase: true));
-
-            // Ensure that the profile was loaded. The profile also checks that $PROFILE was defined.
-            IReadOnlyList<bool> profileLoaded = await psesHost.ExecutePSCommandAsync<bool>(
-                new PSCommand().AddScript("Assert-ProfileLoaded"),
-                CancellationToken.None).ConfigureAwait(true);
-
-            Assert.Collection(profileLoaded, Assert.True);
-        }
-
         [Fact]
         public async Task CanHandleNoProfiles()
         {
@@ -202,6 +169,35 @@ public async Task CanHandleUndefinedPrompt()
             Assert.Equal(PsesInternalHost.DefaultPrompt, prompt);
         }
 
+        [Fact]
+        public async Task CanRunOnIdleTask()
+        {
+            IReadOnlyList<PSObject> task = await psesHost.ExecutePSCommandAsync<PSObject>(
+                new PSCommand().AddScript("$handled = $false; Register-EngineEvent -SourceIdentifier PowerShell.OnIdle -MaxTriggerCount 1 -Action { $global:handled = $true }"),
+                CancellationToken.None).ConfigureAwait(true);
+
+            IReadOnlyList<bool> handled = await psesHost.ExecutePSCommandAsync<bool>(
+                new PSCommand().AddScript("$handled"),
+                CancellationToken.None).ConfigureAwait(true);
+
+            Assert.Collection(handled, (p) => Assert.False(p));
+
+            await psesHost.ExecuteDelegateAsync(
+                nameof(psesHost.OnPowerShellIdle),
+                executionOptions: null,
+                (_, _) => psesHost.OnPowerShellIdle(CancellationToken.None),
+                CancellationToken.None).ConfigureAwait(true);
+
+            // TODO: Why is this racy?
+            Thread.Sleep(2000);
+
+            handled = await psesHost.ExecutePSCommandAsync<bool>(
+                new PSCommand().AddScript("$handled"),
+                CancellationToken.None).ConfigureAwait(true);
+
+            Assert.Collection(handled, (p) => Assert.True(p));
+        }
+
         [Fact]
         public async Task CanLoadPSReadLine()
         {
@@ -240,4 +236,72 @@ public async Task CanHandleBadInitialWorkingDirectory(string path)
             Assert.Collection(getLocation, (d) => Assert.Equal(cwd, d, ignoreCase: true));
         }
     }
+
+    [Trait("Category", "PsesInternalHost")]
+    public class PsesInternalHostWithProfileTests : IDisposable
+    {
+        private readonly PsesInternalHost psesHost;
+
+        public PsesInternalHostWithProfileTests() => psesHost = PsesHostFactory.Create(NullLoggerFactory.Instance, loadProfiles: true);
+
+        public void Dispose()
+        {
+#pragma warning disable VSTHRD002
+            psesHost.StopAsync().Wait();
+#pragma warning restore VSTHRD002
+            GC.SuppressFinalize(this);
+        }
+
+        [Fact]
+        public async Task CanResolveAndLoadProfilesForHostId()
+        {
+            // Ensure that the $PROFILE variable is a string with the value of CurrentUserCurrentHost.
+            IReadOnlyList<string> profileVariable = await psesHost.ExecutePSCommandAsync<string>(
+                new PSCommand().AddScript("$PROFILE"),
+                CancellationToken.None).ConfigureAwait(true);
+
+            Assert.Collection(profileVariable,
+                (p) => Assert.Equal(PsesHostFactory.TestProfilePaths.CurrentUserCurrentHost, p));
+
+            // Ensure that all the profile paths are set in the correct note properties.
+            IReadOnlyList<string> profileProperties = await psesHost.ExecutePSCommandAsync<string>(
+                new PSCommand().AddScript("$PROFILE | Get-Member -Type NoteProperty"),
+                CancellationToken.None).ConfigureAwait(true);
+
+            Assert.Collection(profileProperties,
+                (p) => Assert.Equal($"string AllUsersAllHosts={PsesHostFactory.TestProfilePaths.AllUsersAllHosts}", p, ignoreCase: true),
+                (p) => Assert.Equal($"string AllUsersCurrentHost={PsesHostFactory.TestProfilePaths.AllUsersCurrentHost}", p, ignoreCase: true),
+                (p) => Assert.Equal($"string CurrentUserAllHosts={PsesHostFactory.TestProfilePaths.CurrentUserAllHosts}", p, ignoreCase: true),
+                (p) => Assert.Equal($"string CurrentUserCurrentHost={PsesHostFactory.TestProfilePaths.CurrentUserCurrentHost}", p, ignoreCase: true));
+
+            // Ensure that the profile was loaded. The profile also checks that $PROFILE was defined.
+            IReadOnlyList<bool> profileLoaded = await psesHost.ExecutePSCommandAsync<bool>(
+                new PSCommand().AddScript("Assert-ProfileLoaded"),
+                CancellationToken.None).ConfigureAwait(true);
+
+            Assert.Collection(profileLoaded, Assert.True);
+        }
+
+        // This test specifically relies on a handler registered in the test profile, and on the
+        // test host loading the profiles during startup, that way the pipeline timing is
+        // consistent.
+        [Fact]
+        public async Task CanRunOnIdleInProfileTask()
+        {
+            await psesHost.ExecuteDelegateAsync(
+                nameof(psesHost.OnPowerShellIdle),
+                executionOptions: null,
+                (_, _) => psesHost.OnPowerShellIdle(CancellationToken.None),
+                CancellationToken.None).ConfigureAwait(true);
+
+            // TODO: Why is this racy?
+            Thread.Sleep(2000);
+
+            IReadOnlyList<bool> handled = await psesHost.ExecutePSCommandAsync<bool>(
+                new PSCommand().AddScript("$handledInProfile"),
+                CancellationToken.None).ConfigureAwait(true);
+
+            Assert.Collection(handled, (p) => Assert.True(p));
+        }
+    }
 }