Convert PSRL OnIdle handler to take a CancellationToken

Author: rjmholt

src/PowerShellEditorServices/Services/PowerShell/Console/ConsoleReadLine.cs

@@ -52,7 +52,7 @@ public bool TryOverrideReadKey(Func<bool, ConsoleKeyInfo> readKeyFunc)
             return true;
         }
 
-        public bool TryOverrideIdleHandler(Action idleHandler)
+        public bool TryOverrideIdleHandler(Action<CancellationToken> idleHandler)
         {
             _psrlProxy.OverrideIdleHandler(idleHandler);
             return true;

src/PowerShellEditorServices/Services/PowerShell/Console/IReadLine.cs

@@ -15,6 +15,6 @@ internal interface IReadLine
 
         bool TryOverrideReadKey(Func<bool, ConsoleKeyInfo> readKeyOverride);
 
-        bool TryOverrideIdleHandler(Action idleHandler);
+        bool TryOverrideIdleHandler(Action<CancellationToken> idleHandler);
     }
 }

src/PowerShellEditorServices/Services/PowerShell/Console/PSReadLineProxy.cs

@@ -171,7 +171,7 @@ internal void OverrideReadKey(Func<bool, ConsoleKeyInfo> readKeyFunc)
             _readKeyOverrideField.SetValue(null, readKeyFunc);
         }
 
-        internal void OverrideIdleHandler(Action idleAction)
+        internal void OverrideIdleHandler(Action<CancellationToken> idleAction)
         {
             _handleIdleOverrideField.SetValue(null, idleAction);
         }

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

@@ -611,7 +611,7 @@ private static PowerShell CreatePowerShellForRunspace(Runspace runspace)
             return pwsh;
         }
 
-        public PowerShell CreateInitialPowerShell(
+        private PowerShell CreateInitialPowerShell(
             HostStartupInfo hostStartupInfo,
             ReadLineProvider readLineProvider)
         {
@@ -664,14 +664,14 @@ private Runspace CreateInitialRunspace(InitialSessionState initialSessionState)
             return runspace;
         }
 
-        private void OnPowerShellIdle()
+        private void OnPowerShellIdle(CancellationToken idleCancellationToken)
         {
             if (_taskQueue.Count == 0)
             {
                 return;
             }
 
-            using (CancellationScope cancellationScope = _cancellationContext.EnterScope(isIdleScope: true))
+            using (CancellationScope cancellationScope = _cancellationContext.EnterScope(isIdleScope: true, idleCancellationToken))
             {
                 while (!cancellationScope.CancellationToken.IsCancellationRequested
                     && _taskQueue.TryTake(out ISynchronousTask task))

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

@@ -30,15 +30,17 @@ public CancellationContext()
             _cancellationSourceStack = new ConcurrentStack<CancellationScope>();
         }
 
-        public CancellationScope EnterScope(bool isIdleScope)
+        public CancellationScope EnterScope(bool isIdleScope, CancellationToken cancellationToken)
         {
             CancellationTokenSource newScopeCancellationSource = _cancellationSourceStack.TryPeek(out CancellationScope parentScope)
-                ? CancellationTokenSource.CreateLinkedTokenSource(parentScope.CancellationToken)
-                : new CancellationTokenSource();
+                ? CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, parentScope.CancellationToken)
+                : CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
 
             return EnterScope(isIdleScope, newScopeCancellationSource);
         }
 
+        public CancellationScope EnterScope(bool isIdleScope) => EnterScope(isIdleScope, CancellationToken.None);
+
         public void CancelCurrentTask()
         {
             if (_cancellationSourceStack.TryPeek(out CancellationScope currentCancellationSource))