diff options
| author | Ellie Huxtable <ellie@atuin.sh> | 2026-03-09 14:28:32 -0700 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2026-03-09 14:28:32 -0700 |
| commit | b4a17e4346c97d837d0ee3a3a55c5ceca789a3e8 (patch) | |
| tree | 4be327a9f902455a870232d36e2cd4fb4206804d /crates/atuin-hex/src/lib.rs | |
| parent | chore: update to Rust 1.94 (#3247) (diff) | |
| download | atuin-b4a17e4346c97d837d0ee3a3a55c5ceca789a3e8.zip | |
feat: use pty proxy for rendering tui popups without clearing the terminal (#3234)
It feels much, much nicer this way. This has also been asked for pretty
consistently since we made inline rendering the default. Now we can have
everything :)
Maintains a shadow vt100 renderer so that we can restore the terminal
state upon popup close. This happens on a background thread, so our
impact on terminal performance should still be super minimal, if
anything
## Checks
- [ ] I am happy for maintainers to push small adjustments to this PR,
to speed up the review cycle
- [ ] I have checked that there are no existing pull requests for the
same thing
Diffstat (limited to 'crates/atuin-hex/src/lib.rs')
| -rw-r--r-- | crates/atuin-hex/src/lib.rs | 465 |
1 files changed, 465 insertions, 0 deletions
diff --git a/crates/atuin-hex/src/lib.rs b/crates/atuin-hex/src/lib.rs new file mode 100644 index 00000000..ff37cfe3 --- /dev/null +++ b/crates/atuin-hex/src/lib.rs @@ -0,0 +1,465 @@ +pub mod osc133; + +use clap::{Args, Subcommand, ValueEnum}; + +#[derive(Subcommand, Debug)] +pub enum Cmd { + /// Print shell code to initialize atuin-hex on shell startup + Init(Init), +} + +#[derive(Args, Debug)] +pub struct Init { + /// Shell to generate init for. If omitted, attempt auto-detection + #[arg(value_enum)] + shell: Option<Shell>, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] +#[value(rename_all = "lower")] +#[allow(clippy::enum_variant_names, clippy::doc_markdown)] +enum Shell { + /// Zsh setup + Zsh, + /// Bash setup + Bash, + /// Fish setup + Fish, +} + +impl Shell { + fn as_str(self) -> &'static str { + match self { + Self::Bash => "bash", + Self::Zsh => "zsh", + Self::Fish => "fish", + } + } +} + +impl Init { + fn run(self) -> Result<(), String> { + let shell = detect_shell(self.shell)?; + let script = render_init(shell); + print!("{script}"); + Ok(()) + } +} + +pub fn run(cmd: Option<Cmd>) { + match cmd { + Some(Cmd::Init(init)) => { + if let Err(err) = init.run() { + eprintln!("atuin hex: {err}"); + std::process::exit(1); + } + } + None => app::main(), + } +} + +fn detect_shell(cli_shell: Option<Shell>) -> Result<Shell, String> { + if let Some(shell) = cli_shell { + return Ok(shell); + } + + if let Ok(shell) = std::env::var("ATUIN_SHELL") + && let Some(shell) = shell_from_name(&shell) + { + return Ok(shell); + } + + if let Ok(shell) = std::env::var("SHELL") + && let Some(shell) = shell_from_name(&shell) + { + return Ok(shell); + } + + Err( + "could not detect a supported shell. Please specify one explicitly: bash, zsh, or fish" + .to_string(), + ) +} + +fn shell_from_name(name: &str) -> Option<Shell> { + let shell = name + .trim() + .rsplit('/') + .next() + .unwrap_or(name) + .trim_start_matches('-') + .to_ascii_lowercase(); + + match shell.as_str() { + "bash" => Some(Shell::Bash), + "zsh" => Some(Shell::Zsh), + "fish" => Some(Shell::Fish), + _ => None, + } +} + +fn init_command(shell: Shell) -> String { + format!("atuin init {}", shell.as_str()) +} + +fn render_init(shell: Shell) -> String { + let init_command = init_command(shell); + + match shell { + Shell::Bash | Shell::Zsh => format!( + r#"if [[ "$-" == *i* ]] && [[ -t 0 ]] && [[ -t 1 ]]; then + _atuin_hex_tmux_current="${{TMUX:-}}" + _atuin_hex_tmux_previous="${{ATUIN_HEX_TMUX:-}}" + + if [[ -z "${{ATUIN_HEX_ACTIVE:-}}" ]] || [[ "$_atuin_hex_tmux_current" != "$_atuin_hex_tmux_previous" ]]; then + export ATUIN_HEX_ACTIVE=1 + export ATUIN_HEX_TMUX="$_atuin_hex_tmux_current" + exec atuin hex + fi + + unset _atuin_hex_tmux_current _atuin_hex_tmux_previous +fi + +eval "$({init_command})" +"# + ), + Shell::Fish => format!( + r#"if status is-interactive; and test -t 0; and test -t 1 + set -l _atuin_hex_tmux_current "" + if set -q TMUX + set _atuin_hex_tmux_current "$TMUX" + end + + set -l _atuin_hex_tmux_previous "" + if set -q ATUIN_HEX_TMUX + set _atuin_hex_tmux_previous "$ATUIN_HEX_TMUX" + end + + if not set -q ATUIN_HEX_ACTIVE + set -gx ATUIN_HEX_ACTIVE 1 + set -gx ATUIN_HEX_TMUX "$_atuin_hex_tmux_current" + exec atuin hex + else if test "$_atuin_hex_tmux_current" != "$_atuin_hex_tmux_previous" + set -gx ATUIN_HEX_ACTIVE 1 + set -gx ATUIN_HEX_TMUX "$_atuin_hex_tmux_current" + exec atuin hex + end +end + +{init_command} | source +"# + ), + } +} + +#[cfg(any(not(unix), target_os = "illumos"))] +mod app { + pub(crate) fn main() { + eprintln!("atuin hex currently supports unix platforms excluding illumos"); + std::process::exit(1); + } +} + +#[cfg(all(unix, not(target_os = "illumos")))] +mod app { + use std::io::{Read, Write}; + use std::os::unix::net::UnixListener; + use std::sync::mpsc; + + use crossterm::terminal; + use portable_pty::{CommandBuilder, PtySize, native_pty_system}; + + enum ParserMsg { + Data(Vec<u8>), + Resize { rows: u16, cols: u16 }, + ScreenRequest(mpsc::Sender<Vec<u8>>), + } + + pub(crate) fn main() { + if let Err(e) = run() { + let _ = terminal::disable_raw_mode(); + eprintln!("atuin hex: {e:#}"); + std::process::exit(1); + } + } + + fn socket_path() -> std::path::PathBuf { + let dir = std::env::temp_dir(); + dir.join(format!("atuin-hex-{}.sock", std::process::id())) + } + + /// Wire format written to the Unix socket: + /// + /// ```text + /// [rows: u16 BE][cols: u16 BE][cursor_row: u16 BE][cursor_col: u16 BE] + /// [row_0_len: u32 BE][row_0_bytes...] + /// [row_1_len: u32 BE][row_1_bytes...] + /// ... + /// ``` + /// + /// Each row's bytes come from `screen.rows_formatted(0, cols)` and contain + /// pre-built ANSI escape sequences. The client can write them directly to + /// stdout without needing its own vt100 parser. + fn encode_screen(parser: &vt100::Parser) -> Vec<u8> { + let screen = parser.screen(); + let (rows, cols) = screen.size(); + let (cursor_row, cursor_col) = screen.cursor_position(); + + let mut buf: Vec<u8> = Vec::with_capacity(256 + (rows as usize * cols as usize)); + buf.extend_from_slice(&rows.to_be_bytes()); + buf.extend_from_slice(&cols.to_be_bytes()); + buf.extend_from_slice(&cursor_row.to_be_bytes()); + buf.extend_from_slice(&cursor_col.to_be_bytes()); + + for row_bytes in screen.rows_formatted(0, cols) { + let len = row_bytes.len() as u32; + buf.extend_from_slice(&len.to_be_bytes()); + buf.extend_from_slice(&row_bytes); + } + + buf + } + + fn handle_parser_msg(parser: &mut vt100::Parser, msg: ParserMsg) { + match msg { + ParserMsg::Data(data) => parser.process(&data), + ParserMsg::Resize { rows, cols } => parser.set_size(rows, cols), + ParserMsg::ScreenRequest(reply_tx) => { + let _ = reply_tx.send(encode_screen(parser)); + } + } + } + + fn run() -> eyre::Result<()> { + let (cols, rows) = terminal::size()?; + + let pty_system = native_pty_system(); + let pair = pty_system + .openpty(PtySize { + rows, + cols, + pixel_width: 0, + pixel_height: 0, + }) + .map_err(|e| eyre::eyre!("{e:#}"))?; + + // Set up socket path and expose it to child processes + let sock_path = socket_path(); + // Clean up any stale socket from a previous crash + let _ = std::fs::remove_file(&sock_path); + + let mut cmd = CommandBuilder::new_default_prog(); + cmd.cwd(std::env::current_dir()?); + cmd.env("ATUIN_HEX_SOCKET", sock_path.as_os_str()); + + let mut child = pair + .slave + .spawn_command(cmd) + .map_err(|e| eyre::eyre!("{e:#}"))?; + + // Close slave side in parent process + drop(pair.slave); + + let mut pty_reader = pair + .master + .try_clone_reader() + .map_err(|e| eyre::eyre!("{e:#}"))?; + let mut pty_writer = pair + .master + .take_writer() + .map_err(|e| eyre::eyre!("{e:#}"))?; + + // Channel: stdout/sigwinch/socket threads -> parser thread (bounded, non-blocking send) + let (msg_tx, msg_rx) = mpsc::sync_channel::<ParserMsg>(64); + + // --- Parser thread --- + // Maintains a persistent vt100::Parser fed bytes as they arrive. + // On screen request: reads current state directly (no replay). + std::thread::spawn(move || { + let mut parser = vt100::Parser::new(rows, cols, 0); + + loop { + // Block until at least one message arrives + let first = match msg_rx.recv() { + Ok(msg) => msg, + Err(_) => break, + }; + + handle_parser_msg(&mut parser, first); + + // Drain all remaining pending messages so the parser stays + // caught up during high-throughput bursts (e.g. `cat bigfile`). + // The channel holds at most 64 items, so this is bounded. + while let Ok(msg) = msg_rx.try_recv() { + handle_parser_msg(&mut parser, msg); + } + } + }); + + // --- Socket server thread --- + // Listens on Unix socket; on connection, requests screen state from parser thread. + { + let sock_path_clone = sock_path.clone(); + let screen_tx = msg_tx.clone(); + std::thread::spawn(move || { + let listener = match UnixListener::bind(&sock_path_clone) { + Ok(l) => l, + Err(e) => { + eprintln!("atuin hex: failed to bind socket: {e}"); + return; + } + }; + + for stream in listener.incoming() { + let mut stream = match stream { + Ok(s) => s, + Err(_) => break, + }; + + let (reply_tx, reply_rx) = mpsc::channel(); + if screen_tx.send(ParserMsg::ScreenRequest(reply_tx)).is_err() { + break; + } + if let Ok(data) = reply_rx.recv() { + let _ = stream.write_all(&data); + let _ = stream.flush(); + } + } + }); + } + + // Handle terminal resize via SIGWINCH + { + use signal_hook::consts::SIGWINCH; + use signal_hook::iterator::Signals; + + let master = pair.master; + let resize_tx = msg_tx.clone(); + let mut signals = Signals::new([SIGWINCH])?; + + std::thread::spawn(move || { + for _ in signals.forever() { + if let Ok((cols, rows)) = terminal::size() { + let _ = master.resize(PtySize { + rows, + cols, + pixel_width: 0, + pixel_height: 0, + }); + let _ = resize_tx.try_send(ParserMsg::Resize { rows, cols }); + } + } + }); + } + + terminal::enable_raw_mode()?; + + // PTY -> stdout (with OSC 133 parsing + buffer feed) + let stdout_thread = std::thread::spawn(move || { + let mut stdout = std::io::stdout(); + let mut parser = crate::osc133::Parser::new(); + let mut buf = [0u8; 8192]; + loop { + match pty_reader.read(&mut buf) { + Ok(0) | Err(_) => break, + Ok(n) => { + parser.push(&buf[..n], |_event| { + // Zone transitions are tracked inside the parser. + // Callers can query parser.zone() after push. + }); + + // Feed bytes to the shadow parser. Drops on backpressure — + // the screen snapshot may be stale during bursts, but + // self-corrects once output settles. + let _ = msg_tx.try_send(ParserMsg::Data(buf[..n].to_vec())); + + if stdout.write_all(&buf[..n]).is_err() { + break; + } + let _ = stdout.flush(); + } + } + } + }); + + // stdin -> PTY + std::thread::spawn(move || { + let mut stdin = std::io::stdin(); + let mut buf = [0u8; 8192]; + loop { + match stdin.read(&mut buf) { + Ok(0) | Err(_) => break, + Ok(n) => { + if pty_writer.write_all(&buf[..n]).is_err() { + break; + } + } + } + } + }); + + let status = child.wait()?; + let _ = stdout_thread.join(); + + let _ = terminal::disable_raw_mode(); + + // Clean up socket file + let _ = std::fs::remove_file(&sock_path); + + std::process::exit(process_exit_code(status.exit_code())); + } + + fn process_exit_code(code: u32) -> i32 { + i32::try_from(code).unwrap_or(1) + } + + #[cfg(test)] + mod tests { + use super::process_exit_code; + + #[test] + fn process_exit_code_preserves_valid_values() { + assert_eq!(process_exit_code(0), 0); + assert_eq!(process_exit_code(127), 127); + assert_eq!(process_exit_code(i32::MAX as u32), i32::MAX); + } + + #[test] + fn process_exit_code_defaults_when_out_of_range() { + assert_eq!(process_exit_code(i32::MAX as u32 + 1), 1); + } + } +} + +#[cfg(test)] +mod tests { + use super::{Shell, init_command, render_init, shell_from_name}; + + #[test] + fn shell_from_name_handles_paths() { + assert_eq!(shell_from_name("/bin/zsh"), Some(Shell::Zsh)); + assert_eq!(shell_from_name("/usr/local/bin/bash"), Some(Shell::Bash)); + assert_eq!(shell_from_name("fish"), Some(Shell::Fish)); + } + + #[test] + fn init_command_is_bootstrap_only() { + let command = init_command(Shell::Zsh); + assert_eq!(command, "atuin init zsh"); + } + + #[test] + fn posix_init_uses_exec_and_tmux_guard() { + let script = render_init(Shell::Bash); + assert!(script.contains("exec atuin hex")); + assert!(script.contains("ATUIN_HEX_TMUX")); + assert!(script.contains("eval \"$(atuin init bash)\"")); + } + + #[test] + fn fish_init_uses_source() { + let script = render_init(Shell::Fish); + assert!(script.contains("exec atuin hex")); + assert!(script.contains("atuin init fish | source")); + } +} |
