From 65cf6231f698ff31f590827c26dfb91cc9c93dd6 Mon Sep 17 00:00:00 2001 From: Dongbo Wang Date: Fri, 4 Apr 2025 15:22:19 -0700 Subject: [PATCH 1/3] Improve the reliability of `Start-AIShell` on macOS --- shell/AIShell.Integration/AIShell.psd1 | 1 + shell/AIShell.Integration/Channel.cs | 8 - .../Commands/StartAishCommand.cs | 113 ++++--------- shell/AIShell.Integration/InitAndCleanup.cs | 149 ++++++++++++++++++ 4 files changed, 178 insertions(+), 93 deletions(-) create mode 100644 shell/AIShell.Integration/InitAndCleanup.cs diff --git a/shell/AIShell.Integration/AIShell.psd1 b/shell/AIShell.Integration/AIShell.psd1 index c8dfb0d1..264cc6d2 100644 --- a/shell/AIShell.Integration/AIShell.psd1 +++ b/shell/AIShell.Integration/AIShell.psd1 @@ -8,6 +8,7 @@ Copyright = '(c) Microsoft Corporation. All rights reserved.' Description = 'Integration with the AIShell to provide intelligent shell experience' PowerShellVersion = '7.4.6' + PowerShellHostName = 'ConsoleHost' FunctionsToExport = @() CmdletsToExport = @('Start-AIShell','Invoke-AIShell','Resolve-Error') VariablesToExport = '*' diff --git a/shell/AIShell.Integration/Channel.cs b/shell/AIShell.Integration/Channel.cs index cdae0a4c..47069c5a 100644 --- a/shell/AIShell.Integration/Channel.cs +++ b/shell/AIShell.Integration/Channel.cs @@ -283,11 +283,3 @@ private void PSRLAcceptLine() } internal record CodePostData(string CodeToInsert, List PredictionCandidates); - -public class Init : IModuleAssemblyCleanup -{ - public void OnRemove(PSModuleInfo psModuleInfo) - { - Channel.Singleton?.Dispose(); - } -} diff --git a/shell/AIShell.Integration/Commands/StartAishCommand.cs b/shell/AIShell.Integration/Commands/StartAishCommand.cs index 3501c4e9..12c6683f 100644 --- a/shell/AIShell.Integration/Commands/StartAishCommand.cs +++ b/shell/AIShell.Integration/Commands/StartAishCommand.cs @@ -1,6 +1,5 @@ using System.Diagnostics; using System.Management.Automation; -using System.Text; namespace AIShell.Integration.Commands; @@ -12,6 +11,9 @@ public class StartAIShellCommand : PSCmdlet [ValidateNotNullOrEmpty] public string Path { get; set; } + private string _venvPipPath; + private string _venvPythonPath; + protected override void BeginProcessing() { if (Path is null) @@ -80,25 +82,21 @@ protected override void BeginProcessing() targetObject: null)); } - var python = SessionState.InvokeCommand.GetCommand("python3", CommandTypes.Application); - if (python is null) + try { - ThrowTerminatingError(new( - new NotSupportedException("The executable 'python3' (Windows Terminal) cannot be found. It's required to split a pane in iTerm2 programmatically."), - "Python3Missing", - ErrorCategory.NotInstalled, - targetObject: null)); + InitAndCleanup.CreateVirtualEnvTask.GetAwaiter().GetResult(); } - - var pip3 = SessionState.InvokeCommand.GetCommand("pip3", CommandTypes.Application); - if (pip3 is null) + catch (Exception exception) { ThrowTerminatingError(new( - new NotSupportedException("The executable 'pip3' cannot be found. It's required to split a pane in iTerm2 programmatically."), - "Pip3Missing", - ErrorCategory.NotInstalled, + exception, + "FailedToCreateVirtualEnvironment", + ErrorCategory.InvalidOperation, targetObject: null)); } + + _venvPipPath = System.IO.Path.Join(InitAndCleanup.VirtualEnvPath, "bin", "pip3"); + _venvPythonPath = System.IO.Path.Join(InitAndCleanup.VirtualEnvPath, "bin", "python3"); } else { @@ -112,12 +110,11 @@ protected override void BeginProcessing() protected override void EndProcessing() { - string pipeName = Channel.Singleton.StartChannelSetup(); - if (OperatingSystem.IsWindows()) { ProcessStartInfo startInfo; string wtProfileGuid = Environment.GetEnvironmentVariable("WT_PROFILE_ID"); + string pipeName = Channel.Singleton.StartChannelSetup(); if (wtProfileGuid is null) { @@ -169,22 +166,28 @@ protected override void EndProcessing() } else if (OperatingSystem.IsMacOS()) { - // Install the Python package 'iterm2'. - ProcessStartInfo startInfo = new("pip3") + // Install the Python package 'iterm2' to the venv. + ProcessStartInfo startInfo = new(_venvPipPath) { - ArgumentList = { "install", "-q", "iterm2" }, + ArgumentList = { "install", "-q", "iterm2", "--disable-pip-version-check" }, RedirectStandardError = true, RedirectStandardOutput = true }; - Process proc = new() { StartInfo = startInfo }; - proc.Start(); + Process proc = Process.Start(startInfo); proc.WaitForExit(); if (proc.ExitCode is 1) { + string error = "The Python package 'iterm2' cannot be installed. It's required to split a pane in iTerm2 programmatically."; + string stderr = proc.StandardError.ReadToEnd(); + if (!string.IsNullOrEmpty(stderr)) + { + error = $"{error}\nError details:\n{stderr}"; + } + ThrowTerminatingError(new( - new NotSupportedException("The Python package 'iterm2' cannot be installed. It's required to split a pane in iTerm2 programmatically."), + new NotSupportedException(error), "iterm2Missing", ErrorCategory.NotInstalled, targetObject: null)); @@ -192,71 +195,11 @@ protected override void EndProcessing() proc.Dispose(); - // Write the Python script to a temp file, if not yet. - string pythonScript = System.IO.Path.Combine(System.IO.Path.GetTempPath(), "__aish_split_pane.py"); - if (!File.Exists(pythonScript)) - { - File.WriteAllText(pythonScript, SplitPanePythonCode, Encoding.UTF8); - } - // Run the Python script to split the pane and start AIShell. - startInfo = new("python3") { ArgumentList = { pythonScript, Path, pipeName } }; - proc = new() { StartInfo = startInfo }; - proc.Start(); + string pipeName = Channel.Singleton.StartChannelSetup(); + startInfo = new(_venvPythonPath) { ArgumentList = { InitAndCleanup.PythonScript, Path, pipeName } }; + proc = Process.Start(startInfo); proc.WaitForExit(); } } - - private const string SplitPanePythonCode = """ - import iterm2 - import sys - - # iTerm needs to be running for this to work - async def main(connection): - app = await iterm2.async_get_app(connection) - - # Foreground the app - await app.async_activate() - - window = app.current_terminal_window - if window is not None: - # Get the current pane so that we can split it. - current_tab = window.current_tab - current_pane = current_tab.current_session - - # Get the total width before splitting. - width = current_pane.grid_size.width - - # Split pane vertically - split_pane = await current_pane.async_split_pane(vertical=True) - - # Get the height of the pane after splitting. This value will be - # slightly smaller than its height before splitting. - height = current_pane.grid_size.height - - # Calculate the new width for both panes using the ratio 0.4 for the new pane. - # Then set the preferred size for both pane sessions. - new_current_width = round(width * 0.6); - new_split_width = width - new_current_width; - current_pane.preferred_size = iterm2.Size(new_current_width, height) - split_pane.preferred_size = iterm2.Size(new_split_width, height); - - # Update the layout, which will change the panes to preferred size. - await current_tab.async_update_layout() - - await split_pane.async_send_text(f'{app_path} --channel {channel}\n') - else: - # You can view this message in the script console. - print("No current iTerm2 window. Make sure you are running in iTerm2.") - - if len(sys.argv) > 1: - app_path = sys.argv[1] - channel = sys.argv[2] - - # Do not specify True for retry. It's possible that the user hasn't enable the Python API for iTerm2, - # and in that case, we want it to fail immediately instead of stucking in retries. - iterm2.run_until_complete(main) - else: - print("Please provide the application path as a command line argument.") - """; } diff --git a/shell/AIShell.Integration/InitAndCleanup.cs b/shell/AIShell.Integration/InitAndCleanup.cs new file mode 100644 index 00000000..b5d5e945 --- /dev/null +++ b/shell/AIShell.Integration/InitAndCleanup.cs @@ -0,0 +1,149 @@ +using System.Diagnostics; +using System.Globalization; +using System.Text; +using System.Management.Automation; + +namespace AIShell.Integration; + +public class InitAndCleanup : IModuleAssemblyInitializer, IModuleAssemblyCleanup +{ + private const int ScriptVersion = 1; + private const string ScriptFileTemplate = "aish_split_pane_v{0}.py"; + private const string SplitPanePythonCode = """ + import iterm2 + import sys + + # iTerm needs to be running for this to work + async def main(connection): + app = await iterm2.async_get_app(connection) + + # Foreground the app + await app.async_activate() + + window = app.current_terminal_window + if window is not None: + # Get the current pane so that we can split it. + current_tab = window.current_tab + current_pane = current_tab.current_session + + # Get the total width before splitting. + width = current_pane.grid_size.width + + change = iterm2.LocalWriteOnlyProfile() + change.set_use_custom_command('Yes') + change.set_command(f'{app_path} --channel {channel}') + + # Split pane vertically + split_pane = await current_pane.async_split_pane(vertical=True, profile_customizations=change) + + # Get the height of the pane after splitting. This value will be + # slightly smaller than its height before splitting. + height = current_pane.grid_size.height + + # Calculate the new width for both panes using the ratio 0.4 for the new pane. + # Then set the preferred size for both pane sessions. + new_current_width = round(width * 0.6); + new_split_width = width - new_current_width; + current_pane.preferred_size = iterm2.Size(new_current_width, height) + split_pane.preferred_size = iterm2.Size(new_split_width, height); + + # Update the layout, which will change the panes to preferred size. + await current_tab.async_update_layout() + else: + # You can view this message in the script console. + print("No current iTerm2 window. Make sure you are running in iTerm2.") + + if len(sys.argv) > 1: + app_path = sys.argv[1] + channel = sys.argv[2] + + # Do not specify True for retry. It's possible that the user hasn't enable the Python API for iTerm2, + # and in that case, we want it to fail immediately instead of stucking in retries. + iterm2.run_until_complete(main) + else: + print("Please provide the application path as a command line argument.") + """; + + internal static string CachePath { get; } + internal static string PythonScript { get; } + internal static string VirtualEnvPath { get; } + internal static Task CreateVirtualEnvTask { get; } + + static InitAndCleanup() + { + CachePath = Path.Join(Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), ".aish", ".cache"); + PythonScript = null; + VirtualEnvPath = null; + CreateVirtualEnvTask = null; + + if (OperatingSystem.IsMacOS()) + { + PythonScript = Path.Join(CachePath, string.Format(CultureInfo.InvariantCulture, ScriptFileTemplate, ScriptVersion)); + VirtualEnvPath = Path.Join(CachePath, ".venv"); + CreateVirtualEnvTask = Task.Run(CreatePythonVirtualEnvironment); + } + } + + private static void CreatePythonVirtualEnvironment() + { + // Simply return if the virtual environment was already created. + if (Directory.Exists(VirtualEnvPath)) + { + return; + } + + // Create a virtual environment where we can install the needed pacakges. + ProcessStartInfo startInfo = new("python3") + { + ArgumentList = { "-m", "venv", VirtualEnvPath }, + RedirectStandardError = true, + RedirectStandardOutput = true + }; + + Process proc = Process.Start(startInfo); + proc.WaitForExit(); + + if (proc.ExitCode is 1) + { + string error = $"Failed to create a virtual environment by 'python3 -m venv {VirtualEnvPath}'."; + string stderr = proc.StandardError.ReadToEnd(); + if (!string.IsNullOrEmpty(stderr)) + { + error = $"{error}\nError details:\n{stderr}"; + } + + throw new NotSupportedException(error); + } + + proc.Dispose(); + } + + public void OnImport() + { + if (!OperatingSystem.IsMacOS()) + { + return; + } + + // Remove old scripts, if there is any. + for (int i = 1; i < ScriptVersion; i++) + { + string oldScript = Path.Join(CachePath, string.Format(CultureInfo.InvariantCulture, ScriptFileTemplate, i)); + if (File.Exists(oldScript)) + { + File.Delete(oldScript); + } + } + + // Create the latest script, if not yet. + if (!File.Exists(PythonScript)) + { + File.WriteAllText(PythonScript, SplitPanePythonCode, Encoding.UTF8); + } + } + + public void OnRemove(PSModuleInfo psModuleInfo) + { + Channel.Singleton?.Dispose(); + } +} From f2c6952d3e4c0acdc6c5ec387be3c2b0bb82d87b Mon Sep 17 00:00:00 2001 From: Dongbo Wang Date: Fri, 4 Apr 2025 16:29:50 -0700 Subject: [PATCH 2/3] Avoid running 'pip3 install iterm2' per every 'Invoke-AIShell' call --- .../Commands/StartAishCommand.cs | 51 +++++++++++-------- 1 file changed, 31 insertions(+), 20 deletions(-) diff --git a/shell/AIShell.Integration/Commands/StartAishCommand.cs b/shell/AIShell.Integration/Commands/StartAishCommand.cs index 12c6683f..602bb197 100644 --- a/shell/AIShell.Integration/Commands/StartAishCommand.cs +++ b/shell/AIShell.Integration/Commands/StartAishCommand.cs @@ -13,6 +13,7 @@ public class StartAIShellCommand : PSCmdlet private string _venvPipPath; private string _venvPythonPath; + private static bool s_iterm2Installed = false; protected override void BeginProcessing() { @@ -166,35 +167,45 @@ protected override void EndProcessing() } else if (OperatingSystem.IsMacOS()) { + Process proc; + ProcessStartInfo startInfo; + // Install the Python package 'iterm2' to the venv. - ProcessStartInfo startInfo = new(_venvPipPath) + if (!s_iterm2Installed) { - ArgumentList = { "install", "-q", "iterm2", "--disable-pip-version-check" }, - RedirectStandardError = true, - RedirectStandardOutput = true - }; + startInfo = new(_venvPipPath) + { + ArgumentList = { "install", "-q", "iterm2", "--disable-pip-version-check" }, + RedirectStandardError = true, + RedirectStandardOutput = true + }; - Process proc = Process.Start(startInfo); - proc.WaitForExit(); + proc = Process.Start(startInfo); + proc.WaitForExit(); - if (proc.ExitCode is 1) - { - string error = "The Python package 'iterm2' cannot be installed. It's required to split a pane in iTerm2 programmatically."; - string stderr = proc.StandardError.ReadToEnd(); - if (!string.IsNullOrEmpty(stderr)) + if (proc.ExitCode is 0) { - error = $"{error}\nError details:\n{stderr}"; + s_iterm2Installed = true; + } + else + { + string error = "The Python package 'iterm2' cannot be installed. It's required to split a pane in iTerm2 programmatically."; + string stderr = proc.StandardError.ReadToEnd(); + if (!string.IsNullOrEmpty(stderr)) + { + error = $"{error}\nError details:\n{stderr}"; + } + + ThrowTerminatingError(new( + new NotSupportedException(error), + "iterm2Missing", + ErrorCategory.NotInstalled, + targetObject: null)); } - ThrowTerminatingError(new( - new NotSupportedException(error), - "iterm2Missing", - ErrorCategory.NotInstalled, - targetObject: null)); + proc.Dispose(); } - proc.Dispose(); - // Run the Python script to split the pane and start AIShell. string pipeName = Channel.Singleton.StartChannelSetup(); startInfo = new(_venvPythonPath) { ArgumentList = { InitAndCleanup.PythonScript, Path, pipeName } }; From 7a2b1e25e85c6a86ab3c6f61061ebb91fc9c0ccb Mon Sep 17 00:00:00 2001 From: Dongbo Wang Date: Fri, 4 Apr 2025 17:20:50 -0700 Subject: [PATCH 3/3] Add trusted-host flags for pip3 call --- .../Commands/StartAishCommand.cs | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/shell/AIShell.Integration/Commands/StartAishCommand.cs b/shell/AIShell.Integration/Commands/StartAishCommand.cs index 602bb197..030e7720 100644 --- a/shell/AIShell.Integration/Commands/StartAishCommand.cs +++ b/shell/AIShell.Integration/Commands/StartAishCommand.cs @@ -175,7 +175,19 @@ protected override void EndProcessing() { startInfo = new(_venvPipPath) { - ArgumentList = { "install", "-q", "iterm2", "--disable-pip-version-check" }, + // Make 'pypi.org' and 'files.pythonhosted.org' as trusted hosts, because a security software + // may cause issue to SSL validation for access to/from those two endpoints. + // See https://stackoverflow.com/a/71993364 for details. + ArgumentList = { + "install", + "-q", + "--disable-pip-version-check", + "--trusted-host", + "pypi.org", + "--trusted-host", + "files.pythonhosted.org", + "iterm2" + }, RedirectStandardError = true, RedirectStandardOutput = true };