fix(powershell): add error safety and cleanup (#3040)

This PR does a few things:

- Includes a fix by @sususu98 (added as co-author):
  - Fix: #3032
  - Issue: https://github.com/PowerShell/PSReadLine/issues/5047
- Adds a couple `catch` blocks for better error safety in
`PSConsoleHostReadLine`.
- Allows reloading the Atuin module on v7+ if it's already loaded, which
could be useful when updating Atuin.
- Refactors the code to clean it up.

I'll be AFK next week, so I preferred to submit this sooner than later.

Closes #3032

## Checks
- [x] I am happy for maintainers to push small adjustments to this PR,
to speed up the review cycle
- [x] I have checked that there are no existing pull requests for the
same thing

---------

Co-authored-by: sususu <suchangshan@foxmail.com>

1 files changed, 52 insertions, 30 deletions

diff --git a/crates/atuin/src/shell/atuin.ps1 b/crates/atuin/src/shell/atuin.ps1
index 501c1145..37753f8a 100644
--- a/crates/atuin/src/shell/atuin.ps1
+++ b/crates/atuin/src/shell/atuin.ps1
@@ -10,8 +10,13 @@
 #   It is initialized from the current prompt line count if not set when the first Atuin search is performed.
 
 if (Get-Module Atuin -ErrorAction Ignore) {
-    Write-Warning "The Atuin module is already loaded."
-    return
+    if ($PSVersionTable.PSVersion.Major -ge 7) {
+        Write-Warning "The Atuin module is already loaded, replacing it."
+        Remove-Module Atuin
+    } else {
+        Write-Warning "The Atuin module is already loaded, skipping."
+        return
+    }
 }
 
 if (!(Get-Command atuin -ErrorAction Ignore)) {
@@ -33,6 +38,19 @@ New-Module -Name Atuin -ScriptBlock {
     # The ReadLine overloads changed with breaking changes over time, make sure the one we expect is available.
     $script:hasExpectedReadLineOverload = ([Microsoft.PowerShell.PSConsoleReadLine]::ReadLine).OverloadDefinitions.Contains("static string ReadLine(runspace runspace, System.Management.Automation.EngineIntrinsics engineIntrinsics, System.Threading.CancellationToken cancellationToken, System.Nullable[bool] lastRunStatus)")
 
+    function Get-CommandLine {
+        $commandLine = ""
+        [Microsoft.PowerShell.PSConsoleReadLine]::GetBufferState([ref]$commandLine, [ref]$null)
+        return $commandLine
+    }
+
+    function Set-CommandLine {
+        param([string]$Text)
+
+        $commandLine = Get-CommandLine
+        [Microsoft.PowerShell.PSConsoleReadLine]::Replace(0, $commandLine.Length, $Text)
+    }
+
     # This function name is called by PSReadLine to read the next command line to execute.
     # We replace it with a custom implementation which adds Atuin support.
     function PSConsoleHostReadLine {
@@ -47,25 +65,32 @@ New-Module -Name Atuin -ScriptBlock {
         ## 2. Report the status of the previous command to Atuin (atuin history end).
 
         if ($script:atuinHistoryId) {
-            # The duration is not recorded in old PowerShell versions, let Atuin handle it. $null arguments are ignored.
-            $duration = (Get-History -Count 1).Duration.Ticks * 100
-            $durationArg = if ($duration) { "--duration=$duration" } else { $null }
+            try {
+                # The duration is not recorded in old PowerShell versions, let Atuin handle it. $null arguments are ignored.
+                $duration = (Get-History -Count 1).Duration.Ticks * 100
+                $durationArg = if ($duration) { "--duration=$duration" } else { $null }
 
-            # Fire and forget the atuin history end command to avoid blocking the shell during a potential sync.
-            $process = New-Object System.Diagnostics.Process
-            $process.StartInfo.FileName = "atuin"
-            $process.StartInfo.Arguments = "history end --exit=$exitCode $durationArg -- $script:atuinHistoryId"
-            $process.StartInfo.UseShellExecute = $false
-            $process.StartInfo.CreateNoWindow = $true
-            $process.StartInfo.RedirectStandardInput = $true
-            $process.StartInfo.RedirectStandardOutput = $true
-            $process.StartInfo.RedirectStandardError = $true
-            $process.Start() | Out-Null
-            $process.StandardInput.Close()
-            $process.BeginOutputReadLine()
-            $process.BeginErrorReadLine()
-
-            $script:atuinHistoryId = $null
+                # Fire and forget the atuin history end command to avoid blocking the shell during a potential sync.
+                $process = New-Object System.Diagnostics.Process
+                $process.StartInfo.FileName = "atuin"
+                $process.StartInfo.Arguments = "history end --exit=$exitCode $durationArg -- $script:atuinHistoryId"
+                $process.StartInfo.UseShellExecute = $false
+                $process.StartInfo.CreateNoWindow = $true
+                $process.StartInfo.RedirectStandardInput = $true
+                $process.StartInfo.RedirectStandardOutput = $true
+                $process.StartInfo.RedirectStandardError = $true
+                $process.Start() | Out-Null
+                $process.StandardInput.Close()
+                $process.BeginOutputReadLine()
+                $process.BeginErrorReadLine()
+            }
+            catch {
+                # Ignore errors to avoid breaking the shell.
+                # An error would occur if the user removes atuin from the PATH, for instance.
+            }
+            finally {
+                $script:atuinHistoryId = $null
+            }
         }
 
         ## 3. Read the next command line to execute.
@@ -90,6 +115,9 @@ New-Module -Name Atuin -ScriptBlock {
             $env:ATUIN_COMMAND_LINE = $line
             $script:atuinHistoryId = atuin history start --command-from-env
         }
+        catch {
+            # Ignore errors to avoid breaking the shell, see above.
+        }
         finally {
             $env:ATUIN_COMMAND_LINE = $null
         }
@@ -106,12 +134,9 @@ New-Module -Name Atuin -ScriptBlock {
         try {
             [System.Console]::OutputEncoding = [System.Text.Encoding]::UTF8
 
-            $query = $null
-            [Microsoft.PowerShell.PSConsoleReadLine]::GetBufferState([ref]$query, [ref]$null)
-
             # Atuin is started through Start-Process to avoid interfering with the current shell.
             $env:ATUIN_SHELL = "powershell"
-            $env:ATUIN_QUERY = $query
+            $env:ATUIN_QUERY = Get-CommandLine
             $argString = "search -i --result-file ""$resultFile"" $ExtraArgs"
             Start-Process -PassThru -NoNewWindow -FilePath atuin -ArgumentList $argString | Wait-Process
             $suggestion = (Get-Content -Raw $resultFile -Encoding UTF8 | Out-String).Trim()
@@ -141,12 +166,10 @@ New-Module -Name Atuin -ScriptBlock {
             $acceptPrefix = "__atuin_accept__:"
 
             if ( $suggestion.StartsWith($acceptPrefix)) {
-                [Microsoft.PowerShell.PSConsoleReadLine]::RevertLine()
-                [Microsoft.PowerShell.PSConsoleReadLine]::Insert($suggestion.Substring($acceptPrefix.Length))
+                Set-CommandLine $suggestion.Substring($acceptPrefix.Length)
                 [Microsoft.PowerShell.PSConsoleReadLine]::AcceptLine()
             } else {
-                [Microsoft.PowerShell.PSConsoleReadLine]::RevertLine()
-                [Microsoft.PowerShell.PSConsoleReadLine]::Insert($suggestion)
+                Set-CommandLine $suggestion
             }
         }
         finally {
@@ -168,8 +191,7 @@ New-Module -Name Atuin -ScriptBlock {
 
         if ($UpArrow) {
             Set-PSReadLineKeyHandler -Chord "UpArrow" -BriefDescription "Runs Atuin search" -ScriptBlock {
-                $line = $null
-                [Microsoft.PowerShell.PSConsoleReadLine]::GetBufferState([ref]$line, [ref]$null)
+                $line = Get-CommandLine
 
                 if (!$line.Contains("`n")) {
                     Invoke-AtuinSearch -ExtraArgs "--shell-up-key-binding"