Support posting code from the sidecar AIShell to PowerShell with Invoke-AIShell -PostCode

shell/AIShell.Integration/AIShell.psd1

@@ -1,7 +1,7 @@
 @{
     RootModule = 'AIShell.psm1'
     NestedModules = @("AIShell.Integration.dll")
-    ModuleVersion = '1.0.3'
+    ModuleVersion = '1.0.4'
     GUID = 'ECB8BEE0-59B9-4DAE-9D7B-A990B480279A'
     Author = 'Microsoft Corporation'
     CompanyName = 'Microsoft Corporation'
@@ -13,5 +13,5 @@
     VariablesToExport = '*'
     AliasesToExport = @('aish', 'askai', 'fixit')
     HelpInfoURI = 'https://aka.ms/aishell-help'
-    PrivateData = @{ PSData = @{ Prerelease = 'preview3'; ProjectUri = 'https://github.com/PowerShell/AIShell' } }
+    PrivateData = @{ PSData = @{ Prerelease = 'preview4'; ProjectUri = 'https://github.com/PowerShell/AIShell' } }
 }

shell/AIShell.Integration/AIShell.psm1

@@ -1,3 +1,7 @@
+$module = Get-Module -Name PSReadLine
+if ($null -eq $module -or $module.Version -lt [version]"2.4.1") {
+    throw "The PSReadLine v2.4.1-beta1 or higher is required for the AIShell module to work properly."
+}
 
 ## Create the channel singleton when loading the module.
 $null = [AIShell.Integration.Channel]::CreateSingleton($host.Runspace, [Microsoft.PowerShell.PSConsoleReadLine])

shell/AIShell.Integration/Channel.cs

@@ -18,12 +18,14 @@ public class Channel : IDisposable
     private readonly MethodInfo _psrlInsert, _psrlRevertLine, _psrlAcceptLine;
     private readonly ManualResetEvent _connSetupWaitHandler;
     private readonly Predictor _predictor;
+    private readonly ScriptBlock _onIdleAction;
 
     private ShellClientPipe _clientPipe;
     private ShellServerPipe _serverPipe;
     private bool? _setupSuccess;
     private Exception _exception;
     private Thread _serverThread;
+    private CodePostData _pendingPostCodeData;
 
     private Channel(Runspace runspace, Type psConsoleReadLineType)
     {
@@ -44,6 +46,7 @@ private Channel(Runspace runspace, Type psConsoleReadLineType)
         _psrlAcceptLine = _psrlType.GetMethod("AcceptLine", bindingFlags);
 
         _predictor = new Predictor();
+        _onIdleAction = ScriptBlock.Create("[AIShell.Integration.Channel]::Singleton.OnIdleHandler()");
     }
 
     public static Channel CreateSingleton(Runspace runspace, Type psConsoleReadLineType)
@@ -165,9 +168,22 @@ private void ThrowIfNotConnected()
         }
     }
 
+    [Hidden()]
+    public void OnIdleHandler()
+    {
+        if (_pendingPostCodeData is not null)
+        {
+            PSRLInsert(_pendingPostCodeData.CodeToInsert);
+            _predictor.SetCandidates(_pendingPostCodeData.PredictionCandidates);
+            _pendingPostCodeData = null;
+        }
+    }
+
     private void OnPostCode(PostCodeMessage postCodeMessage)
     {
-        if (!Console.TreatControlCAsInput || postCodeMessage.CodeBlocks.Count is 0)
+        // 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)
         {
             return;
         }
@@ -201,12 +217,31 @@ private void OnPostCode(PostCodeMessage postCodeMessage)
             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);
         }
+        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);
+        }
     }
 
     private PostContextMessage OnAskContext(AskContextMessage askContextMessage)
@@ -247,6 +282,8 @@ private void PSRLAcceptLine()
     }
 }
 
