Skip to content

Commit 221dcde

Browse files
authored
Fix code posting on macOS (#366)
This PR fixes the following scenarios on macOS 1. Running `/code post` from the `AIShell` pane to post code to the `PowerShell` pane 2. Running `Invoke-AIShell -PostCode` from the `PowerShell` to request code to be posted from the `AIShell` pane. **[NOTE]** Now the `AIShell` module depends on the **v2.4.2-beta2** PSReadLine, because a private field was added in PSReadLine to accurately indicate if PSReadLine is initialized and ready to render (PowerShell/PSReadLine#4706). The `AIShell` module starts to depend on that field to be more deterministic about if PSReadLine is running and ready. Before this change, we relied on `Console.TreatControlCAsInput` which is not accurate and may cause race condition: `Console.TreatControlCAsInput` is set to true, but `Initialize(...)` is not called yet. Rendering the posted code before initialization could result in corrupted state -- posted code may be rendered on wrong position or even cause exception.
1 parent a2b23fe commit 221dcde

File tree

2 files changed

+57
-9
lines changed

2 files changed

+57
-9
lines changed

shell/AIShell.Integration/AIShell.psm1

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
$module = Get-Module -Name PSReadLine
2-
if ($null -eq $module -or $module.Version -lt [version]"2.4.1") {
3-
throw "The PSReadLine v2.4.1-beta1 or higher is required for the AIShell module to work properly."
2+
if ($null -eq $module -or $module.Version -lt [version]"2.4.2") {
3+
throw "The PSReadLine v2.4.2-beta2 or higher is required for the AIShell module to work properly."
44
}
55

66
## Create the channel singleton when loading the module.

shell/AIShell.Integration/Channel.cs

Lines changed: 55 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ public class Channel : IDisposable
1616
private readonly Type _psrlType;
1717
private readonly Runspace _runspace;
1818
private readonly MethodInfo _psrlInsert, _psrlRevertLine, _psrlAcceptLine;
19+
private readonly FieldInfo _psrlHandleResizing, _psrlReadLineReady;
20+
private readonly object _psrlSingleton;
1921
private readonly ManualResetEvent _connSetupWaitHandler;
2022
private readonly Predictor _predictor;
2123
private readonly ScriptBlock _onIdleAction;
@@ -40,10 +42,17 @@ private Channel(Runspace runspace, Type psConsoleReadLineType)
4042
.Append(Path.GetFileNameWithoutExtension(Environment.ProcessPath))
4143
.ToString();
4244

43-
BindingFlags bindingFlags = BindingFlags.Static | BindingFlags.Public;
44-
_psrlInsert = _psrlType.GetMethod("Insert", bindingFlags, [typeof(string)]);
45-
_psrlRevertLine = _psrlType.GetMethod("RevertLine", bindingFlags);
46-
_psrlAcceptLine = _psrlType.GetMethod("AcceptLine", bindingFlags);
45+
BindingFlags methodFlags = BindingFlags.Static | BindingFlags.Public;
46+
_psrlInsert = _psrlType.GetMethod("Insert", methodFlags, [typeof(string)]);
47+
_psrlRevertLine = _psrlType.GetMethod("RevertLine", methodFlags);
48+
_psrlAcceptLine = _psrlType.GetMethod("AcceptLine", methodFlags);
49+
50+
FieldInfo singletonInfo = _psrlType.GetField("_singleton", BindingFlags.Static | BindingFlags.NonPublic);
51+
_psrlSingleton = singletonInfo.GetValue(null);
52+
53+
BindingFlags fieldFlags = BindingFlags.Instance | BindingFlags.NonPublic;
54+
_psrlReadLineReady = _psrlType.GetField("_readLineReady", fieldFlags);
55+
_psrlHandleResizing = _psrlType.GetField("_handlePotentialResizing", fieldFlags);
4756

4857
_predictor = new Predictor();
4958
_onIdleAction = ScriptBlock.Create("[AIShell.Integration.Channel]::Singleton.OnIdleHandler()");
@@ -217,10 +226,9 @@ private void OnPostCode(PostCodeMessage postCodeMessage)
217226
codeToInsert = sb.ToString();
218227
}
219228

220-
// When PSReadLine is actively running, 'TreatControlCAsInput' would be set to 'true' because
221-
// it handles 'Ctrl+c' as regular input.
229+
// When PSReadLine is actively running, its '_readLineReady' field should be set to 'true'.
222230
// When the value is 'false', it means PowerShell is still busy running scripts or commands.
223-
if (Console.TreatControlCAsInput)
231+
if (_psrlReadLineReady.GetValue(_psrlSingleton) is true)
224232
{
225233
PSRLRevertLine();
226234
PSRLInsert(codeToInsert);
@@ -268,18 +276,58 @@ private void OnAskConnection(ShellClientPipe clientPipe, Exception exception)
268276

269277
private void PSRLInsert(string text)
270278
{
279+
using var _ = new NoWindowResizingCheck();
271280
_psrlInsert.Invoke(null, [text]);
272281
}
273282

274283
private void PSRLRevertLine()
275284
{
285+
using var _ = new NoWindowResizingCheck();
276286
_psrlRevertLine.Invoke(null, [null, null]);
277287
}
278288

279289
private void PSRLAcceptLine()
280290
{
291+
using var _ = new NoWindowResizingCheck();
281292
_psrlAcceptLine.Invoke(null, [null, null]);
282293
}
294+
295+
/// <summary>
296+
/// We assume the terminal window will not resize during the code-post operation and hence disable the window resizing check on macOS.
297+
/// This is to avoid reading console cursor positions while PSReadLine is already blocked on 'Console.ReadKey', because on Unix system,
298+
/// when we are already blocked on key input, reading cursor position on another thread will be blocked too until a key is pressed.
299+
///
300+
/// We do need window resizing check on Windows due to how 'Start-AIShell' works differently:
301+
/// - On Windows, 'Start-AIShell' returns way BEFORE the current tab gets splitted for the sidecar pane, and PowerShell has already
302+
/// called into PSReadLine when the splitting actually happens. So, it's literally a window resizing for PSReadLine at that point
303+
/// and hence we need the window resizing check to correct the initial coordinates ('_initialX' and '_initialY').
304+
/// - On macOS, however, 'Start-AIShell' returns AFTER the current tab gets splitted for the sidecar pane. So, window resizing will
305+
/// be done before PowerShell calls into PSReadLine and hence there is no need for window resizing check on macOS.
306+
/// Also, On Windows we can read cursor position without blocking even if another thread is blocked on calling 'ReadKey'.
307+
/// </summary>
308+
private class NoWindowResizingCheck : IDisposable
309+
{
310+
private readonly object _originalValue;
311+
312+
internal NoWindowResizingCheck()
313+
{
314+
if (OperatingSystem.IsMacOS())
315+
{
316+
Channel channel = Singleton;
317+
_originalValue = channel._psrlHandleResizing.GetValue(channel._psrlSingleton);
318+
channel._psrlHandleResizing.SetValue(channel._psrlSingleton, false);
319+
}
320+
}
321+
322+
public void Dispose()
323+
{
324+
if (OperatingSystem.IsMacOS())
325+
{
326+
Channel channel = Singleton;
327+
channel._psrlHandleResizing.SetValue(channel._psrlSingleton, _originalValue);
328+
}
329+
}
330+
}
283331
}
284332

285333
internal record CodePostData(string CodeToInsert, List<PredictionCandidate> PredictionCandidates);

0 commit comments

Comments
 (0)