Skip to content

Commit d9c5e06

Browse files
daxian-dbwdaxian-dbw
committed
Support posting code by Invoke-AIShell
1 parent 4d688ba commit d9c5e06

File tree

2 files changed

+115
-22
lines changed

2 files changed

+115
-22
lines changed

‎shell/AIShell.Integration/Channel.cs‎

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,14 @@ public class Channel : IDisposable
1818
private readonly MethodInfo _psrlInsert, _psrlRevertLine, _psrlAcceptLine;
1919
private readonly ManualResetEvent _connSetupWaitHandler;
2020
private readonly Predictor _predictor;
21+
private readonly ScriptBlock _onIdleAction;
2122

2223
private ShellClientPipe _clientPipe;
2324
private ShellServerPipe _serverPipe;
2425
private bool? _setupSuccess;
2526
private Exception _exception;
2627
private Thread _serverThread;
28+
private CodePostData _pendingPostCodeData;
2729

2830
private Channel(Runspace runspace, Type psConsoleReadLineType)
2931
{
@@ -44,6 +46,7 @@ private Channel(Runspace runspace, Type psConsoleReadLineType)
4446
_psrlAcceptLine = _psrlType.GetMethod("AcceptLine", bindingFlags);
4547

4648
_predictor = new Predictor();
49+
_onIdleAction = ScriptBlock.Create("[AIShell.Integration.Channel]::Singleton.OnIdleHandler()");
4750
}
4851

4952
public static Channel CreateSingleton(Runspace runspace, Type psConsoleReadLineType)
@@ -165,9 +168,22 @@ private void ThrowIfNotConnected()
165168
}
166169
}
167170

171+
[Hidden()]
172+
public void OnIdleHandler()
173+
{
174+
if (_pendingPostCodeData is not null)
175+
{
176+
PSRLInsert(_pendingPostCodeData.CodeToInsert);
177+
_predictor.SetCandidates(_pendingPostCodeData.PredictionCandidates);
178+
_pendingPostCodeData = null;
179+
}
180+
}
181+
168182
private void OnPostCode(PostCodeMessage postCodeMessage)
169183
{
170-
if (!Console.TreatControlCAsInput || postCodeMessage.CodeBlocks.Count is 0)
184+
// Ignore 'code post' request when a posting operation is on-going.
185+
// This most likely would happen when user run 'code post' mutliple times to post the same code, which is safe to ignore.
186+
if (_pendingPostCodeData is not null || postCodeMessage.CodeBlocks.Count is 0)
171187
{
172188
return;
173189
}
@@ -201,12 +217,31 @@ private void OnPostCode(PostCodeMessage postCodeMessage)
201217
codeToInsert = sb.ToString();
202218
}
203219

220+
// When PSReadLine is actively running, 'TreatControlCAsInput' would be set to 'true' because
221+
// it handles 'Ctrl+c' as regular input.
222+
// When the value is 'false', it means PowerShell is still busy running scripts or commands.
204223
if (Console.TreatControlCAsInput)
205224
{
206225
PSRLRevertLine();
207226
PSRLInsert(codeToInsert);
208227
_predictor.SetCandidates(predictionCandidates);
209228
}
229+
else
230+
{
231+
_pendingPostCodeData = new CodePostData(codeToInsert, predictionCandidates);
232+
// We use script block handler instead of a delegate handler because the latter will run
233+
// in a background thread, while the former will run in the pipeline thread, which is way
234+
// more predictable.
235+
_runspace.Events.SubscribeEvent(
236+
source: null,
237+
eventName: null,
238+
sourceIdentifier: PSEngineEvent.OnIdle,
239+
data: null,
240+
action: _onIdleAction,
241+
supportEvent: true,
242+
forwardEvent: false,
243+
maxTriggerCount: 1);
244+
}
210245
}
211246

212247
private PostContextMessage OnAskContext(AskContextMessage askContextMessage)
@@ -247,6 +282,8 @@ private void PSRLAcceptLine()
247282
}
248283
}
249284

