From 9210fbdc687c7ff74f4815644cf8129cc7242634 Mon Sep 17 00:00:00 2001
From: John McPherson <johnmcp@microsoft.com>
Date: Tue, 23 Jul 2024 16:39:48 -0700
Subject: [PATCH] use ContinueWith rather than finally for coroutine result
 types

---
 src/PowerShell/CommonFiles/PowerShellCmdlet.cs | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/PowerShell/CommonFiles/PowerShellCmdlet.cs b/src/PowerShell/CommonFiles/PowerShellCmdlet.cs
index bd66db5ef7..1006e2cb58 100644
--- a/src/PowerShell/CommonFiles/PowerShellCmdlet.cs
+++ b/src/PowerShell/CommonFiles/PowerShellCmdlet.cs
@@ -97,11 +97,14 @@ internal Task RunOnMTA(Func<Task> func)
                 this.Write(StreamType.Verbose, "Already running on MTA");
                 try
                 {
-                    return func();
+                    Task result = func();
+                    result.ContinueWith((task) => this.Complete(), TaskContinuationOptions.ExecuteSynchronously);
+                    return result;
                 }
-                finally
+                catch
                 {
                     this.Complete();
+                    throw;
                 }
             }
 
@@ -150,11 +153,14 @@ internal Task<TResult> RunOnMTA<TResult>(Func<Task<TResult>> func)
                 this.Write(StreamType.Verbose, "Already running on MTA");
                 try
                 {
-                    return func();
+                    Task<TResult> result = func();
+                    result.ContinueWith((task) => this.Complete(), TaskContinuationOptions.ExecuteSynchronously);
+                    return result;
                 }
-                finally
+                catch
                 {
                     this.Complete();
+                    throw;
                 }
             }