Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use a provided shell integration script directly #2156

Merged
merged 1 commit into from
Apr 3, 2024
Merged

Use a provided shell integration script directly #2156

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions src/PowerShellEditorServices/Server/PsesLanguageServer.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,8 +153,9 @@ public async Task StartAsync()
InitialWorkingDirectory = initializationOptions?.GetValue("initialWorkingDirectory")?.Value<string>()
?? workspaceService.WorkspaceFolders.FirstOrDefault()?.Uri.GetFileSystemPath()
?? Directory.GetCurrentDirectory(),
ShellIntegrationEnabled = initializationOptions?.GetValue("shellIntegrationEnabled")?.Value<bool>()
?? false
// If a shell integration script path is provided, that implies the feature is enabled.
ShellIntegrationScript = initializationOptions?.GetValue("shellIntegrationScript")?.Value<string>()
?? "",
};

workspaceService.InitialWorkingDirectory = hostStartOptions.InitialWorkingDirectory;
Expand Down
2 changes: 1 addition & 1 deletion src/PowerShellEditorServices/Services/PowerShell/Host/HostStartOptions.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,6 @@ internal struct HostStartOptions

public string InitialWorkingDirectory { get; set; }

public bool ShellIntegrationEnabled { get; set; }
public string ShellIntegrationScript { get; set; }
}
}
256 changes: 45 additions & 211 deletions src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -315,16 +315,28 @@ public async Task<bool> TryStartAsync(HostStartOptions startOptions, Cancellatio
_logger.LogDebug("Profiles loaded!");
}

if (startOptions.ShellIntegrationEnabled)
if (!string.IsNullOrEmpty(startOptions.ShellIntegrationScript))
{
_logger.LogDebug("Enabling shell integration...");
_logger.LogDebug("Enabling Terminal Shell Integration...");
_shellIntegrationEnabled = true;
await EnableShellIntegrationAsync(cancellationToken).ConfigureAwait(false);
// TODO: Make the __psEditorServices prefix shared (it's used elsewhere too).
string setupShellIntegration = $$"""
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like C# 11's triple-quote double-dollar thing. Very useful.

# Setup Terminal Shell Integration.

# Define a fake PSConsoleHostReadLine so the integration script's wrapper
# can execute it to get the user's input.
$global:__psEditorServices_userInput = "";
function global:PSConsoleHostReadLine { $global:__psEditorServices_userInput }

# Execute the provided shell integration script.
try { . '{{startOptions.ShellIntegrationScript}}' } catch {}
""";
await EnableShellIntegrationAsync(setupShellIntegration, cancellationToken).ConfigureAwait(false);
_logger.LogDebug("Shell integration enabled!");
}
else
{
_logger.LogDebug("Shell integration not enabled!");
_logger.LogDebug("Terminal Shell Integration not enabled!");
}

await _started.Task.ConfigureAwait(false);
Expand Down Expand Up @@ -495,6 +507,7 @@ public Task ExecuteDelegateAsync(
new SynchronousDelegateTask(_logger, representation, executionOptions, action, cancellationToken));
}

// TODO: One day fix these so the cancellation token is last.
public Task<IReadOnlyList<TResult>> ExecutePSCommandAsync<TResult>(
PSCommand psCommand,
CancellationToken cancellationToken,
Expand Down Expand Up @@ -581,209 +594,12 @@ internal Task LoadHostProfilesAsync(CancellationToken cancellationToken)
cancellationToken);
}

