Skip to content

Commit 282c015

Browse files
daxian-dbwdaxian-dbw
authored
Improve the reliability of Start-AIShell on macOS (#362)
1 parent 8367575 commit 282c015

File tree

4 files changed

+212
-104
lines changed

4 files changed

+212
-104
lines changed

‎shell/AIShell.Integration/AIShell.psd1

+1
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
Copyright = '(c) Microsoft Corporation. All rights reserved.'
99
Description = 'Integration with the AIShell to provide intelligent shell experience'
1010
PowerShellVersion = '7.4.6'
11+
PowerShellHostName = 'ConsoleHost'
1112
FunctionsToExport = @()
1213
CmdletsToExport = @('Start-AIShell','Invoke-AIShell','Resolve-Error')
1314
VariablesToExport = '*'

‎shell/AIShell.Integration/Channel.cs

-8
Original file line numberDiff line numberDiff line change
@@ -283,11 +283,3 @@ private void PSRLAcceptLine()
283283
}
284284

285285
internal record CodePostData(string CodeToInsert, List<PredictionCandidate> PredictionCandidates);
286-
287-
public class Init : IModuleAssemblyCleanup
288-
{
289-
public void OnRemove(PSModuleInfo psModuleInfo)
290-
{
291-
Channel.Singleton?.Dispose();
292-
}
293-
}

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

+62-96
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
using System.Diagnostics;
22
using System.Management.Automation;
3-
using System.Text;
43

54
namespace AIShell.Integration.Commands;
65

@@ -12,6 +11,10 @@ public class StartAIShellCommand : PSCmdlet
1211
[ValidateNotNullOrEmpty]
1312
public string Path { get; set; }
1413

14+
private string _venvPipPath;
15+
private string _venvPythonPath;
16+
private static bool s_iterm2Installed = false;
17+
1518
protected override void BeginProcessing()
1619
{
1720
if (Path is null)
@@ -80,25 +83,21 @@ protected override void BeginProcessing()
8083
targetObject: null));
8184
}
8285