+internal record CodePostData(string CodeToInsert, List<PredictionCandidate> PredictionCandidates);
+
 public class Init : IModuleAssemblyCleanup
 {
     public void OnRemove(PSModuleInfo psModuleInfo)

shell/AIShell.Integration/Commands/InvokeAishCommand.cs

@@ -8,19 +8,57 @@ namespace AIShell.Integration.Commands;
 [Cmdlet(VerbsLifecycle.Invoke, "AIShell", DefaultParameterSetName = "Default")]
 public class InvokeAIShellCommand : PSCmdlet
 {
-    [Parameter(Mandatory = true, ValueFromRemainingArguments = true)]
+    private const string DefaultSet = "Default";
+    private const string ClipboardSet = "Clipboard";
+    private const string PostCodeSet = "PostCode";
+    private const string CopyCodeSet = "CopyCode";
+    private const string ExitSet = "Exit";
+
+    /// <summary>
+    /// Sets and gets the query to be sent to AIShell
+    /// </summary>
+    [Parameter(Mandatory = true, ValueFromRemainingArguments = true, ParameterSetName = DefaultSet)]
+    [Parameter(Mandatory = true, ValueFromRemainingArguments = true, ParameterSetName = ClipboardSet)]
     public string[] Query { get; set; }
 
-    [Parameter]
+    /// <summary>
+    /// Sets and gets the agent to use for the query.
+    /// </summary>
+    [Parameter(ParameterSetName = DefaultSet)]
+    [Parameter(ParameterSetName = ClipboardSet)]
     [ValidateNotNullOrEmpty]
     public string Agent { get; set; }
 
-    [Parameter(ParameterSetName = "Default", Mandatory = false, ValueFromPipeline = true)]
+    /// <summary>
+    /// Sets and gets the context information for the query.
+    /// </summary>
+    [Parameter(ValueFromPipeline = true, ParameterSetName = DefaultSet)]
     public PSObject Context { get; set; }
 
-    [Parameter(ParameterSetName = "Clipboard", Mandatory = true)]
+    /// <summary>
+    /// Indicates getting context information from clipboard.
+    /// </summary>
+    [Parameter(Mandatory = true, ParameterSetName = ClipboardSet)]
     public SwitchParameter ContextFromClipboard { get; set; }
 
+    /// <summary>
+    /// Indicates running '/code post' from the AIShell.
+    /// </summary>
+    [Parameter(ParameterSetName = PostCodeSet)]
+    public SwitchParameter PostCode { get; set; }
+
+    /// <summary>
+    /// Indicates running '/code copy' from the AIShell.
+    /// </summary>
+    [Parameter(ParameterSetName = CopyCodeSet)]
+    public SwitchParameter CopyCode { get; set; }
+
+    /// <summary>
+    /// Indicates running '/exit' from the AIShell.
+    /// </summary>
+    [Parameter(ParameterSetName = ExitSet)]
+    public SwitchParameter Exit { get; set; }
+
     private List<PSObject> _contextObjects;
 
     protected override void ProcessRecord()
@@ -36,25 +74,43 @@ protected override void ProcessRecord()
 
     protected override void EndProcessing()
     {
-        Collection<string> results = null;
-        if (_contextObjects is not null)
-        {
-            using PowerShell pwsh = PowerShell.Create(RunspaceMode.CurrentRunspace);
-            results = pwsh
-                .AddCommand("Out-String")
-                .AddParameter("InputObject", _contextObjects)
-                .Invoke<string>();
-        }
-        else if (ContextFromClipboard)
+        string message, context = null;
+
+        switch (ParameterSetName)
         {
-            using PowerShell pwsh = PowerShell.Create(RunspaceMode.CurrentRunspace);
-            results = pwsh
-                .AddCommand("Get-Clipboard")
-                .AddParameter("Raw")
-                .Invoke<string>();
+            case PostCodeSet:
+                message = "/code post";
+                break;
+            case CopyCodeSet:
+                message = "/code copy";
+                break;
+            case ExitSet:
+                message = "/exit";
+                break;
+            default:
+                Collection<string> results = null;
+                if (_contextObjects is not null)
+                {
+                    using PowerShell pwsh = PowerShell.Create(RunspaceMode.CurrentRunspace);
+                    results = pwsh
+                        .AddCommand("Out-String")
+                        .AddParameter("InputObject", _contextObjects)
+                        .Invoke<string>();
+                }
+                else if (ContextFromClipboard)
+                {
+                    using PowerShell pwsh = PowerShell.Create(RunspaceMode.CurrentRunspace);
+                    results = pwsh
+                        .AddCommand("Get-Clipboard")
+                        .AddParameter("Raw")
+                        .Invoke<string>();
+                }
+
+                context = results?.Count > 0 ? results[0] : null;
+                message = string.Join(' ', Query);
+                break;
         }
 
-        string context = results?.Count > 0 ? results[0] : null;
-        Channel.Singleton.PostQuery(new PostQueryMessage(string.Join(' ', Query), context, Agent));
+        Channel.Singleton.PostQuery(new PostQueryMessage(message, context, Agent));
     }
 }