Use VS Code's built-in shell integration script

andyleejordan · andyleejordan · commit f85e2f8dc821 · 2024-04-02T13:09:26.000-07:00
Instead of having to maintain an edited copy (which was really annoying)
I stubbed out `PSConsoleHostReadLine` to do what they expect.
So now we can just use their script directly!
diff --git a/src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs b/src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs
@@ -35,6 +35,16 @@ internal class PsesInternalHost : PSHost, IHostSupportsInteractiveSession, IRuns
     {
         internal const string DefaultPrompt = "> ";
 
+        // TODO: Make the __psEditorServices prefix shared (it's used elsewhere too).
+        private const string ShellIntegrationScript = @"
+# Define fake PSConsoleHostReadLine so VS Code's script can execute it to get the user's input.
+$global:__psEditorServices_userInput = """";
+function global:PSConsoleHostReadLine { $global:__psEditorServices_userInput }
+
+# Execute VS Code's built-in shell integration script.
+try { . ""$(code --locate-shell-integration-path pwsh)"" } catch {}
+        ";
+
         private static readonly PSCommand s_promptCommand = new PSCommand().AddCommand("prompt");
 
         private static readonly PropertyInfo s_scriptDebuggerTriggerObjectProperty;
@@ -317,14 +327,14 @@ public async Task<bool> TryStartAsync(HostStartOptions startOptions, Cancellatio
 
             if (startOptions.ShellIntegrationEnabled)
             {
-                _logger.LogDebug("Enabling shell integration...");
+                _logger.LogDebug("Enabling VS Code's shell integration...");
                 _shellIntegrationEnabled = true;
                 await EnableShellIntegrationAsync(cancellationToken).ConfigureAwait(false);
                 _logger.LogDebug("Shell integration enabled!");
             }
             else
             {
-                _logger.LogDebug("Shell integration not enabled!");
+                _logger.LogDebug("VS Code's shell integration not enabled!");
             }
 
             await _started.Task.ConfigureAwait(false);
@@ -495,6 +505,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,
@@ -581,210 +592,14 @@ internal Task LoadHostProfilesAsync(CancellationToken cancellationToken)
                 cancellationToken);
         }
 
-        private Task EnableShellIntegrationAsync(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);
-        }
+        private Task EnableShellIntegrationAsync(CancellationToken cancellationToken) => ExecutePSCommandAsync(
+            new PSCommand().AddScript(ShellIntegrationScript),
+            cancellationToken,
+            new PowerShellExecutionOptions
+            {
+                AddToHistory = true, // VS Code adds theirs to the history, so we do too.
+                ThrowOnError = false
+            });
 
         public Task SetInitialWorkingDirectoryAsync(string path, CancellationToken cancellationToken)
         {
@@ -1262,16 +1077,36 @@ 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 VS Code's shell integration feature, we call PSConsoleHostReadLine specially as they've wrapped it.
+                // 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 their 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 = \"${input}\""),
+                        new PowerShellExecutionOptions { ThrowOnError = false },
+                        cancellationToken);
+
+                    // Invoke VS Code's PSConsoleHostReadLine wrapper.
+                    InvokePSCommand(
+                        new PSCommand().AddScript("PSConsoleHostReadLine"),
+                        new PowerShellExecutionOptions
+                        {
+                            ThrowOnError = false,
+                            WriteOutputToHost = true,
+                        },
+                        cancellationToken);
+
+                    // Reset our global variable.
+                    InvokePSCommand(
+                        new PSCommand().AddScript("$global:__psEditorServices_userInput = \"\""),
+                        new PowerShellExecutionOptions { ThrowOnError = false },
+                        cancellationToken);
                 }
 
                 InvokePSCommand(
-                    new PSCommand().AddScript(input, useLocalScope: false),
+                    new PSCommand().AddScript(input),
                     new PowerShellExecutionOptions
                     {
                         AddToHistory = true,
