Prevent StreamingHost feedback loops in Receive-Job

help/Start-ThreadJob.md

@@ -171,6 +171,9 @@ Accept wildcard characters: False
 
 ### -StreamingHost
 Optional host object to which job output can directed.
+Progress, verbose, warning, and debug messages are streamed to this host instead of
+being buffered for `Receive-Job`. Success output, errors, and information records
+remain buffered in the job.
 
 ```yaml
 Type: PSHost

src/code/Microsoft.PowerShell.ThreadJob.cs

@@ -648,14 +648,7 @@ public ThreadJob(
                 iss.LanguageMode = enforceLockdown ? PSLanguageMode.ConstrainedLanguage : PSLanguageMode.FullLanguage;
             }
 
-            if (_streamingHost != null)
-            {
-                _rs = RunspaceFactory.CreateRunspace(_streamingHost, iss);
-            }
-            else
-            {
-                _rs = RunspaceFactory.CreateRunspace(iss);
-            }
+            _rs = RunspaceFactory.CreateRunspace(iss);
             _ps = System.Management.Automation.PowerShell.Create();
             _ps.Runspace = _rs;
             _ps.InvocationStateChanged += (sender, psStateChanged) =>
@@ -708,25 +701,26 @@ public ThreadJob(
             this.Error = _ps.Streams.Error;
             this.Error.EnumeratorNeverBlocks = true;
 
-            this.Progress = _ps.Streams.Progress;
-            this.Progress.EnumeratorNeverBlocks = true;
+            // Host-streamed records must not be replayed by Receive-Job.
+            if (_streamingHost == null)
+            {
+                this.Progress = _ps.Streams.Progress;
+                this.Verbose = _ps.Streams.Verbose;
+                this.Warning = _ps.Streams.Warning;
+                this.Debug = _ps.Streams.Debug;
+            }
 
-            this.Verbose = _ps.Streams.Verbose;
+            this.Information = _ps.Streams.Information;
+            this.Progress.EnumeratorNeverBlocks = true;
             this.Verbose.EnumeratorNeverBlocks = true;
-
-            this.Warning = _ps.Streams.Warning;
             this.Warning.EnumeratorNeverBlocks = true;
+            this.Debug.EnumeratorNeverBlocks = true;
+            this.Information.EnumeratorNeverBlocks = true;
             if (lockdownWarning != null)
             {
                 this.Warning.Add(lockdownWarning);
             }
 
-            this.Debug = _ps.Streams.Debug;
-            this.Debug.EnumeratorNeverBlocks = true;
-
-            this.Information = _ps.Streams.Information;
-            this.Information.EnumeratorNeverBlocks = true;
-
             // Create the JobManager job definition and job specification, and add to the JobManager.
             ThreadJobDefinition = new JobDefinition(typeof(ThreadJobSourceAdapter), "", Name);
             Dictionary<string, object> parameterCollection = new Dictionary<string, object>();
@@ -770,7 +764,8 @@ public override void StartJob()
                 _ps.Commands.Clear();
                 _ps.AddScript(_initSb.ToString());
                 _runningInitScript = true;
-                _ps.BeginInvoke<object, PSObject>(_input, _output);
+                var settings = _streamingHost == null ? null : new PSInvocationSettings { Host = _streamingHost };
+                _ps.BeginInvoke<object, PSObject>(_input, _output, settings, null, null);
             }
             else
             {
@@ -1078,7 +1073,8 @@ private void RunScript()
                 _ps.AddParameter(VERBATIM_ARGUMENT, _usingValuesMap);
             }
 
-            _ps.BeginInvoke<object, PSObject>(_input, _output);
+            var settings = _streamingHost == null ? null : new PSInvocationSettings { Host = _streamingHost };
+            _ps.BeginInvoke<object, PSObject>(_input, _output, settings, null, null);
         }
 
         private ScriptBlock GetScriptBlockFromFile(string filePath, PSCmdlet psCmdlet)

test/Microsoft.PowerShell.ThreadJob.Tests.ps1

@@ -23,6 +23,69 @@ function Wait-ForJobRunning
     }
 }
 
+Describe 'ThreadJob StreamingHost Tests' -Tags 'CI' {
+
+    It 'Streams <Stream> output once while Receive-Job waits (initialization: <Initialization>)' -TestCases @(
+        foreach ($initialization in $false, $true)
+        {
+            @{ Stream = 'Verbose'; Prefix = 'VERBOSE: '; Command = 'Write-Verbose "StreamingHostMessage" -Verbose'; Initialization = $initialization }
+            @{ Stream = 'Debug'; Prefix = 'DEBUG: '; Command = 'Write-Debug "StreamingHostMessage" -Debug'; Initialization = $initialization }
+            @{ Stream = 'Warning'; Prefix = 'WARNING: '; Command = 'Write-Warning "StreamingHostMessage"'; Initialization = $initialization }
+        }
+    ) {
+        param ($Stream, $Prefix, $Command, $Initialization)
+
+        if ($Stream -eq 'Debug' -and $PSVersionTable.PSVersion.Major -eq 5)
+        {
+            # Windows PowerShell's -Debug parameter prompts instead of continuing.
+            $Command = '$DebugPreference = "Continue"; Write-Debug "StreamingHostMessage"'
+        }
+
+        $modulePath = (Get-Command Start-ThreadJob).Module.Path.Replace("'", "''")
+        if ($Initialization)
+        {
+            $scriptParameters = "-InitializationScript { $Command; Start-Sleep -Milliseconds 300 } -ScriptBlock { 'SuccessOutput' }"
+        }
+        else
+        {
+            $scriptParameters = "-ScriptBlock { $Command; Start-Sleep -Milliseconds 300; 'SuccessOutput' }"
+        }
+
+        $script = @"
+if (`$null -ne `$PSStyle) { `$PSStyle.OutputRendering = 'PlainText' }
+Import-Module '$modulePath'
+Start-ThreadJob -StreamingHost `$Host $scriptParameters | Receive-Job -Wait -Verbose -Debug
+"@
+
+        $results = & (Get-Process -Id $PID).Path -NoProfile -NonInteractive -Command $script
+        $LASTEXITCODE | Should -Be 0
+        @($results | Where-Object { $_ -eq 'SuccessOutput' }) | Should -HaveCount 1
+        @($results | Where-Object { $_ -eq "${Prefix}StreamingHostMessage" }) |
+            Should -HaveCount 1
+    }
+
+    It 'Retains success output, errors, and silent information when using StreamingHost' {
+        $job = Start-ThreadJob -StreamingHost $Host -ScriptBlock {
+            'SuccessOutput'
+            Write-Error 'ErrorOutput'
+            Write-Information 'InformationOutput' -Tags 'TestTag'
+        } | Wait-Job
+
+        try
+        {
+            $job.Output | Should -Be 'SuccessOutput'
+            $job.Error | Should -Be 'ErrorOutput'
+            $job.Information | Should -HaveCount 1
+            $job.Information[0].MessageData | Should -Be 'InformationOutput'
+            $job.Information[0].Tags | Should -Contain 'TestTag'
+        }
+        finally
+        {
+            $job | Remove-Job -Force
+        }
+    }
+}
+
 # Basic thread job tests
 
 Describe 'Basic ThreadJob Tests' -Tags 'CI' {