From 59e311193546c3e4579cf523463e5f91faf7ba80 Mon Sep 17 00:00:00 2001
From: Christian Bankester <cbankester@immense.net>
Date: Tue, 19 Dec 2023 14:50:12 -0600
Subject: [PATCH] Fixed InvalidOperationException by wrapping async operations
 in a JoinableTaskContext

---
 Engine/Engine.csproj     |  1 +
 Engine/ScriptAnalyzer.cs | 23 +++++++++++++----------
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/Engine/Engine.csproj b/Engine/Engine.csproj
index ac8f389e7..f38a4f4d7 100644
--- a/Engine/Engine.csproj
+++ b/Engine/Engine.csproj
@@ -46,6 +46,7 @@
 
   <ItemGroup>
     <PackageReference Include="Microsoft.CSharp" Version="4.7.0" />
+    <PackageReference Include="Microsoft.VisualStudio.Threading" Version="17.8.14" />
   </ItemGroup>
 
   <ItemGroup>
diff --git a/Engine/ScriptAnalyzer.cs b/Engine/ScriptAnalyzer.cs
index 1a885eabe..da324dfe5 100644
--- a/Engine/ScriptAnalyzer.cs
+++ b/Engine/ScriptAnalyzer.cs
@@ -22,6 +22,7 @@
 using System.Collections;
 using System.Diagnostics;
 using System.Text;
+using Microsoft.VisualStudio.Threading;
 using Microsoft.Windows.PowerShell.ScriptAnalyzer.Extensions;
 
 namespace Microsoft.Windows.PowerShell.ScriptAnalyzer
@@ -2172,18 +2173,20 @@ public IEnumerable<DiagnosticRecord> AnalyzeSyntaxTree(
                             result.Add(new ErrorRecord(scriptRuleException, Strings.RuleErrorMessage, ErrorCategory.InvalidOperation, scriptAst.Extent.File));
                         }
 
-                        verboseOrErrors.Add(result);
-                    }));
-                    Task.Factory.ContinueWhenAll(tasks.ToArray(), t => verboseOrErrors.CompleteAdding());
-                    while (!verboseOrErrors.IsCompleted)
+                        return result;
+                    }).WithTimeout(TimeSpan.FromSeconds(10))).ToList();
+                    while (tasks.Count > 0)
                     {
-                        List<object> data = null;
-                        try
+                        var data = new JoinableTaskContext().Factory.Run(async () =>
                         {
-                            data = verboseOrErrors.Take();
-                        }
-                        catch (InvalidOperationException) { }
-
+                            var task = await Task.WhenAny(tasks.ToArray());
+                            tasks.Remove(task);
+                            if (task.IsCanceled || task.IsFaulted)
+                            {
+                                return null;
+                            }
+                            return await task;
+                        });
                         if (data != null)
                         {
                             this.outputWriter.WriteVerbose(data[0] as string);