Make sure we don't process code post message in parallel

Author: daxian-dbw

shell/AIShell.Integration/Channel.cs

@@ -27,6 +27,13 @@ public class Channel : IDisposable
     private Thread _serverThread;
     private CodePostData _pendingPostCodeData;
 
+    /// <summary>
+    /// Indicates if a code-post message is being processed:
+    ///  - value is 1: yes;
+    ///  - value is 0: no.
+    /// </summary>
+    private int _processingCodePost;
+
     private Channel(Runspace runspace, Type psConsoleReadLineType)
     {
         _runspace = runspace;
@@ -183,64 +190,79 @@ private void OnPostCode(PostCodeMessage postCodeMessage)
     {
         // Ignore 'code post' request when a posting operation is on-going.
         // This most likely would happen when user run 'code post' mutliple times to post the same code, which is safe to ignore.
-        if (_pendingPostCodeData is not null || postCodeMessage.CodeBlocks.Count is 0)
+        if (Interlocked.CompareExchange(ref _processingCodePost, 1, 0) is 1)
         {
             return;
         }
 
-        string codeToInsert;
-        List<string> codeBlocks = postCodeMessage.CodeBlocks;
-        List<PredictionCandidate> predictionCandidates = null;
-
-        if (codeBlocks.Count is 1)
-        {
-            codeToInsert = codeBlocks[0];
-        }
-        else if (Predictor.TryProcessForPrediction(codeBlocks, out predictionCandidates))
-        {
-            codeToInsert = predictionCandidates[0].Code;
-        }
-        else
+        try
         {
-            // Use LF as line ending to be consistent with the response from LLM.
-            StringBuilder sb = new(capacity: 50);
-            for (int i = 0; i < codeBlocks.Count; i++)
+            // If '_pendingPostCodeData' is already set, then we are still in an on-going posting operation,
+            // waiting the code to be posted on the 'OnIdle' event.
+            if (_pendingPostCodeData is not null || postCodeMessage.CodeBlocks.Count is 0)
+            {
+                return;
+            }
+
+            string codeToInsert;
+            List<string> codeBlocks = postCodeMessage.CodeBlocks;
+            List<PredictionCandidate> predictionCandidates = null;
+
+            if (codeBlocks.Count is 1)
+            {
+                codeToInsert = codeBlocks[0];
+            }
+            else if (Predictor.TryProcessForPrediction(codeBlocks, out predictionCandidates))
+            {
+                codeToInsert = predictionCandidates[0].Code;
+            }
+            else
             {
-                if (i > 0)
+                // Use LF as line ending to be consistent with the response from LLM.
+                StringBuilder sb = new(capacity: 50);
+                for (int i = 0; i < codeBlocks.Count; i++)
                 {
-                    sb.Append('\n');
+                    if (i > 0)
+                    {
+                        sb.Append('\n');
+                    }
+
+                    sb.Append(codeBlocks[i]).Append('\n');
                 }
 
-                sb.Append(codeBlocks[i]).Append('\n');
+                codeToInsert = sb.ToString();
             }
 
-            codeToInsert = sb.ToString();
-        }
-
-        // When PSReadLine is actively running, 'TreatControlCAsInput' would be set to 'true' because
-        // it handles 'Ctrl+c' as regular input.
-        // When the value is 'false', it means PowerShell is still busy running scripts or commands.
-        if (Console.TreatControlCAsInput)
-        {
-            PSRLRevertLine();
-            PSRLInsert(codeToInsert);
-            _predictor.SetCandidates(predictionCandidates);
+            // When PSReadLine is actively running, 'TreatControlCAsInput' would be set to 'true' because
+            // it handles 'Ctrl+c' as regular input.
+            // When the value is 'false', it means PowerShell is still busy running scripts or commands.
+            if (Console.TreatControlCAsInput)
+            {
+                PSRLRevertLine();
+                PSRLInsert(codeToInsert);
+                _predictor.SetCandidates(predictionCandidates);
+            }
+            else
+            {
+                _pendingPostCodeData = new CodePostData(codeToInsert, predictionCandidates);
+                // We use script block handler instead of a delegate handler because the latter will run
+                // in a background thread, while the former will run in the pipeline thread, which is way
+                // more predictable.
+                _runspace.Events.SubscribeEvent(
+                    source: null,
+                    eventName: null,
+                    sourceIdentifier: PSEngineEvent.OnIdle,
+                    data: null,
+                    action: _onIdleAction,
+                    supportEvent: true,
+                    forwardEvent: false,
+                    maxTriggerCount: 1);
+            }
         }
-        else
+        finally
         {
-            _pendingPostCodeData = new CodePostData(codeToInsert, predictionCandidates);
-            // We use script block handler instead of a delegate handler because the latter will run
-            // in a background thread, while the former will run in the pipeline thread, which is way
-            // more predictable.
-            _runspace.Events.SubscribeEvent(
-                source: null,
-                eventName: null,
-                sourceIdentifier: PSEngineEvent.OnIdle,
-                data: null,
-                action: _onIdleAction,
-                supportEvent: true,
-                forwardEvent: false,
-                maxTriggerCount: 1);
+            // Reset the state.
+            _processingCodePost = 0;
         }
     }