285+
internal record CodePostData(string CodeToInsert, List<PredictionCandidate> PredictionCandidates);
286+
250287
public class Init : IModuleAssemblyCleanup
251288
{
252289
public void OnRemove(PSModuleInfo psModuleInfo)

‎shell/AIShell.Integration/Commands/InvokeAishCommand.cs‎

Lines changed: 77 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -8,19 +8,57 @@ namespace AIShell.Integration.Commands;
88
[Cmdlet(VerbsLifecycle.Invoke, "AIShell", DefaultParameterSetName = "Default")]
99
public class InvokeAIShellCommand : PSCmdlet
1010
{
11-
[Parameter(Mandatory = true, ValueFromRemainingArguments = true)]
11+
private const string DefaultSet = "Default";
12+
private const string ClipboardSet = "Clipboard";
13+
private const string PostCodeSet = "PostCode";
14+
private const string CopyCodeSet = "CopyCode";
15+
private const string ExitSet = "Exit";
16+
17+
/// <summary>
18+
/// Sets and gets the query to be sent to AIShell
19+
/// </summary>
20+
[Parameter(Mandatory = true, ValueFromRemainingArguments = true, ParameterSetName = DefaultSet)]
21+
[Parameter(Mandatory = true, ValueFromRemainingArguments = true, ParameterSetName = ClipboardSet)]
1222
public string[] Query { get; set; }
1323

14-
[Parameter]
24+
/// <summary>
25+
/// Sets and gets the agent to use for the query.
26+
/// </summary>
27+
[Parameter(ParameterSetName = DefaultSet)]
28+
[Parameter(ParameterSetName = ClipboardSet)]
1529
[ValidateNotNullOrEmpty]
1630
public string Agent { get; set; }
1731

18-
[Parameter(ParameterSetName = "Default", Mandatory = false, ValueFromPipeline = true)]
32+
/// <summary>
33+
/// Sets and gets the context information for the query.
34+
/// </summary>
35+
[Parameter(ValueFromPipeline = true, ParameterSetName = DefaultSet)]
1936
public PSObject Context { get; set; }
2037

21-
[Parameter(ParameterSetName = "Clipboard", Mandatory = true)]
38+
/// <summary>
39+
/// Indicates getting context information from clipboard.
40+
/// </summary>
41+
[Parameter(Mandatory = true, ParameterSetName = ClipboardSet)]
2242
public SwitchParameter ContextFromClipboard { get; set; }
2343

44+
/// <summary>
45+
/// Indicates running '/code post' from the AIShell.
46+
/// </summary>
47+
[Parameter(ParameterSetName = PostCodeSet)]
48+
public SwitchParameter PostCode { get; set; }
49+
50+
/// <summary>
51+
/// Indicates running '/code copy' from the AIShell.
52+
/// </summary>
53+
[Parameter(ParameterSetName = CopyCodeSet)]
54+
public SwitchParameter CopyCode { get; set; }
55+
56+
/// <summary>
57+
/// Indicates running '/exit' from the AIShell.
58+
/// </summary>
59+
[Parameter(ParameterSetName = ExitSet)]
60+
public SwitchParameter Exit { get; set; }
61+
2462
private List<PSObject> _contextObjects;
2563

2664
protected override void ProcessRecord()
@@ -36,25 +74,43 @@ protected override void ProcessRecord()
3674

3775
protected override void EndProcessing()
3876
{
39-
Collection<string> results = null;
40-
if (_contextObjects is not null)
41-
{
42-
using PowerShell pwsh = PowerShell.Create(RunspaceMode.CurrentRunspace);
43-
results = pwsh
44-
.AddCommand("Out-String")
45-
.AddParameter("InputObject", _contextObjects)
46-
.Invoke<string>();
47-
}
48-
else if (ContextFromClipboard)
77+
string message, context = null;
78+
79+
switch (ParameterSetName)
4980
{
50-
using PowerShell pwsh = PowerShell.Create(RunspaceMode.CurrentRunspace);
51-
results = pwsh
52-
.AddCommand("Get-Clipboard")
53-
.AddParameter("Raw")
54-
.Invoke<string>();
81+
case PostCodeSet:
82+
message = "/code post";
83+
break;
84+
case CopyCodeSet:
85+
message = "/code copy";
86+
break;
87+
case ExitSet:
88+
message = "/exit";
89+
break;
90+
default:
91+
Collection<string> results = null;
92+
if (_contextObjects is not null)
93+
{
94+
using PowerShell pwsh = PowerShell.Create(RunspaceMode.CurrentRunspace);
95+
results = pwsh
96+
.AddCommand("Out-String")
97+
.AddParameter("InputObject", _contextObjects)
98+
.Invoke<string>();
99+
}
100+
else if (ContextFromClipboard)
101+
{
102+
using PowerShell pwsh = PowerShell.Create(RunspaceMode.CurrentRunspace);
103+
results = pwsh
104+
.AddCommand("Get-Clipboard")
105+
.AddParameter("Raw")
106+
.Invoke<string>();
107+
}
108+
109+
context = results?.Count > 0 ? results[0] : null;
110+
message = string.Join(' ', Query);
111+
break;
55112
}
56113

57-
string context = results?.Count > 0 ? results[0] : null;
58-
Channel.Singleton.PostQuery(new PostQueryMessage(string.Join(' ', Query), context, Agent));
114+
Channel.Singleton.PostQuery(new PostQueryMessage(message, context, Agent));
59115
}
60116
}

0 commit comments

Comments
 (0)