diff options
Diffstat (limited to 'crates/turtle/src/shell/atuin.ps1')
| -rw-r--r-- | crates/turtle/src/shell/atuin.ps1 | 240 |
1 files changed, 240 insertions, 0 deletions
diff --git a/crates/turtle/src/shell/atuin.ps1 b/crates/turtle/src/shell/atuin.ps1 new file mode 100644 index 00000000..431ee2c3 --- /dev/null +++ b/crates/turtle/src/shell/atuin.ps1 @@ -0,0 +1,240 @@ +# Atuin PowerShell module +# +# This should support PowerShell 5.1 (which is shipped with Windows) and later versions, on Windows and Linux. +# +# Usage: atuin init powershell | Out-String | Invoke-Expression +# +# Settings: +# - $env:ATUIN_POWERSHELL_PROMPT_OFFSET - Number of lines to offset the prompt position after exiting search. +# This is useful when using a multi-line prompt: e.g. set this to -1 when using a 2-line prompt. +# 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) { + 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)) { + Write-Error "The 'atuin' executable needs to be available in the PATH." + return +} + +if (!(Get-Module PSReadLine -ErrorAction Ignore)) { + Write-Error "Atuin requires the PSReadLine module to be installed." + return +} + +New-Module -Name Atuin -ScriptBlock { + if (-not $env:ATUIN_SESSION -or $env:ATUIN_PID -ne $PID) { + $env:ATUIN_SESSION = atuin uuid + $env:ATUIN_PID = $PID + } + + $script:atuinHistoryId = $null + $script:previousPSConsoleHostReadLine = $Function:PSConsoleHostReadLine + + # 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 { + ## 1. Collect the exit code of the previous command. + + # This needs to be done as the first thing because any script run will flush $?. + $lastRunStatus = $? + + # Exit statuses are maintained separately for native and PowerShell commands, this needs to be taken into account. + $lastNativeExitCode = $global:LASTEXITCODE + $exitCode = if ($lastRunStatus) { 0 } elseif ($lastNativeExitCode) { $lastNativeExitCode } else { 1 } + + ## 2. Report the status of the previous command to Atuin (atuin history end). + + if ($script:atuinHistoryId) { + 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() + } + 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. + + # PSConsoleHostReadLine implementation from PSReadLine, adjusted to support old versions. + Microsoft.PowerShell.Core\Set-StrictMode -Off + + $line = if ($script:hasExpectedReadLineOverload) { + # When the overload we expect is available, we can pass $lastRunStatus to it. + [Microsoft.PowerShell.PSConsoleReadLine]::ReadLine($Host.Runspace, $ExecutionContext, [System.Threading.CancellationToken]::None, $lastRunStatus) + } else { + # Either PSReadLine is older than v2.2.0-beta3, or maybe newer than we expect, so use the function from PSReadLine as-is. + & $script:previousPSConsoleHostReadLine + } + + ## 4. Report the next command line to Atuin (atuin history start). + + # PowerShell doesn't handle double quotes in native command line arguments the same way depending on its version, + # and the value of $PSNativeCommandArgumentPassing - see the about_Parsing help page which explains the breaking changes. + # This makes it unreliable, so we go through an environment variable, which should always be consistent across versions. + try { + $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 + } + + $global:LASTEXITCODE = $lastNativeExitCode + return $line + } + + function Invoke-AtuinSearch { + param([string]$ExtraArgs = "") + + $previousOutputEncoding = [System.Console]::OutputEncoding + $resultFile = New-TemporaryFile + $suggestion = "" + $errorOutput = "" + + try { + [System.Console]::OutputEncoding = [System.Text.Encoding]::UTF8 + + # Start-Process does some crazy stuff, just use the Process class directly to have more control. + $process = New-Object System.Diagnostics.Process + $process.StartInfo.FileName = "atuin" + $process.StartInfo.Arguments = "search -i --result-file ""$($resultFile.FullName)"" $ExtraArgs" + $process.StartInfo.UseShellExecute = $false + $process.StartInfo.RedirectStandardError = $true + $process.StartInfo.StandardErrorEncoding = [System.Text.Encoding]::UTF8 + $process.StartInfo.EnvironmentVariables["ATUIN_SHELL"] = "powershell" + $process.StartInfo.EnvironmentVariables["ATUIN_QUERY"] = Get-CommandLine + # PowerShell's Set-Location (cd) doesn't update the process-level working directory, set it explicitly + $process.StartInfo.WorkingDirectory = (Get-Location -PSProvider FileSystem).ProviderPath + + try { + $process.Start() | Out-Null + + # A single stream is redirected, so we can read it synchronously, but we have to start reading it + # before waiting for the process to exit, otherwise the buffer could fill up and cause a deadlock. + $errorOutput = $process.StandardError.ReadToEnd().Trim() + $process.WaitForExit() + + $suggestion = (Get-Content -LiteralPath $resultFile.FullName -Raw -Encoding UTF8 | Out-String).Trim() + } + catch { + $errorOutput = $_ + } + + if ($errorOutput) { + Write-Host -ForegroundColor Red "Atuin error:" + Write-Host -ForegroundColor DarkRed $errorOutput + } + + # If no shell prompt offset is set, initialize it from the current prompt line count. + if ($null -eq $env:ATUIN_POWERSHELL_PROMPT_OFFSET) { + try { + $promptLines = (& $Function:prompt | Out-String | Measure-Object -Line).Lines + $env:ATUIN_POWERSHELL_PROMPT_OFFSET = -1 * ($promptLines - 1) + } + catch { + $env:ATUIN_POWERSHELL_PROMPT_OFFSET = 0 + } + } + + # PSReadLine maintains its own cursor position, which will no longer be valid if Atuin scrolls the display in inline mode. + # Fortunately, InvokePrompt can receive a new Y position and reset the internal state. + $y = $Host.UI.RawUI.CursorPosition.Y + [int]$env:ATUIN_POWERSHELL_PROMPT_OFFSET + $y = [System.Math]::Max([System.Math]::Min($y, [System.Console]::BufferHeight - 1), 0) + [Microsoft.PowerShell.PSConsoleReadLine]::InvokePrompt($null, $y) + + if ($suggestion -eq "") { + # The previous input was already rendered by InvokePrompt + return + } + + $acceptPrefix = "__atuin_accept__:" + + if ( $suggestion.StartsWith($acceptPrefix)) { + Set-CommandLine $suggestion.Substring($acceptPrefix.Length) + [Microsoft.PowerShell.PSConsoleReadLine]::AcceptLine() + } else { + Set-CommandLine $suggestion + } + } + finally { + [System.Console]::OutputEncoding = $previousOutputEncoding + $resultFile.Delete() + } + } + + function Enable-AtuinSearchKeys { + param([bool]$CtrlR = $true, [bool]$UpArrow = $true) + + if ($CtrlR) { + Set-PSReadLineKeyHandler -Chord "Ctrl+r" -BriefDescription "Runs Atuin search" -ScriptBlock { + Invoke-AtuinSearch + } + } + + if ($UpArrow) { + Set-PSReadLineKeyHandler -Chord "UpArrow" -BriefDescription "Runs Atuin search" -ScriptBlock { + $line = Get-CommandLine + + if (!$line.Contains("`n")) { + Invoke-AtuinSearch -ExtraArgs "--shell-up-key-binding" + } else { + [Microsoft.PowerShell.PSConsoleReadLine]::PreviousLine() + } + } + } + } + + $ExecutionContext.SessionState.Module.OnRemove += { + $env:ATUIN_SESSION = $null + $Function:PSConsoleHostReadLine = $script:previousPSConsoleHostReadLine + } + + Export-ModuleMember -Function @("Enable-AtuinSearchKeys", "PSConsoleHostReadLine") +} | Import-Module -Global |