From 65f96e0298045f41f4c12ebe4ff3cb2218f1dedb Mon Sep 17 00:00:00 2001
From: Chunqing Chen <cchen@microsoft.com>
Date: Wed, 14 Jun 2017 06:44:06 +0800
Subject: [PATCH] Fix the issue that PS only stops transcription when all
 runspaces are closed (#3896)

Only stop transcription when all runspaces are closed
---
 .../engine/hostifaces/LocalConnection.cs      | 19 ++++++++----
 .../Start-Transcript.Tests.ps1                | 30 +++++++++++++++++++
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/src/System.Management.Automation/engine/hostifaces/LocalConnection.cs b/src/System.Management.Automation/engine/hostifaces/LocalConnection.cs
index 46333bb3d..34188c9af 100644
--- a/src/System.Management.Automation/engine/hostifaces/LocalConnection.cs
+++ b/src/System.Management.Automation/engine/hostifaces/LocalConnection.cs
@@ -863,12 +863,24 @@ namespace System.Management.Automation.Runspaces
                 }
 
                 if ((hostRunspace == null) || (this == hostRunspace))
-                {
+                {                   
+                    // We should close transcripting only if we are closing the last opened runspace.
+                    foreach (Runspace runspace in RunspaceList)
+                    {
+                        // At this stage, the last opened runspace should be at closing state.
+                        if (runspace.RunspaceStateInfo.State == RunspaceState.Opened)
+                        {
+                            return;
+                        }
+                    }
+
                     PSHostUserInterface host = executionContext.EngineHostInterface.UI;
                     if (host != null)
                     {
                         host.StopAllTranscribing();
                     }
+
+                    AmsiUtils.Uninitialize();
                 }
             }
 
@@ -913,10 +925,7 @@ namespace System.Management.Automation.Runspaces
 
             //Log engine lifecycle event.
             MshLog.LogEngineLifecycleEvent(_engine.Context, EngineState.Stopped);
-
-            // Uninitialize the AMSI scan interface
-            AmsiUtils.Uninitialize();
-
+  
             //All pipelines have been canceled. Close the runspace.
             _engine = null;
             _commandFactory = null;
diff --git a/test/powershell/Modules/Microsoft.Powershell.Host/Start-Transcript.Tests.ps1 b/test/powershell/Modules/Microsoft.Powershell.Host/Start-Transcript.Tests.ps1
index 4383c9c6e..cfcb7de91 100644
--- a/test/powershell/Modules/Microsoft.Powershell.Host/Start-Transcript.Tests.ps1
+++ b/test/powershell/Modules/Microsoft.Powershell.Host/Start-Transcript.Tests.ps1
@@ -106,4 +106,34 @@ Describe "Start-Transcript, Stop-Transcript tests" -tags "CI" {
         $expectedError = "CannotStartTranscription,Microsoft.PowerShell.Commands.StartTranscriptCommand"
         ValidateTranscription -scriptToExecute $script -outputFilePath $null -expectedError $expectedError
     }
+    It "Transcription should remain active if other runspace in the host get closed" {
+        try{
+            $ps = [powershell]::Create()
+            $ps.addscript("Start-Transcript -path $transcriptFilePath").Invoke()
+            $ps.addscript('$rs = [system.management.automation.runspaces.runspacefactory]::CreateRunspace()').Invoke()      
+            $ps.addscript('$rs.open()').Invoke()
+            $ps.addscript('$rs.Dispose()').Invoke()
+            $ps.addscript('Write-Host "After Dispose"').Invoke()
+            $ps.addscript("Stop-Transcript").Invoke()
+            } finally {
+                if ($ps -ne $null) {
+                    $ps.Dispose()
+                }
+            }
+
+        
+        Test-Path $transcriptFilePath | Should be $true
+        $transcriptFilePath | Should contain "After Dispose"
+    }
+
+    It "Transcription should be closed if the only runspace gets closed" {
+        $powerShellPath = [System.Diagnostics.Process]::GetCurrentProcess().Path
+        $powerShellCommand = $powerShellPath + ' "start-transcript $transcriptFilePath; Write-Host ''Before Dispose'';"'
+        Invoke-Expression $powerShellCommand
+        
+        Test-Path $transcriptFilePath | Should be $true
+        $transcriptFilePath | Should contain "Before Dispose"
+        $transcriptFilePath | Should contain "Windows PowerShell transcript end"
+    }
+
 }
\ No newline at end of file