Implement PowerShell engine OnIdle events

Author: rjmholt

src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs

@@ -55,6 +55,8 @@ internal class PsesInternalHost : PSHost, IHostSupportsInteractiveSession, IRuns
 
         private readonly IdempotentLatch _isRunningLatch = new();
 
+        private EngineIntrinsics _mainRunspaceEngineIntrinsics;
+
         private bool _shouldExit = false;
 
         private string _localComputerName;
@@ -345,17 +347,18 @@ public Task SetInitialWorkingDirectoryAsync(string path, CancellationToken cance
 
         private void Run()
         {
-            (PowerShell pwsh, RunspaceInfo localRunspaceInfo) = CreateInitialPowerShellSession();
+            (PowerShell pwsh, RunspaceInfo localRunspaceInfo, EngineIntrinsics engineIntrinsics) = CreateInitialPowerShellSession();
+            _mainRunspaceEngineIntrinsics = engineIntrinsics;
             _localComputerName = localRunspaceInfo.SessionDetails.ComputerName;
             _runspaceStack.Push(new RunspaceFrame(pwsh.Runspace, localRunspaceInfo));
             PushPowerShellAndRunLoop(pwsh, PowerShellFrameType.Normal, localRunspaceInfo);
         }
 
-        private (PowerShell, RunspaceInfo) CreateInitialPowerShellSession()
+        private (PowerShell, RunspaceInfo, EngineIntrinsics) CreateInitialPowerShellSession()
         {
-            PowerShell pwsh = CreateInitialPowerShell(_hostInfo, _readLineProvider);
+            (PowerShell pwsh, EngineIntrinsics engineIntrinsics) = CreateInitialPowerShell(_hostInfo, _readLineProvider);
             RunspaceInfo localRunspaceInfo = RunspaceInfo.CreateFromLocalPowerShell(_logger, pwsh);
-            return (pwsh, localRunspaceInfo);
+            return (pwsh, localRunspaceInfo, engineIntrinsics);
         }
 
         private void PushPowerShellAndRunLoop(SMA.PowerShell pwsh, PowerShellFrameType frameType, RunspaceInfo newRunspaceInfo = null)
@@ -611,7 +614,7 @@ private static PowerShell CreatePowerShellForRunspace(Runspace runspace)
             return pwsh;
         }
 
-        public PowerShell CreateInitialPowerShell(
+        public (PowerShell, EngineIntrinsics) CreateInitialPowerShell(
             HostStartupInfo hostStartupInfo,
             ReadLineProvider readLineProvider)
         {
@@ -647,7 +650,7 @@ public PowerShell CreateInitialPowerShell(
                 }
             }
 
-            return pwsh;
+            return (pwsh, engineIntrinsics);
         }
 
         private Runspace CreateInitialRunspace(InitialSessionState initialSessionState)
@@ -666,7 +669,27 @@ private Runspace CreateInitialRunspace(InitialSessionState initialSessionState)
 
         private void OnPowerShellIdle()
         {
-            if (_taskQueue.Count == 0)
+            IReadOnlyList<PSEventSubscriber> eventSubscribers = _mainRunspaceEngineIntrinsics.Events.Subscribers;
+
+            // Go through pending event subscribers and:
+            // - if we have any subscribers, ensure we process any events
+            // - if we have any idle events, generate an idle event and process that
+            bool runPipelineForEventProcessing = false;
+            foreach (PSEventSubscriber subscriber in eventSubscribers)
+            {
+                runPipelineForEventProcessing = true;
+
+                if (string.Equals(subscriber.SourceIdentifier, PSEngineEvent.OnIdle, StringComparison.OrdinalIgnoreCase))
+                {
+                    // PowerShell thinks we're in a call (the ReadLine call) rather than idle,
+                    // but we know we're sitting in the prompt.
+                    // So we need to generate the idle event ourselves
+                    _mainRunspaceEngineIntrinsics.Events.GenerateEvent(PSEngineEvent.OnIdle, sender: null, args: null, extraData: null);
+                    break;
+                }
+            }
+
+            if (!runPipelineForEventProcessing && _taskQueue.Count == 0)
             {
                 return;
             }
@@ -686,9 +709,21 @@ private void OnPowerShellIdle()
                         return;
                     }
 
+                    // If we're executing a task, we don't need to run an extra pipeline later for events
+                    // TODO: This may not be a PowerShell task, so ideally we can differentiate that here.
+                    //       For now it's mostly true and an easy assumption to make.
+                    runPipelineForEventProcessing = false;
                     task.ExecuteSynchronously(cancellationScope.CancellationToken);
                 }
             }
+
+            // We didn't end up executinng anything in the background,
+            // so we need to run a small artificial pipeline instead
+            // to force event processing
+            if (runPipelineForEventProcessing)
+            {
+                InvokePSCommand(new PSCommand().AddScript("0", useLocalScope: true), PowerShellExecutionOptions.Default, CancellationToken.None);
+            }
         }
 
         private void OnCancelKeyPress(object sender, ConsoleCancelEventArgs args)
@@ -769,7 +804,8 @@ private Task PopOrReinitializeRunspaceAsync()
                         // If our main runspace was corrupted,
                         // we must re-initialize our state.
                         // TODO: Use runspace.ResetRunspaceState() here instead
-                        (PowerShell pwsh, RunspaceInfo runspaceInfo) = CreateInitialPowerShellSession();
+                        (PowerShell pwsh, RunspaceInfo runspaceInfo, EngineIntrinsics engineIntrinsics) = CreateInitialPowerShellSession();
+                        _mainRunspaceEngineIntrinsics = engineIntrinsics;
                         PushPowerShell(new PowerShellContextFrame(pwsh, runspaceInfo, PowerShellFrameType.Normal));
 
                         _logger.LogError($"Top level runspace entered state '{oldRunspaceState.State}' for reason '{oldRunspaceState.Reason}' and was reinitialized."