83-
var python = SessionState.InvokeCommand.GetCommand("python3", CommandTypes.Application);
84-
if (python is null)
86+
try
8587
{
86-
ThrowTerminatingError(new(
87-
new NotSupportedException("The executable 'python3' (Windows Terminal) cannot be found. It's required to split a pane in iTerm2 programmatically."),
88-
"Python3Missing",
89-
ErrorCategory.NotInstalled,
90-
targetObject: null));
88+
InitAndCleanup.CreateVirtualEnvTask.GetAwaiter().GetResult();
9189
}
92-
93-
var pip3 = SessionState.InvokeCommand.GetCommand("pip3", CommandTypes.Application);
94-
if (pip3 is null)
90+
catch (Exception exception)
9591
{
9692
ThrowTerminatingError(new(
97-
new NotSupportedException("The executable 'pip3' cannot be found. It's required to split a pane in iTerm2 programmatically."),
98-
"Pip3Missing",
99-
ErrorCategory.NotInstalled,
93+
exception,
94+
"FailedToCreateVirtualEnvironment",
95+
ErrorCategory.InvalidOperation,
10096
targetObject: null));
10197
}
98+
99+
_venvPipPath = System.IO.Path.Join(InitAndCleanup.VirtualEnvPath, "bin", "pip3");
100+
_venvPythonPath = System.IO.Path.Join(InitAndCleanup.VirtualEnvPath, "bin", "python3");
102101
}
103102
else
104103
{
@@ -112,12 +111,11 @@ protected override void BeginProcessing()
112111

113112
protected override void EndProcessing()
114113
{
115-
string pipeName = Channel.Singleton.StartChannelSetup();
116-
117114
if (OperatingSystem.IsWindows())
118115
{
119116
ProcessStartInfo startInfo;
120117
string wtProfileGuid = Environment.GetEnvironmentVariable("WT_PROFILE_ID");
118+
string pipeName = Channel.Singleton.StartChannelSetup();
121119

122120
if (wtProfileGuid is null)
123121
{
@@ -169,94 +167,62 @@ protected override void EndProcessing()
169167
}
170168
else if (OperatingSystem.IsMacOS())
171169
{
172-
// Install the Python package 'iterm2'.
173-
ProcessStartInfo startInfo = new("pip3")
174-
{
175-
ArgumentList = { "install", "-q", "iterm2" },
176-
RedirectStandardError = true,
177-
RedirectStandardOutput = true
178-
};
179-
180-
Process proc = new() { StartInfo = startInfo };
181-
proc.Start();
182-
proc.WaitForExit();
170+
Process proc;
171+
ProcessStartInfo startInfo;
183172

184-
if (proc.ExitCode is 1)
173+
// Install the Python package 'iterm2' to the venv.
174+
if (!s_iterm2Installed)
185175
{
186-
ThrowTerminatingError(new(
187-
new NotSupportedException("The Python package 'iterm2' cannot be installed. It's required to split a pane in iTerm2 programmatically."),
188-
"iterm2Missing",
189-
ErrorCategory.NotInstalled,
190-
targetObject: null));
191-
}
176+
startInfo = new(_venvPipPath)
177+
{
178+
// Make 'pypi.org' and 'files.pythonhosted.org' as trusted hosts, because a security software
179+
// may cause issue to SSL validation for access to/from those two endpoints.
180+
// See https://stackoverflow.com/a/71993364 for details.
181+
ArgumentList = {
182+
"install",
183+
"-q",
184+
"--disable-pip-version-check",
185+
"--trusted-host",
186+
"pypi.org",
187+
"--trusted-host",
188+
"files.pythonhosted.org",
189+
"iterm2"
190+
},
191+
RedirectStandardError = true,
192+
RedirectStandardOutput = true
193+
};
192194

193-
proc.Dispose();
195+
proc = Process.Start(startInfo);
196+
proc.WaitForExit();
194197

195-
// Write the Python script to a temp file, if not yet.
196-
string pythonScript = System.IO.Path.Combine(System.IO.Path.GetTempPath(), "__aish_split_pane.py");
197-
if (!File.Exists(pythonScript))
198-
{
199-
File.WriteAllText(pythonScript, SplitPanePythonCode, Encoding.UTF8);
198+
if (proc.ExitCode is 0)
199+
{
200+
s_iterm2Installed = true;
201+
}
202+
else
203+
{
204+
string error = "The Python package 'iterm2' cannot be installed. It's required to split a pane in iTerm2 programmatically.";
205+
string stderr = proc.StandardError.ReadToEnd();
206+
if (!string.IsNullOrEmpty(stderr))
207+
{
208+
error = $"{error}\nError details:\n{stderr}";
209+
}
210+
211+
ThrowTerminatingError(new(
212+
new NotSupportedException(error),
213+
"iterm2Missing",
214+
ErrorCategory.NotInstalled,
215+
targetObject: null));
216+
}
217+
218+
proc.Dispose();
200219
}
201220

202221
// Run the Python script to split the pane and start AIShell.
203-
startInfo = new("python3") { ArgumentList = { pythonScript, Path, pipeName } };
204-
proc = new() { StartInfo = startInfo };
205-
proc.Start();
222+
string pipeName = Channel.Singleton.StartChannelSetup();
223+
startInfo = new(_venvPythonPath) { ArgumentList = { InitAndCleanup.PythonScript, Path, pipeName } };
224+
proc = Process.Start(startInfo);
206225
proc.WaitForExit();
207226
}
208227
}
209-
210-
private const string SplitPanePythonCode = """
211-
import iterm2
212-
import sys
213-
214-
# iTerm needs to be running for this to work
215-
async def main(connection):
216-
app = await iterm2.async_get_app(connection)
217-
218-
# Foreground the app
219-
await app.async_activate()
220-
221-
window = app.current_terminal_window
222-
if window is not None:
223-
# Get the current pane so that we can split it.
224-
current_tab = window.current_tab
225-
current_pane = current_tab.current_session
226-
227-
# Get the total width before splitting.
228-
width = current_pane.grid_size.width
229-
230-
# Split pane vertically
231-
split_pane = await current_pane.async_split_pane(vertical=True)
232-
233-
# Get the height of the pane after splitting. This value will be
234-
# slightly smaller than its height before splitting.
235-
height = current_pane.grid_size.height
236-
237-
# Calculate the new width for both panes using the ratio 0.4 for the new pane.
238-
# Then set the preferred size for both pane sessions.
239-
new_current_width = round(width * 0.6);
240-
new_split_width = width - new_current_width;
241-
current_pane.preferred_size = iterm2.Size(new_current_width, height)
242-
split_pane.preferred_size = iterm2.Size(new_split_width, height);
243-
244-
# Update the layout, which will change the panes to preferred size.
245-
await current_tab.async_update_layout()
246-
247-
await split_pane.async_send_text(f'{app_path} --channel {channel}\n')
248-
else:
249-
# You can view this message in the script console.
250-
print("No current iTerm2 window. Make sure you are running in iTerm2.")
251-
252-
if len(sys.argv) > 1:
253-
app_path = sys.argv[1]
254-
channel = sys.argv[2]
255-
256-
# Do not specify True for retry. It's possible that the user hasn't enable the Python API for iTerm2,
257-
# and in that case, we want it to fail immediately instead of stucking in retries.
258-
iterm2.run_until_complete(main)
259-
else:
260-
print("Please provide the application path as a command line argument.")
261-
""";
262228
}

‎shell/AIShell.Integration/InitAndCleanup.cs

+149
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
using System.Diagnostics;
2+
using System.Globalization;
3+
using System.Text;
4+
using System.Management.Automation;
5+
6+
namespace AIShell.Integration;
7+
8+
public class InitAndCleanup : IModuleAssemblyInitializer, IModuleAssemblyCleanup
9+
{
10+
private const int ScriptVersion = 1;
11+
private const string ScriptFileTemplate = "aish_split_pane_v{0}.py";
12+
private const string SplitPanePythonCode = """
13+
import iterm2
14+
import sys
15+
16+
# iTerm needs to be running for this to work
17+
async def main(connection):
18+
app = await iterm2.async_get_app(connection)
19+
20+
# Foreground the app
21+
await app.async_activate()
22+
23+
window = app.current_terminal_window
24+
if window is not None:
25+
# Get the current pane so that we can split it.
26+
current_tab = window.current_tab
27+
current_pane = current_tab.current_session
28+
29+
# Get the total width before splitting.
30+
width = current_pane.grid_size.width
31+
32+
change = iterm2.LocalWriteOnlyProfile()
33+
change.set_use_custom_command('Yes')
34+
change.set_command(f'{app_path} --channel {channel}')
35+
36+
# Split pane vertically
37+
split_pane = await current_pane.async_split_pane(vertical=True, profile_customizations=change)
38+
39+
# Get the height of the pane after splitting. This value will be
40+
# slightly smaller than its height before splitting.
41+
height = current_pane.grid_size.height
42+
43+
# Calculate the new width for both panes using the ratio 0.4 for the new pane.
44+
# Then set the preferred size for both pane sessions.
45+
new_current_width = round(width * 0.6);
46+
new_split_width = width - new_current_width;
47+
current_pane.preferred_size = iterm2.Size(new_current_width, height)
48+
split_pane.preferred_size = iterm2.Size(new_split_width, height);
49+
50+
# Update the layout, which will change the panes to preferred size.
51+
await current_tab.async_update_layout()
52+
else:
53+
# You can view this message in the script console.
54+
print("No current iTerm2 window. Make sure you are running in iTerm2.")
55+
56+
if len(sys.argv) > 1:
57+
app_path = sys.argv[1]
58+
channel = sys.argv[2]
59+
60+
# Do not specify True for retry. It's possible that the user hasn't enable the Python API for iTerm2,
61+
# and in that case, we want it to fail immediately instead of stucking in retries.
62+
iterm2.run_until_complete(main)
63+
else:
64+
print("Please provide the application path as a command line argument.")
65+
""";
66+
67+
internal static string CachePath { get; }
68+
internal static string PythonScript { get; }
69+
internal static string VirtualEnvPath { get; }
70+
internal static Task CreateVirtualEnvTask { get; }
71+
72+
static InitAndCleanup()
73+
{
74+
CachePath = Path.Join(Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), ".aish", ".cache");
75+
PythonScript = null;
76+
VirtualEnvPath = null;
77+
CreateVirtualEnvTask = null;
78+
79+
if (OperatingSystem.IsMacOS())
80+
{
81+
PythonScript = Path.Join(CachePath, string.Format(CultureInfo.InvariantCulture, ScriptFileTemplate, ScriptVersion));
82+
VirtualEnvPath = Path.Join(CachePath, ".venv");
83+
CreateVirtualEnvTask = Task.Run(CreatePythonVirtualEnvironment);
84+
}
85+
}
86+
87+
private static void CreatePythonVirtualEnvironment()
88+
{
89+
// Simply return if the virtual environment was already created.
90+
if (Directory.Exists(VirtualEnvPath))
91+
{
92+
return;
93+
}
94+
95+
// Create a virtual environment where we can install the needed pacakges.
96+
ProcessStartInfo startInfo = new("python3")
97+
{
98+
ArgumentList = { "-m", "venv", VirtualEnvPath },
99+
RedirectStandardError = true,
100+
RedirectStandardOutput = true
101+
};
102+
103+
Process proc = Process.Start(startInfo);
104+
proc.WaitForExit();
105+
106+
if (proc.ExitCode is 1)
107+
{
108+
string error = $"Failed to create a virtual environment by 'python3 -m venv {VirtualEnvPath}'.";
109+
string stderr = proc.StandardError.ReadToEnd();
110+
if (!string.IsNullOrEmpty(stderr))
111+
{
112+
error = $"{error}\nError details:\n{stderr}";
113+
}
114+
115+
throw new NotSupportedException(error);
116+
}
117+
118+
proc.Dispose();
119+
}
120+
121+
public void OnImport()
122+
{
123+
if (!OperatingSystem.IsMacOS())
124+
{
125+
return;
126+
}
127+
128+
// Remove old scripts, if there is any.
129+
for (int i = 1; i < ScriptVersion; i++)
130+
{
131+
string oldScript = Path.Join(CachePath, string.Format(CultureInfo.InvariantCulture, ScriptFileTemplate, i));
132+
if (File.Exists(oldScript))
133+
{
134+
File.Delete(oldScript);
135+
}
136+
}
137+
138+
// Create the latest script, if not yet.
139+
if (!File.Exists(PythonScript))
140+
{
141+
File.WriteAllText(PythonScript, SplitPanePythonCode, Encoding.UTF8);
142+
}
143+
}
144+
145+
public void OnRemove(PSModuleInfo psModuleInfo)
146+
{
147+
Channel.Singleton?.Dispose();
148+
}
149+
}

0 commit comments

Comments
 (0)