Use a provided shell integration script directly

Instead of having to maintain an edited copy (which was really annoying)
I stubbed out `PSConsoleHostReadLine` to do what's expected.
So now we can just use the existing shell integration script directly!

Since we can't reliably find the script using
`code --locate-shell-integration-path pwsh` we now rely on it being sent
by the client on initialization. Its presence implies the feature is on.

This is pretty VS Code specific, but not necessarily so.

Author: andyleejordan

src/PowerShellEditorServices/Server/PsesLanguageServer.cs

@@ -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;

src/PowerShellEditorServices/Services/PowerShell/Host/HostStartOptions.cs

@@ -9,6 +9,6 @@ internal struct HostStartOptions
 
         public string InitialWorkingDirectory { get; set; }
 
-        public bool ShellIntegrationEnabled { get; set; }
+        public string ShellIntegrationScript { get; set; }
 }
 }

src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs

@@ -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 = $$"""
+                    # 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);
@@ -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,
@@ -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)
@@ -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 = \"{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);
+
+                    // 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),
+                    new PSCommand().AddScript(input),
                     new PowerShellExecutionOptions
                     {
                         AddToHistory = true,