From d9c5e06072d8f6e7f76e7b5b72f8daf35e59d5ce Mon Sep 17 00:00:00 2001 From: Dongbo Wang Date: Wed, 9 Oct 2024 15:18:58 -0700 Subject: [PATCH 1/2] Support posting code by Invoke-AIShell --- shell/AIShell.Integration/Channel.cs | 39 +++++++- .../Commands/InvokeAishCommand.cs | 98 +++++++++++++++---- 2 files changed, 115 insertions(+), 22 deletions(-) diff --git a/shell/AIShell.Integration/Channel.cs b/shell/AIShell.Integration/Channel.cs index afbf18f6..cdae0a4c 100644 --- a/shell/AIShell.Integration/Channel.cs +++ b/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 PredictionCandidates); + public class Init : IModuleAssemblyCleanup { public void OnRemove(PSModuleInfo psModuleInfo) diff --git a/shell/AIShell.Integration/Commands/InvokeAishCommand.cs b/shell/AIShell.Integration/Commands/InvokeAishCommand.cs index 845aa093..9c1d8f70 100644 --- a/shell/AIShell.Integration/Commands/InvokeAishCommand.cs +++ b/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"; + + /// + /// Sets and gets the query to be sent to AIShell + /// + [Parameter(Mandatory = true, ValueFromRemainingArguments = true, ParameterSetName = DefaultSet)] + [Parameter(Mandatory = true, ValueFromRemainingArguments = true, ParameterSetName = ClipboardSet)] public string[] Query { get; set; } - [Parameter] + /// + /// Sets and gets the agent to use for the query. + /// + [Parameter(ParameterSetName = DefaultSet)] + [Parameter(ParameterSetName = ClipboardSet)] [ValidateNotNullOrEmpty] public string Agent { get; set; } - [Parameter(ParameterSetName = "Default", Mandatory = false, ValueFromPipeline = true)] + /// + /// Sets and gets the context information for the query. + /// + [Parameter(ValueFromPipeline = true, ParameterSetName = DefaultSet)] public PSObject Context { get; set; } - [Parameter(ParameterSetName = "Clipboard", Mandatory = true)] + /// + /// Indicates getting context information from clipboard. + /// + [Parameter(Mandatory = true, ParameterSetName = ClipboardSet)] public SwitchParameter ContextFromClipboard { get; set; } + /// + /// Indicates running '/code post' from the AIShell. + /// + [Parameter(ParameterSetName = PostCodeSet)] + public SwitchParameter PostCode { get; set; } + + /// + /// Indicates running '/code copy' from the AIShell. + /// + [Parameter(ParameterSetName = CopyCodeSet)] + public SwitchParameter CopyCode { get; set; } + + /// + /// Indicates running '/exit' from the AIShell. + /// + [Parameter(ParameterSetName = ExitSet)] + public SwitchParameter Exit { get; set; } + private List _contextObjects; protected override void ProcessRecord() @@ -36,25 +74,43 @@ protected override void ProcessRecord() protected override void EndProcessing() { - Collection results = null; - if (_contextObjects is not null) - { - using PowerShell pwsh = PowerShell.Create(RunspaceMode.CurrentRunspace); - results = pwsh - .AddCommand("Out-String") - .AddParameter("InputObject", _contextObjects) - .Invoke(); - } - else if (ContextFromClipboard) + string message, context = null; + + switch (ParameterSetName) { - using PowerShell pwsh = PowerShell.Create(RunspaceMode.CurrentRunspace); - results = pwsh - .AddCommand("Get-Clipboard") - .AddParameter("Raw") - .Invoke(); + case PostCodeSet: + message = "/code post"; + break; + case CopyCodeSet: + message = "/code copy"; + break; + case ExitSet: + message = "/exit"; + break; + default: + Collection results = null; + if (_contextObjects is not null) + { + using PowerShell pwsh = PowerShell.Create(RunspaceMode.CurrentRunspace); + results = pwsh + .AddCommand("Out-String") + .AddParameter("InputObject", _contextObjects) + .Invoke(); + } + else if (ContextFromClipboard) + { + using PowerShell pwsh = PowerShell.Create(RunspaceMode.CurrentRunspace); + results = pwsh + .AddCommand("Get-Clipboard") + .AddParameter("Raw") + .Invoke(); + } + + 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)); } } From 94bb70759e2c9ca07c18b15f225bdac9e7c76709 Mon Sep 17 00:00:00 2001 From: Dongbo Wang Date: Tue, 1 Apr 2025 14:25:15 -0700 Subject: [PATCH 2/2] Update the module manifest and check the PSReadLine version --- shell/AIShell.Integration/AIShell.psd1 | 4 ++-- shell/AIShell.Integration/AIShell.psm1 | 4 ++++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/shell/AIShell.Integration/AIShell.psd1 b/shell/AIShell.Integration/AIShell.psd1 index d1419cde..c8dfb0d1 100644 --- a/shell/AIShell.Integration/AIShell.psd1 +++ b/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' } } } diff --git a/shell/AIShell.Integration/AIShell.psm1 b/shell/AIShell.Integration/AIShell.psm1 index e27916c7..bac0bc66 100644 --- a/shell/AIShell.Integration/AIShell.psm1 +++ b/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])