private Task EnableShellIntegrationAsync(CancellationToken cancellationToken)
private Task EnableShellIntegrationAsync(string shellIntegrationScript, CancellationToken cancellationToken)
{
// Imported on 01/03/24 from
// https://github.com/microsoft/vscode/blob/main/src/vs/workbench/contrib/terminal/browser/media/shellIntegration.ps1
// with quotes escaped, `__VSCodeOriginalPSConsoleHostReadLine` removed (as it's done
// in our own ReadLine function), and `[Console]::Write` replaced with `Write-Host`.
const string shellIntegrationScript = @"
# Prevent installing more than once per session
if (Test-Path variable:global:__VSCodeOriginalPrompt) {
return;
}

# Disable shell integration when the language mode is restricted
if ($ExecutionContext.SessionState.LanguageMode -ne ""FullLanguage"") {
return;
}

$Global:__VSCodeOriginalPrompt = $function:Prompt

$Global:__LastHistoryId = -1

# Store the nonce in script scope and unset the global
$Nonce = $env:VSCODE_NONCE
$env:VSCODE_NONCE = $null

if ($env:VSCODE_ENV_REPLACE) {
$Split = $env:VSCODE_ENV_REPLACE.Split("":"")
foreach ($Item in $Split) {
$Inner = $Item.Split('=')
[Environment]::SetEnvironmentVariable($Inner[0], $Inner[1].Replace('\x3a', ':'))
}
$env:VSCODE_ENV_REPLACE = $null
}
if ($env:VSCODE_ENV_PREPEND) {
$Split = $env:VSCODE_ENV_PREPEND.Split("":"")
foreach ($Item in $Split) {
$Inner = $Item.Split('=')
[Environment]::SetEnvironmentVariable($Inner[0], $Inner[1].Replace('\x3a', ':') + [Environment]::GetEnvironmentVariable($Inner[0]))
}
$env:VSCODE_ENV_PREPEND = $null
}
if ($env:VSCODE_ENV_APPEND) {
$Split = $env:VSCODE_ENV_APPEND.Split("":"")
foreach ($Item in $Split) {
$Inner = $Item.Split('=')
[Environment]::SetEnvironmentVariable($Inner[0], [Environment]::GetEnvironmentVariable($Inner[0]) + $Inner[1].Replace('\x3a', ':'))
}
$env:VSCODE_ENV_APPEND = $null
}

function Global:__VSCode-Escape-Value([string]$value) {
# NOTE: In PowerShell v6.1+, this can be written `$value -replace '…', { … }` instead of `[regex]::Replace`.
# Replace any non-alphanumeric characters.
[regex]::Replace($value, '[\\\n;]', { param($match)
# Encode the (ascii) matches as `\x<hex>`
-Join (
[System.Text.Encoding]::UTF8.GetBytes($match.Value) | ForEach-Object { '\x{0:x2}' -f $_ }
)
})
}

function Global:Prompt() {
$FakeCode = [int]!$global:?
# NOTE: We disable strict mode for the scope of this function because it unhelpfully throws an
# error when $LastHistoryEntry is null, and is not otherwise useful.
Set-StrictMode -Off
$LastHistoryEntry = Get-History -Count 1
# Skip finishing the command if the first command has not yet started
if ($Global:__LastHistoryId -ne -1) {
if ($LastHistoryEntry.Id -eq $Global:__LastHistoryId) {
# Don't provide a command line or exit code if there was no history entry (eg. ctrl+c, enter on no command)
$Result = ""$([char]0x1b)]633;E`a""
$Result += ""$([char]0x1b)]633;D`a""
}
else {
# Command finished command line
# OSC 633 ; E ; <CommandLine?> ; <Nonce?> ST
$Result = ""$([char]0x1b)]633;E;""
# Sanitize the command line to ensure it can get transferred to the terminal and can be parsed
# correctly. This isn't entirely safe but good for most cases, it's important for the Pt parameter
# to only be composed of _printable_ characters as per the spec.
if ($LastHistoryEntry.CommandLine) {
$CommandLine = $LastHistoryEntry.CommandLine
}
else {
$CommandLine = """"
}
$Result += $(__VSCode-Escape-Value $CommandLine)
$Result += "";$Nonce""
$Result += ""`a""
# Command finished exit code
# OSC 633 ; D [; <ExitCode>] ST
$Result += ""$([char]0x1b)]633;D;$FakeCode`a""
}
}
# Prompt started
# OSC 633 ; A ST
$Result += ""$([char]0x1b)]633;A`a""
# Current working directory
# OSC 633 ; <Property>=<Value> ST
$Result += if ($pwd.Provider.Name -eq 'FileSystem') { ""$([char]0x1b)]633;P;Cwd=$(__VSCode-Escape-Value $pwd.ProviderPath)`a"" }
# Before running the original prompt, put $? back to what it was:
if ($FakeCode -ne 0) {
Write-Error ""failure"" -ea ignore
}
# Run the original prompt
$Result += $Global:__VSCodeOriginalPrompt.Invoke()
# Write command started
$Result += ""$([char]0x1b)]633;B`a""
$Global:__LastHistoryId = $LastHistoryEntry.Id
return $Result
}

# Set IsWindows property
if ($PSVersionTable.PSVersion -lt ""6.0"") {
# Windows PowerShell is only available on Windows
Write-Host -NoNewLine ""$([char]0x1b)]633;P;IsWindows=$true`a""
}
else {
Write-Host -NoNewLine ""$([char]0x1b)]633;P;IsWindows=$IsWindows`a""
}

# Set always on key handlers which map to default VS Code keybindings
function Set-MappedKeyHandler {
param ([string[]] $Chord, [string[]]$Sequence)
try {
$Handler = Get-PSReadLineKeyHandler -Chord $Chord | Select-Object -First 1
}
catch [System.Management.Automation.ParameterBindingException] {
# PowerShell 5.1 ships with PSReadLine 2.0.0 which does not have -Chord,
# so we check what's bound and filter it.
$Handler = Get-PSReadLineKeyHandler -Bound | Where-Object -FilterScript { $_.Key -eq $Chord } | Select-Object -First 1
}
if ($Handler) {
Set-PSReadLineKeyHandler -Chord $Sequence -Function $Handler.Function
}
}

$Global:__VSCodeHaltCompletions = $false
function Set-MappedKeyHandlers {
Set-MappedKeyHandler -Chord Ctrl+Spacebar -Sequence 'F12,a'
Set-MappedKeyHandler -Chord Alt+Spacebar -Sequence 'F12,b'
Set-MappedKeyHandler -Chord Shift+Enter -Sequence 'F12,c'
Set-MappedKeyHandler -Chord Shift+End -Sequence 'F12,d'

# Conditionally enable suggestions
if ($env:VSCODE_SUGGEST -eq '1') {
Remove-Item Env:VSCODE_SUGGEST

# VS Code send completions request (may override Ctrl+Spacebar)
Set-PSReadLineKeyHandler -Chord 'F12,e' -ScriptBlock {
Send-Completions
}

# Suggest trigger characters
Set-PSReadLineKeyHandler -Chord ""-"" -ScriptBlock {
[Microsoft.PowerShell.PSConsoleReadLine]::Insert(""-"")
if (!$Global:__VSCodeHaltCompletions) {
Send-Completions
}
}

Set-PSReadLineKeyHandler -Chord 'F12,y' -ScriptBlock {
$Global:__VSCodeHaltCompletions = $true
}

Set-PSReadLineKeyHandler -Chord 'F12,z' -ScriptBlock {
$Global:__VSCodeHaltCompletions = $false
}
}
}

function Send-Completions {
$commandLine = """"
$cursorIndex = 0
# TODO: Since fuzzy matching exists, should completions be provided only for character after the
# last space and then filter on the client side? That would let you trigger ctrl+space
# anywhere on a word and have full completions available
[Microsoft.PowerShell.PSConsoleReadLine]::GetBufferState([ref]$commandLine, [ref]$cursorIndex)
$completionPrefix = $commandLine

# Get completions
$result = ""`e]633;Completions""
if ($completionPrefix.Length -gt 0) {
# Get and send completions
$completions = TabExpansion2 -inputScript $completionPrefix -cursorColumn $cursorIndex
if ($null -ne $completions.CompletionMatches) {
$result += "";$($completions.ReplacementIndex);$($completions.ReplacementLength);$($cursorIndex);""
$result += $completions.CompletionMatches | ConvertTo-Json -Compress
}
}
$result += ""`a""

Write-Host -NoNewLine $result
}

# Register key handlers if PSReadLine is available
if (Get-Module -Name PSReadLine) {
Set-MappedKeyHandlers
}
";

return ExecutePSCommandAsync(new PSCommand().AddScript(shellIntegrationScript), cancellationToken);
return ExecutePSCommandAsync(
new PSCommand().AddScript(shellIntegrationScript),
cancellationToken,
new PowerShellExecutionOptions { AddToHistory = false, ThrowOnError = false });
}

public Task SetInitialWorkingDirectoryAsync(string path, CancellationToken cancellationToken)
Expand Down Expand Up @@ -1262,16 +1078,34 @@ private void InvokeInput(string input, CancellationToken cancellationToken)

try
{
// For VS Code's shell integration feature, this replaces their
// PSConsoleHostReadLine function wrapper, as that global function is not available
// to users of PSES, since we already wrap ReadLine ourselves.
// For the terminal shell integration feature, we call PSConsoleHostReadLine specially as it's been wrapped.
// Normally it would not be available (since we wrap ReadLine ourselves),
// but in this case we've made the original just emit the user's input so that the wrapper works as intended.
if (_shellIntegrationEnabled)
{
System.Console.Write("\x1b]633;C\a");
// Save the user's input to our special global variable so PSConsoleHostReadLine can read it.
InvokePSCommand(
new PSCommand().AddScript("$global:__psEditorServices_userInput = $args[0]").AddArgument(input),
new PowerShellExecutionOptions { ThrowOnError = false, WriteOutputToHost = false },
cancellationToken);

// Invoke the PSConsoleHostReadLine wrapper. We don't write the output because it
// returns the command line (user input) which would then be duplicate noise. Fortunately
// it writes the shell integration sequences directly using [Console]::Write.
InvokePSCommand(
new PSCommand().AddScript("PSConsoleHostReadLine"),
new PowerShellExecutionOptions { ThrowOnError = false, WriteOutputToHost = false },
cancellationToken);
Comment on lines +1092 to +1098
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Tyriar just a heads-up here that your PSConsoleHostReadLine wrapper must continue to write those escape sequences directly, as we're intentionally not emitting the output (like the returned value) to the host. Write-Host might work (though you don't use it), but we'd need to test.


// Reset our global variable.
InvokePSCommand(
new PSCommand().AddScript("$global:__psEditorServices_userInput = \"\""),
new PowerShellExecutionOptions { ThrowOnError = false, WriteOutputToHost = false },
cancellationToken);
}

InvokePSCommand(
new PSCommand().AddScript(input, useLocalScope: false),
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No real change, it's the default.

new PSCommand().AddScript(input),
new PowerShellExecutionOptions
{
AddToHistory = true,
Expand Down
Loading