Update shell integration script to fix command decorations

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

@@ -583,7 +583,7 @@ internal Task LoadHostProfilesAsync(CancellationToken cancellationToken)
 
         private Task EnableShellIntegrationAsync(CancellationToken cancellationToken)
         {
-            // Imported on 01/03/23 from
+            // 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`.
@@ -602,42 +602,74 @@ private Task EnableShellIntegrationAsync(CancellationToken cancellationToken)
 
 $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 $_ }
-		)
-	})
+			# 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
-	$FakeCode = [int]!$global:?
 	$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;E`a""
 			$Result += ""$([char]0x1b)]633;D`a""
-		} else {
+		}
+		else {
 			# Command finished command line
-			# OSC 633 ; A ; <CommandLine?> ST
-			$Result  = ""$([char]0x1b)]633;E;""
+			# 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 {
+			}
+			else {
 				$CommandLine = """"
 			}
 			$Result += $(__VSCode-Escape-Value $CommandLine)
+			$Result += "";$Nonce""
 			$Result += ""`a""
 			# Command finished exit code
 			# OSC 633 ; D [; <ExitCode>] ST
@@ -649,7 +681,7 @@ private Task EnableShellIntegrationAsync(CancellationToken cancellationToken)
 	$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""}
+	$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
@@ -664,28 +696,91 @@ private Task EnableShellIntegrationAsync(CancellationToken cancellationToken)
 
 # Set IsWindows property
 if ($PSVersionTable.PSVersion -lt ""6.0"") {
-	[Console]::Write(""$([char]0x1b)]633;P;IsWindows=$true`a"")
-} else {
-	[Console]::Write(""$([char]0x1b)]633;P;IsWindows=$IsWindows`a"")
+	# 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)
-	$Handler = $(Get-PSReadLineKeyHandler -Chord $Chord | Select-Object -First 1)
+	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
 }
 
-Set-MappedKeyHandlers
+# Register key handlers if PSReadLine is available
+if (Get-Module -Name PSReadLine) {
+	Set-MappedKeyHandlers
+}
             ";
 
             return ExecutePSCommandAsync(new PSCommand().AddScript(shellIntegrationScript), cancellationToken);