diff --git a/shell/AIShell.Integration/AIShell.psd1 b/shell/AIShell.Integration/AIShell.psd1 index c8dfb0d..264cc6d 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 cdae0a4..47069c5 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 3501c4e..030e772 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,10 @@ public class StartAIShellCommand : PSCmdlet [ValidateNotNullOrEmpty] public string Path { get; set; } + private string _venvPipPath; + private string _venvPythonPath; + private static bool s_iterm2Installed = false; + protected override void BeginProcessing() { if (Path is null) @@ -80,25 +83,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 +111,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,94 +167,62 @@ protected override void EndProcessing() } else if (OperatingSystem.IsMacOS()) { - // Install the Python package 'iterm2'. - ProcessStartInfo startInfo = new("pip3") - { - ArgumentList = { "install", "-q", "iterm2" }, - RedirectStandardError = true, - RedirectStandardOutput = true - }; - - Process proc = new() { StartInfo = startInfo }; - proc.Start(); - proc.WaitForExit(); + Process proc; + ProcessStartInfo startInfo; - if (proc.ExitCode is 1) + // Install the Python package 'iterm2' to the venv. + if (!s_iterm2Installed) { - ThrowTerminatingError(new( - new NotSupportedException("The Python package 'iterm2' cannot be installed. It's required to split a pane in iTerm2 programmatically."), - "iterm2Missing", - ErrorCategory.NotInstalled, - targetObject: null)); - } + startInfo = new(_venvPipPath) + { + // 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 + }; - proc.Dispose(); + proc = Process.Start(startInfo); + proc.WaitForExit(); - // 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); + if (proc.ExitCode is 0) + { + 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)); + } + + proc.Dispose(); } // 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 0000000..b5d5e94 --- /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(); + } +}