diff options
| author | Benedikt Peetz <benedikt.peetz@b-peetz.de> | 2026-06-11 00:54:30 +0200 |
|---|---|---|
| committer | Benedikt Peetz <benedikt.peetz@b-peetz.de> | 2026-06-11 00:54:30 +0200 |
| commit | 5c39e7cf284a1f6e9a1657f2deb44e359fc47eb8 (patch) | |
| tree | c64baa8d5866c8e339eaf660dd3f94f30a3f7d8a /crates/turtle/src/atuin_pty_proxy | |
| parent | chore: Somewhat simplify sync code (diff) | |
| download | atuin-5c39e7cf284a1f6e9a1657f2deb44e359fc47eb8.zip | |
chore: Move everything into one big crate
That helps remove duplicated code and rustc/cargo will now also show
dead code correctly.
Diffstat (limited to 'crates/turtle/src/atuin_pty_proxy')
| -rw-r--r-- | crates/turtle/src/atuin_pty_proxy/capture.rs | 467 | ||||
| -rw-r--r-- | crates/turtle/src/atuin_pty_proxy/debug.rs | 53 | ||||
| -rw-r--r-- | crates/turtle/src/atuin_pty_proxy/mod.rs | 17 | ||||
| -rw-r--r-- | crates/turtle/src/atuin_pty_proxy/osc133.rs | 900 | ||||
| -rw-r--r-- | crates/turtle/src/atuin_pty_proxy/pty_proxy.rs | 231 | ||||
| -rw-r--r-- | crates/turtle/src/atuin_pty_proxy/runtime.rs | 184 | ||||
| -rw-r--r-- | crates/turtle/src/atuin_pty_proxy/screen.rs | 104 |
7 files changed, 1956 insertions, 0 deletions
diff --git a/crates/turtle/src/atuin_pty_proxy/capture.rs b/crates/turtle/src/atuin_pty_proxy/capture.rs new file mode 100644 index 00000000..97ac9b8f --- /dev/null +++ b/crates/turtle/src/atuin_pty_proxy/capture.rs @@ -0,0 +1,467 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicU16, Ordering}; + +use crate::atuin_pty_proxy::osc133::{Event, Params, Parser, Zone}; + +const HISTORY_ID_PARAM: &str = "history_id"; +const SESSION_ID_PARAM: &str = "session_id"; +const MAX_OUTPUT_CAPTURE_BYTES: usize = 1024 * 1024; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CommandCapture { + pub prompt: String, + pub command: String, + pub output: String, + pub exit_code: Option<i32>, + pub history_id: Option<String>, + pub session_id: Option<String>, + pub output_truncated: bool, + pub output_observed_bytes: u64, +} + +pub type CommandCaptureSink = Box<dyn Fn(CommandCapture) + Send + 'static>; + +#[derive(Default)] +struct CaptureBuffers { + prompt: Vec<u8>, + command: Vec<u8>, + output: Vec<u8>, + output_observed_bytes: u64, + output_truncated: bool, + exit_code: Option<i32>, + history_id: Option<String>, + session_id: Option<String>, +} + +pub(crate) struct CommandCaptureTracker { + parser: Parser, + zone: Zone, + buffers: CaptureBuffers, + cols: Arc<AtomicU16>, +} + +impl CommandCaptureTracker { + pub(crate) fn new(cols: Arc<AtomicU16>) -> Self { + Self { + parser: Parser::new(), + zone: Zone::Unknown, + buffers: CaptureBuffers::default(), + cols, + } + } + + pub(crate) fn push(&mut self, data: &[u8], mut on_capture: impl FnMut(CommandCapture)) { + let mut events = Vec::new(); + self.parser + .push_located(data, |located| events.push(located)); + + let mut start = 0; + for located in events { + let marker_start = located.start_offset.min(data.len()).max(start); + let offset = located.offset.min(data.len()); + self.append(&data[start..marker_start]); + self.handle_event(located.event, &located.params, &mut on_capture); + self.zone = located.zone; + start = offset; + } + + let append_end = self + .parser + .incomplete_osc_sequence_start() + .map_or(data.len(), |sequence_start| { + sequence_start.min(data.len()).max(start) + }); + if start < append_end { + self.append(&data[start..append_end]); + } + } + + fn append(&mut self, data: &[u8]) { + match self.zone { + Zone::Prompt => self.buffers.prompt.extend_from_slice(data), + Zone::Input => self.buffers.command.extend_from_slice(data), + Zone::Output => self.append_output(data), + Zone::Unknown => {} + } + } + + fn append_output(&mut self, data: &[u8]) { + self.buffers.output_observed_bytes = self + .buffers + .output_observed_bytes + .saturating_add(data.len() as u64); + + if self.buffers.output_truncated { + return; + } + + let remaining = MAX_OUTPUT_CAPTURE_BYTES.saturating_sub(self.buffers.output.len()); + let retained = data.len().min(remaining); + self.buffers.output_truncated = retained < data.len(); + + if retained > 0 { + self.buffers.output.extend_from_slice(&data[..retained]); + } + } + + fn handle_event( + &mut self, + event: Event, + params: &Params, + on_capture: &mut impl FnMut(CommandCapture), + ) { + match event { + Event::PromptStart => { + if self.zone != Zone::Prompt { + self.buffers = CaptureBuffers::default(); + } + } + Event::CommandStart | Event::CommandExecuted => {} + Event::CommandFinished { exit_code } => { + let Some(history_id) = params.get(HISTORY_ID_PARAM).map(str::to_owned) else { + return; + }; + + if exit_code.is_some() || self.buffers.exit_code.is_none() { + self.buffers.exit_code = exit_code; + } + self.buffers.history_id = Some(history_id); + self.buffers.session_id = params.get(SESSION_ID_PARAM).map(str::to_owned); + + if let Some(capture) = self.finish_capture() { + on_capture(capture); + } + } + } + } + + fn finish_capture(&mut self) -> Option<CommandCapture> { + let buffers = std::mem::take(&mut self.buffers); + let cols = self.cols.load(Ordering::Relaxed).max(1); + let prompt = render_plain_text(&buffers.prompt, cols); + let command = render_plain_text(&buffers.command, cols) + .trim_matches(|c| c == '\r' || c == '\n') + .to_string(); + let output = render_plain_text(&buffers.output, cols); + let output_truncated = buffers.output_truncated; + let output_observed_bytes = buffers.output_observed_bytes; + let exit_code = buffers.exit_code; + let history_id = buffers.history_id; + let session_id = buffers.session_id; + + if command.is_empty() && output.is_empty() { + return None; + } + + Some(CommandCapture { + prompt, + command, + output, + exit_code, + history_id, + session_id, + output_truncated, + output_observed_bytes, + }) + } +} + +const CLEAN_TEXT_MAX_ROWS: usize = 10_000; + +fn render_plain_text(bytes: &[u8], cols: u16) -> String { + if bytes.is_empty() { + return String::new(); + } + + let cols = cols.max(1); + let mut parser = vt100::Parser::new(estimated_rows(bytes, cols), cols, 0); + parser.process(bytes); + normalize_screen_contents(&parser.screen().contents()) +} + +fn normalize_screen_contents(contents: &str) -> String { + let mut lines = contents.lines().map(str::trim_end).collect::<Vec<_>>(); + while lines.last().is_some_and(|line| line.is_empty()) { + lines.pop(); + } + lines.join("\n") +} + +fn estimated_rows(bytes: &[u8], cols: u16) -> u16 { + let newline_rows = bytes.iter().filter(|byte| **byte == b'\n').count() + 1; + let wrapped_rows = bytes.len() / cols as usize; + newline_rows + .saturating_add(wrapped_rows) + .saturating_add(1) + .clamp(1, CLEAN_TEXT_MAX_ROWS) as u16 +} + +#[cfg(test)] +mod tests { + use super::*; + + fn tracker(cols: u16) -> CommandCaptureTracker { + CommandCaptureTracker::new(Arc::new(AtomicU16::new(cols))) + } + + fn assert_no_terminal_controls(text: &str) { + assert!( + !text + .chars() + .any(|ch| ch.is_control() && ch != '\n' && ch != '\t'), + "text still contains terminal controls: {text:?}" + ); + } + + #[test] + fn command_text_collapses_terminal_echo_edits() { + assert_eq!(render_plain_text(b"e\x08echo hi", 80), "echo hi"); + assert_eq!( + render_plain_text( + b"e\x08echo\x08 \x08\x08 \x08\x08\x08e \x08\x08 \x08e\x08echo hi", + 80 + ), + "echo hi" + ); + assert_eq!(render_plain_text(b"echo hi", 80), "echo hi"); + } + + #[test] + fn text_cleaning_strips_ansi_and_terminal_controls() { + let text = render_plain_text( + b"\x1b[32mhi\x1b[0m\r\n% \r \r", + 80, + ); + + assert_eq!(text, "hi"); + assert_no_terminal_controls(&text); + } + + #[test] + fn text_cleaning_preserves_valid_utf8_after_backspace() { + let text = render_plain_text("🦀x\x08 \x08 crab".as_bytes(), 80); + + assert_eq!(text, "🦀 crab"); + assert_no_terminal_controls(&text); + } + + #[test] + fn command_text_replays_backspaces() { + let mut tracker = tracker(80); + let mut captures = Vec::new(); + + let input = + b"\x1b]133;A\x07$ \x1b]133;B\x07e\x08echo hi\r\n\x1b]133;C\x07hi\r\n\x1b]133;D;0;history_id=hist;session_id=sess\x07\x1b]133;A\x07$ "; + tracker.push(input, |capture| captures.push(capture)); + + assert_eq!(captures.len(), 1); + assert_eq!(captures[0].command, "echo hi"); + assert_eq!(captures[0].output, "hi"); + assert_no_terminal_controls(&captures[0].command); + assert_no_terminal_controls(&captures[0].output); + } + + #[test] + fn captures_complete_command() { + let mut tracker = tracker(80); + let mut captures = Vec::new(); + + tracker.push( + b"\x1b]133;A\x07$ \x1b]133;B\x07echo hi\r\n\x1b]133;C\x07hi\r\n\x1b]133;D;0;history_id=hist;session_id=sess\x07\x1b]133;A\x07$ ", + |capture| captures.push(capture), + ); + + assert_eq!( + captures, + vec![CommandCapture { + prompt: "$".to_string(), + command: "echo hi".to_string(), + output: "hi".to_string(), + exit_code: Some(0), + history_id: Some("hist".to_string()), + session_id: Some("sess".to_string()), + output_truncated: false, + output_observed_bytes: 4, + }] + ); + } + + #[test] + fn strips_ansi_and_split_markers() { + let mut tracker = tracker(80); + let mut captures = Vec::new(); + + tracker.push(b"\x1b]133;A\x07\x1b[32m%\x1b[0m ", |_| {}); + tracker.push(b"\x1b]133;B\x07ls\x1b]133;C", |_| {}); + tracker.push( + b"\x07\x1b[31mfile\x1b[0m\r\n\x1b]133;D;1;history_id=hist;session_id=sess\x07\x1b]133;A\x07% ", + |capture| { + captures.push(capture); + }, + ); + + assert_eq!( + captures, + vec![CommandCapture { + prompt: "%".to_string(), + command: "ls".to_string(), + output: "file".to_string(), + exit_code: Some(1), + history_id: Some("hist".to_string()), + session_id: Some("sess".to_string()), + output_truncated: false, + output_observed_bytes: 15, + }] + ); + } + + #[test] + fn duplicate_prompt_start_does_not_reset_prompt_capture() { + let mut tracker = tracker(80); + let mut captures = Vec::new(); + + tracker.push( + b"\x1b]133;A\x07$ \x1b]133;A\x07continued \x1b]133;B\x07echo hi\r\n\x1b]133;C\x07hi\r\n\x1b]133;D;0;history_id=hist;session_id=sess\x07\x1b]133;A\x07$ ", + |capture| captures.push(capture), + ); + + assert_eq!(captures.len(), 1); + assert_eq!(captures[0].prompt, "$ continued"); + assert_eq!(captures[0].command, "echo hi"); + assert_eq!(captures[0].output, "hi"); + } + + #[test] + fn bare_finish_without_metadata_is_ignored() { + let mut tracker = tracker(80); + let mut captures = Vec::new(); + + tracker.push(b"\x1b]133;C\x07line one\r\n\x1b]133;D;0\x07", |capture| { + captures.push(capture); + }); + + tracker.push(b"\x1b]133;A\x07$ ", |capture| captures.push(capture)); + + assert!(captures.is_empty()); + } + + #[test] + fn bare_finish_before_metadata_in_same_push_ignored() { + let mut tracker = tracker(80); + let mut captures = Vec::new(); + + tracker.push( + b"\x1b]133;C\x07line one\r\n\x1b]133;D;1\x07\x1b]133;D;0;history_id=018f;session_id=abcd\x07", + |capture| captures.push(capture), + ); + + assert_eq!(captures.len(), 1); + assert_eq!(captures[0].output, "line one"); + assert_eq!(captures[0].exit_code, Some(0)); + assert_eq!(captures[0].history_id.as_deref(), Some("018f")); + assert_eq!(captures[0].session_id.as_deref(), Some("abcd")); + } + + #[test] + fn metadata_arriving_after_bare_finish_across_pushes() { + let mut tracker = tracker(80); + let mut captures = Vec::new(); + + tracker.push(b"\x1b]133;C\x07line one\r\n\x1b]133;D;0\x07", |capture| { + captures.push(capture); + }); + tracker.push(b"\x1b]133;D;0;history_id=018f", |capture| { + captures.push(capture) + }); + + assert!(captures.is_empty()); + + tracker.push(b";session_id=abcd\x07", |capture| captures.push(capture)); + + assert_eq!(captures.len(), 1); + assert_eq!(captures[0].output, "line one"); + assert_eq!(captures[0].exit_code, Some(0)); + assert_eq!(captures[0].history_id.as_deref(), Some("018f")); + assert_eq!(captures[0].session_id.as_deref(), Some("abcd")); + } + + #[test] + fn split_finish_marker_is_not_counted_as_output() { + let mut tracker = tracker(80); + let mut captures = Vec::new(); + + tracker.push( + b"\x1b]133;C\x07line one\r\n\x1b]133;D;0;history_id=018f", + |capture| { + captures.push(capture); + }, + ); + assert!(captures.is_empty()); + + tracker.push(b";session_id=abcd\x07", |capture| captures.push(capture)); + + assert_eq!(captures.len(), 1); + assert_eq!(captures[0].output, "line one"); + assert_eq!(captures[0].output_observed_bytes, 10); + } + + #[test] + fn captures_output_with_history_metadata_from_d_marker() { + let mut tracker = tracker(80); + let mut captures = Vec::new(); + + tracker.push( + b"\x1b]133;C\x07line one\r\n\x1b]133;D;0;history_id=018f;session_id=abcd\x07", + |capture| captures.push(capture), + ); + + assert_eq!( + captures, + vec![CommandCapture { + prompt: String::new(), + command: String::new(), + output: "line one".to_string(), + exit_code: Some(0), + history_id: Some("018f".to_string()), + session_id: Some("abcd".to_string()), + output_truncated: false, + output_observed_bytes: 10, + }] + ); + } + + #[test] + fn output_capture_is_capped_and_reports_observed_bytes() { + let mut tracker = tracker(80); + let mut captures = Vec::new(); + let mut input = b"\x1b]133;C\x07".to_vec(); + input.extend(std::iter::repeat_n(b'x', MAX_OUTPUT_CAPTURE_BYTES + 10)); + input.extend_from_slice(b"\x1b]133;D;0;history_id=big;session_id=session-1\x07"); + + tracker.push(&input, |capture| captures.push(capture)); + + assert_eq!(captures.len(), 1); + assert!(captures[0].output_truncated); + assert_eq!( + captures[0].output_observed_bytes, + (MAX_OUTPUT_CAPTURE_BYTES + 10) as u64 + ); + } + + #[test] + fn resets_buffers_between_c_d_only_captures() { + let mut tracker = tracker(80); + let mut captures = Vec::new(); + + tracker.push( + b"\x1b]133;C\x07first\r\n\x1b]133;D;0;history_id=one\x07\x1b]133;C\x07second\r\n\x1b]133;D;1;history_id=two\x07", + |capture| captures.push(capture), + ); + + assert_eq!(captures.len(), 2); + assert_eq!(captures[0].output, "first"); + assert_eq!(captures[0].history_id.as_deref(), Some("one")); + assert_eq!(captures[1].output, "second"); + assert_eq!(captures[1].history_id.as_deref(), Some("two")); + } +} diff --git a/crates/turtle/src/atuin_pty_proxy/debug.rs b/crates/turtle/src/atuin_pty_proxy/debug.rs new file mode 100644 index 00000000..bf311281 --- /dev/null +++ b/crates/turtle/src/atuin_pty_proxy/debug.rs @@ -0,0 +1,53 @@ +use crate::atuin_pty_proxy::osc133::{Event, Parser}; + +pub(crate) const RESET: &[u8] = b"\x1b[0m"; + +pub(crate) struct Osc133DebugHighlighter { + parser: Parser, +} + +impl Osc133DebugHighlighter { + pub(crate) fn new() -> Self { + Self { + parser: Parser::new(), + } + } + + pub(crate) fn render(&mut self, data: &[u8]) -> Vec<u8> { + let mut events = Vec::new(); + self.parser + .push_located(data, |located| events.push(located)); + + if events.is_empty() { + return data.to_vec(); + } + + let mut rendered = Vec::with_capacity(data.len() + (events.len() * 64)); + let mut start = 0; + + for located in events { + let offset = located.offset.min(data.len()); + if offset > start { + rendered.extend_from_slice(&data[start..offset]); + } + + rendered.extend_from_slice(event_label(&located.event)); + rendered.extend_from_slice(RESET); + start = offset; + } + + rendered.extend_from_slice(&data[start..]); + rendered + } +} + +fn event_label(event: &Event) -> &'static [u8] { + match event { + Event::PromptStart => b"\x1b[1;37;45m[OSC133:A prompt]\x1b[0m", + Event::CommandStart => b"\x1b[1;30;43m[OSC133:B input]\x1b[0m", + Event::CommandExecuted => b"\x1b[1;30;46m[OSC133:C output]\x1b[0m", + Event::CommandFinished { exit_code: Some(0) } => b"\x1b[1;37;42m[OSC133:D exit=0]\x1b[0m", + Event::CommandFinished { exit_code: Some(_) } => b"\x1b[1;37;41m[OSC133:D exit!=0]\x1b[0m", + Event::CommandFinished { exit_code: None } => b"\x1b[1;37;44m[OSC133:D exit=?]\x1b[0m", + } +} diff --git a/crates/turtle/src/atuin_pty_proxy/mod.rs b/crates/turtle/src/atuin_pty_proxy/mod.rs new file mode 100644 index 00000000..612943fa --- /dev/null +++ b/crates/turtle/src/atuin_pty_proxy/mod.rs @@ -0,0 +1,17 @@ +#[cfg(unix)] +mod capture; +#[cfg(unix)] +mod debug; +#[cfg(unix)] +mod osc133; +#[cfg(unix)] +mod pty_proxy; +#[cfg(unix)] +mod runtime; +#[cfg(unix)] +mod screen; + +#[cfg(unix)] +pub use capture::{CommandCapture, CommandCaptureSink}; +#[cfg(unix)] +pub use pty_proxy::PtyProxy; diff --git a/crates/turtle/src/atuin_pty_proxy/osc133.rs b/crates/turtle/src/atuin_pty_proxy/osc133.rs new file mode 100644 index 00000000..5b70f0aa --- /dev/null +++ b/crates/turtle/src/atuin_pty_proxy/osc133.rs @@ -0,0 +1,900 @@ +//! Streaming parser for OSC 133 (FinalTerm semantic prompt) escape sequences. +//! +//! OSC 133 marks four regions of a shell interaction: +//! +//! | Marker | Meaning | +//! |--------|--------------------------------------| +//! | A | Prompt is about to be printed | +//! | B | Prompt ended — command input begins | +//! | C | Command submitted — output begins | +//! | D[;n] | Command finished with exit code *n* | +//! +//! The wire format is `ESC ] 133 ; <cmd> [; <params>] ST` where ST is BEL +//! (0x07), ESC \ (0x1B 0x5C), or C1 ST (0x9C). +//! +//! # Design goals +//! +//! * **Transparent** — the parser observes the byte stream without modifying it; +//! the caller remains responsible for forwarding bytes to their destination. +//! * **Bounded** — OSC parameter buffering is capped so malformed output cannot +//! grow memory without limit. +//! * **Non-blocking** — [`Parser::push`] processes whatever bytes are available +//! and returns immediately. +//! * **Extensible** — marker parameters are preserved so Atuin-specific metadata +//! can ride alongside standard OSC 133 markers. + +/// Events emitted when an OSC 133 marker is detected. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Event { + /// `ESC ] 133 ; A ST` — the shell is about to display its prompt. + PromptStart, + /// `ESC ] 133 ; B ST` — the prompt has ended; the user may type a command. + CommandStart, + /// `ESC ] 133 ; C ST` — the command has been submitted for execution. + CommandExecuted, + /// `ESC ] 133 ; D [; <exit_code>] ST` — command output is complete. + CommandFinished { + /// The exit code reported after the `;`, if present and valid. + exit_code: Option<i32>, + }, +} + +/// Parameters attached to an OSC 133 marker. +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub struct Params { + items: Vec<Param>, +} + +impl Params { + /// Iterate over all marker parameters in order. + #[cfg(test)] + #[inline] + pub fn iter(&self) -> impl Iterator<Item = &Param> { + self.items.iter() + } + + /// Return the value for the first `key=value` parameter with this key. + #[inline] + pub fn get(&self, key: &str) -> Option<&str> { + self.items.iter().find_map(|item| match item { + Param::KeyValue { + key: item_key, + value, + } if item_key == key => Some(value.as_str()), + Param::Value(_) | Param::KeyValue { .. } => None, + }) + } +} + +/// A single OSC 133 marker parameter. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Param { + /// A positional parameter without an equals sign. + Value(String), + /// A `key=value` parameter. + KeyValue { key: String, value: String }, +} + +/// An OSC 133 event with its position in the most recent input chunk. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LocatedEvent { + /// The OSC 133 event that was parsed. + pub event: Event, + /// Offset where this marker starts in the current chunk. + /// + /// If a marker started in an earlier [`Parser::push_located`] call, this is + /// `0` in the chunk that completed the marker. + pub start_offset: usize, + /// Offset immediately after this marker's terminator in the current chunk. + /// + /// If a marker spans multiple [`Parser::push_located`] calls, this is still + /// the offset in the chunk that completed the marker. + pub offset: usize, + /// The semantic zone after applying this event. + pub zone: Zone, + /// Metadata parameters attached to this marker. + pub params: Params, +} + +/// The current semantic zone as determined by the most recent OSC 133 marker. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +#[expect(dead_code)] +pub enum Zone { + /// No marker seen yet, or after a `D` marker (between commands). + #[default] + Unknown, + /// Between `A` and `B` — the shell is rendering its prompt. + Prompt, + /// Between `B` and `C` — the user is editing a command line. + Input, + /// Between `C` and `D` — command output is being produced. + Output, +} + +// --------------------------------------------------------------------------- +// Internal constants +// --------------------------------------------------------------------------- + +const ESC: u8 = 0x1B; +const BEL: u8 = 0x07; +const C1_ST: u8 = 0x9C; +const BACKSLASH: u8 = b'\\'; +const RIGHT_BRACKET: u8 = b']'; + +/// Maximum bytes we'll buffer for the OSC parameter string. This is large enough +/// for Atuin metadata such as history/session IDs while still bounding malformed +/// OSC sequences. +const PARAM_BUF_CAP: usize = 512; + +// --------------------------------------------------------------------------- +// State machine +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum State { + /// Normal pass-through. + Ground, + /// Saw ESC (0x1B). + Esc, + /// Inside an OSC sequence (`ESC ]`), accumulating parameter bytes. + OscParam, + /// Inside an OSC sequence, saw ESC — next byte decides if this is `ESC \` + /// (string terminator) or something else. + OscEsc, +} + +/// A streaming, zero-allocation parser for OSC 133 escape sequences. +/// +/// Feed arbitrary byte slices into [`Parser::push`]. The parser detects +/// OSC 133 markers and reports [`Event`]s through a caller-supplied callback +/// without modifying the data. It can sit transparently between a PTY reader +/// and stdout. +pub struct Parser { + state: State, + zone: Zone, + sequence_start: Option<usize>, + param_buf: [u8; PARAM_BUF_CAP], + param_len: usize, +} + +impl Default for Parser { + fn default() -> Self { + Self::new() + } +} + +impl Parser { + /// Create a new parser in the initial (ground / unknown-zone) state. + #[inline] + pub fn new() -> Self { + Self { + state: State::Ground, + zone: Zone::Unknown, + sequence_start: None, + param_buf: [0u8; PARAM_BUF_CAP], + param_len: 0, + } + } + + /// The current semantic zone based on markers seen so far. + #[inline] + #[expect(dead_code)] + pub fn zone(&self) -> Zone { + self.zone + } + + /// Start offset of an incomplete OSC sequence in the most recent chunk. + #[inline] + pub(crate) fn incomplete_osc_sequence_start(&self) -> Option<usize> { + matches!(self.state, State::OscParam | State::OscEsc) + .then(|| self.sequence_start.unwrap_or(0)) + } + + /// Process a chunk of bytes, calling `on_event` for every OSC 133 marker + /// found. + /// + /// All bytes in `data` should still be forwarded to the terminal by the + /// caller — this method only *observes* the stream. + #[cfg(test)] + #[inline] + pub fn push(&mut self, data: &[u8], mut on_event: impl FnMut(Event)) { + self.push_located(data, |located| on_event(located.event)); + } + + /// Process a chunk of bytes, calling `on_event` for every OSC 133 marker + /// found with its byte offset in this chunk. + /// + /// The offset points to the first byte after the marker terminator, making + /// it suitable for callers that need to split the original chunk at marker + /// boundaries. + #[inline] + pub fn push_located(&mut self, data: &[u8], mut on_event: impl FnMut(LocatedEvent)) { + self.sequence_start = (self.state != State::Ground).then_some(0); + + for (offset, &byte) in data.iter().enumerate() { + match self.state { + State::Ground => { + if byte == ESC { + self.state = State::Esc; + self.sequence_start = Some(offset); + } + } + State::Esc => { + if byte == RIGHT_BRACKET { + self.state = State::OscParam; + self.param_len = 0; + } else { + self.state = State::Ground; + self.sequence_start = None; + } + } + State::OscParam => { + if byte == BEL || byte == C1_ST { + self.dispatch(offset + 1, &mut on_event); + self.state = State::Ground; + self.sequence_start = None; + } else if byte == ESC { + self.state = State::OscEsc; + } else if self.param_len < PARAM_BUF_CAP { + self.param_buf[self.param_len] = byte; + self.param_len += 1; + } + // If param_len == PARAM_BUF_CAP we silently stop + // accumulating — dispatch will ignore non-133 sequences. + } + State::OscEsc => { + if byte == BACKSLASH { + self.dispatch(offset + 1, &mut on_event); + } + // Whether we got a valid ST or not, return to ground. + // (A new ESC ] would restart accumulation via the Ground + // -> Esc -> OscParam path on the *next* byte.) + self.state = State::Ground; + self.sequence_start = None; + } + } + } + } + + /// Inspect the accumulated parameter buffer. If it holds an OSC 133 + /// payload, emit the corresponding [`Event`] and update the zone. + #[inline] + fn dispatch(&mut self, offset: usize, on_event: &mut impl FnMut(LocatedEvent)) { + let payload = &self.param_buf[..self.param_len]; + + if payload.len() < 5 || &payload[..4] != b"133;" { + return; + } + + if payload.len() > 5 && payload[5] != b';' { + return; + } + + let metadata = payload.get(6..).unwrap_or_default(); + let cmd = payload[4]; + let (event, params) = match cmd { + b'A' => { + self.zone = Zone::Prompt; + (Event::PromptStart, parse_params(metadata)) + } + b'B' => { + self.zone = Zone::Input; + (Event::CommandStart, parse_params(metadata)) + } + b'C' => { + self.zone = Zone::Output; + (Event::CommandExecuted, parse_params(metadata)) + } + b'D' => { + let (exit_code, params) = parse_command_finished_params(metadata); + self.zone = Zone::Unknown; + (Event::CommandFinished { exit_code }, params) + } + _ => return, + }; + + on_event(LocatedEvent { + event, + start_offset: self.sequence_start.unwrap_or(0), + offset, + zone: self.zone, + params, + }); + } +} + +fn parse_command_finished_params(metadata: &[u8]) -> (Option<i32>, Params) { + if metadata.is_empty() { + return (None, Params::default()); + } + + let Some(separator) = metadata.iter().position(|byte| *byte == b';') else { + return parse_exit_code(metadata).map_or_else( + || (None, parse_params(metadata)), + |exit_code| (Some(exit_code), Params::default()), + ); + }; + + let (first, rest) = metadata.split_at(separator); + let rest = &rest[1..]; + + parse_exit_code(first).map_or_else( + || (None, parse_params(metadata)), + |exit_code| (Some(exit_code), parse_params(rest)), + ) +} + +fn parse_exit_code(code: &[u8]) -> Option<i32> { + if code.is_empty() { + return None; + } + + std::str::from_utf8(code) + .ok() + .and_then(|code| code.parse::<i32>().ok()) +} + +fn parse_params(metadata: &[u8]) -> Params { + let items = metadata + .split(|byte| *byte == b';') + .filter(|part| !part.is_empty()) + .map(parse_param) + .collect(); + + Params { items } +} + +fn parse_param(param: &[u8]) -> Param { + let param = String::from_utf8_lossy(param); + + if let Some((key, value)) = param.split_once('=') { + return Param::KeyValue { + key: key.to_string(), + value: value.to_string(), + }; + } + + Param::Value(param.into_owned()) +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + /// Collect all events from a single `push` call. + fn parse_events(data: &[u8]) -> Vec<Event> { + let mut parser = Parser::new(); + let mut events = Vec::new(); + parser.push(data, |e| events.push(e)); + events + } + + // -- Basic event detection ------------------------------------------------ + + #[test] + fn detect_prompt_start_bel() { + let data = b"\x1b]133;A\x07"; + assert_eq!(parse_events(data), vec![Event::PromptStart]); + } + + #[test] + fn detect_prompt_start_st() { + let data = b"\x1b]133;A\x1b\\"; + assert_eq!(parse_events(data), vec![Event::PromptStart]); + } + + #[test] + fn detect_command_start_bel() { + let data = b"\x1b]133;B\x07"; + assert_eq!(parse_events(data), vec![Event::CommandStart]); + } + + #[test] + fn detect_command_start_st() { + let data = b"\x1b]133;B\x1b\\"; + assert_eq!(parse_events(data), vec![Event::CommandStart]); + } + + #[test] + fn detect_command_executed_bel() { + let data = b"\x1b]133;C\x07"; + assert_eq!(parse_events(data), vec![Event::CommandExecuted]); + } + + #[test] + fn detect_command_executed_st() { + let data = b"\x1b]133;C\x1b\\"; + assert_eq!(parse_events(data), vec![Event::CommandExecuted]); + } + + #[test] + fn detect_command_finished_no_exit_code() { + let data = b"\x1b]133;D\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { exit_code: None }] + ); + } + + #[test] + fn detect_command_finished_exit_zero() { + let data = b"\x1b]133;D;0\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { exit_code: Some(0) }] + ); + } + + #[test] + fn detect_command_finished_exit_nonzero() { + let data = b"\x1b]133;D;127\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { + exit_code: Some(127) + }] + ); + } + + #[test] + fn detect_command_finished_negative_exit_code() { + let data = b"\x1b]133;D;-1\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { + exit_code: Some(-1) + }] + ); + } + + #[test] + fn detect_command_finished_exit_code_st() { + let data = b"\x1b]133;D;42\x1b\\"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { + exit_code: Some(42) + }] + ); + } + + #[test] + fn invalid_exit_code_yields_none() { + let data = b"\x1b]133;D;abc\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { exit_code: None }] + ); + } + + // -- Zone tracking -------------------------------------------------------- + + #[test] + fn zone_starts_unknown() { + let parser = Parser::new(); + assert_eq!(parser.zone(), Zone::Unknown); + } + + #[test] + fn full_zone_cycle() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push(b"\x1b]133;A\x07", |e| events.push(e)); + assert_eq!(parser.zone(), Zone::Prompt); + + parser.push(b"\x1b]133;B\x07", |e| events.push(e)); + assert_eq!(parser.zone(), Zone::Input); + + parser.push(b"\x1b]133;C\x07", |e| events.push(e)); + assert_eq!(parser.zone(), Zone::Output); + + parser.push(b"\x1b]133;D;0\x07", |e| events.push(e)); + assert_eq!(parser.zone(), Zone::Unknown); + + assert_eq!( + events, + vec![ + Event::PromptStart, + Event::CommandStart, + Event::CommandExecuted, + Event::CommandFinished { exit_code: Some(0) }, + ] + ); + } + + // -- Multiple events in one push ------------------------------------------ + + #[test] + fn multiple_events_single_push() { + let data = b"\x1b]133;A\x07$ \x1b]133;B\x07ls\n\x1b]133;C\x07file.txt\n\x1b]133;D;0\x07"; + let events = parse_events(data); + assert_eq!( + events, + vec![ + Event::PromptStart, + Event::CommandStart, + Event::CommandExecuted, + Event::CommandFinished { exit_code: Some(0) }, + ] + ); + } + + // -- Split across push boundaries ----------------------------------------- + + #[test] + fn split_esc_and_bracket() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push(b"\x1b", |e| events.push(e)); + assert!(events.is_empty()); + + parser.push(b"]133;A\x07", |e| events.push(e)); + assert_eq!(events, vec![Event::PromptStart]); + } + + #[test] + fn split_mid_param() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push(b"\x1b]13", |e| events.push(e)); + assert!(events.is_empty()); + + parser.push(b"3;D;42\x07", |e| events.push(e)); + assert_eq!( + events, + vec![Event::CommandFinished { + exit_code: Some(42) + }] + ); + } + + #[test] + fn split_before_terminator() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push(b"\x1b]133;B", |e| events.push(e)); + assert!(events.is_empty()); + + parser.push(b"\x07", |e| events.push(e)); + assert_eq!(events, vec![Event::CommandStart]); + } + + #[test] + fn split_esc_backslash_terminator() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push(b"\x1b]133;C\x1b", |e| events.push(e)); + assert!(events.is_empty()); + + parser.push(b"\\", |e| events.push(e)); + assert_eq!(events, vec![Event::CommandExecuted]); + } + + // -- Interleaved normal text ---------------------------------------------- + + #[test] + fn normal_text_before_and_after() { + let data = b"hello world\x1b]133;A\x07prompt text\x1b]133;B\x07command"; + let events = parse_events(data); + assert_eq!(events, vec![Event::PromptStart, Event::CommandStart]); + } + + // -- Non-133 OSC sequences (should be ignored) ---------------------------- + + #[test] + fn non_133_osc_ignored() { + let data = b"\x1b]0;window title\x07\x1b]133;A\x07"; + let events = parse_events(data); + assert_eq!(events, vec![Event::PromptStart]); + } + + #[test] + fn osc_7_ignored() { + let data = b"\x1b]7;file:///home/user\x07"; + assert!(parse_events(data).is_empty()); + } + + // -- Unknown command letter ----------------------------------------------- + + #[test] + fn unknown_command_ignored() { + let data = b"\x1b]133;Z\x07"; + assert!(parse_events(data).is_empty()); + } + + #[test] + fn marker_with_unexpected_trailing_bytes_ignored() { + let data = b"\x1b]133;ABC\x07"; + assert!(parse_events(data).is_empty()); + } + + // -- Malformed sequences -------------------------------------------------- + + #[test] + fn esc_followed_by_non_bracket() { + let data = b"\x1b[31m\x1b]133;A\x07"; + let events = parse_events(data); + assert_eq!(events, vec![Event::PromptStart]); + } + + #[test] + fn lone_esc_at_end_of_chunk() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push(b"\x1b", |e| events.push(e)); + assert!(events.is_empty()); + + // Feed non-bracket to abort the escape, then a real sequence. + parser.push(b"x\x1b]133;A\x07", |e| events.push(e)); + assert_eq!(events, vec![Event::PromptStart]); + } + + #[test] + fn truncated_133_prefix() { + // "13" followed by terminator — not "133;" so no event. + let data = b"\x1b]13\x07"; + assert!(parse_events(data).is_empty()); + } + + #[test] + fn empty_osc() { + let data = b"\x1b]\x07"; + assert!(parse_events(data).is_empty()); + } + + // -- Buffer overflow (very long non-133 OSC) ------------------------------ + + #[test] + fn very_long_osc_does_not_panic() { + let mut data = Vec::new(); + data.extend_from_slice(b"\x1b]"); + data.extend(std::iter::repeat_n(b'x', 1000)); + data.push(BEL); + // Should not panic and should produce no event. + assert!(parse_events(&data).is_empty()); + } + + // -- Empty input ---------------------------------------------------------- + + #[test] + fn empty_input() { + assert!(parse_events(b"").is_empty()); + } + + #[test] + fn only_normal_text() { + let data = b"just some regular terminal output\r\n"; + assert!(parse_events(data).is_empty()); + } + + // -- Repeated prompts (empty command) ------------------------------------ + + #[test] + fn repeated_prompt_cycle() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + // User hits enter on an empty prompt twice. + let data = b"\x1b]133;A\x07$ \x1b]133;B\x07\x1b]133;D\x07\x1b]133;A\x07$ \x1b]133;B\x07"; + parser.push(data, |e| events.push(e)); + + assert_eq!( + events, + vec![ + Event::PromptStart, + Event::CommandStart, + Event::CommandFinished { exit_code: None }, + Event::PromptStart, + Event::CommandStart, + ] + ); + assert_eq!(parser.zone(), Zone::Input); + } + + // -- Byte-at-a-time feeding ----------------------------------------------- + + #[test] + fn byte_at_a_time() { + let data = b"\x1b]133;D;99\x07"; + let mut parser = Parser::new(); + let mut events = Vec::new(); + + for &byte in data { + parser.push(&[byte], |e| events.push(e)); + } + + assert_eq!( + events, + vec![Event::CommandFinished { + exit_code: Some(99) + }] + ); + } + + // -- Mixed terminators ---------------------------------------------------- + + #[test] + fn mixed_bel_and_st_terminators() { + let data = b"\x1b]133;A\x07\x1b]133;B\x1b\\\x1b]133;C\x07\x1b]133;D;1\x1b\\"; + let events = parse_events(data); + assert_eq!( + events, + vec![ + Event::PromptStart, + Event::CommandStart, + Event::CommandExecuted, + Event::CommandFinished { exit_code: Some(1) }, + ] + ); + } + + #[test] + fn detects_c1_st_terminator() { + let data = b"\x1b]133;A\x9c"; + assert_eq!(parse_events(data), vec![Event::PromptStart]); + } + + // -- Located event offsets ------------------------------------------------ + + #[test] + fn located_event_reports_offset_after_marker() { + let data = b"before\x1b]133;A\x07prompt"; + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push_located(data, |e| events.push(e)); + + assert_eq!( + events, + vec![LocatedEvent { + event: Event::PromptStart, + start_offset: b"before".len(), + offset: b"before\x1b]133;A\x07".len(), + zone: Zone::Prompt, + params: Params::default(), + }] + ); + } + + #[test] + fn located_event_offset_is_relative_to_completing_chunk() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push_located(b"\x1b]133;", |e| events.push(e)); + parser.push_located(b"D;42\x07after", |e| events.push(e)); + + assert_eq!( + events, + vec![LocatedEvent { + event: Event::CommandFinished { + exit_code: Some(42) + }, + start_offset: 0, + offset: b"D;42\x07".len(), + zone: Zone::Unknown, + params: Params::default(), + }] + ); + } + + #[test] + fn located_event_preserves_metadata_params() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push_located( + b"\x1b]133;D;127;history_id=018f;session_id=abcd;flag\x07", + |event| events.push(event), + ); + + assert_eq!(events.len(), 1); + let event = &events[0]; + assert_eq!( + event.event, + Event::CommandFinished { + exit_code: Some(127) + } + ); + assert_eq!(event.params.get("history_id"), Some("018f")); + assert_eq!(event.params.get("session_id"), Some("abcd")); + assert!( + event + .params + .iter() + .any(|param| param == &Param::Value("flag".to_string())) + ); + } + + #[test] + fn command_finished_metadata_without_exit_code_is_preserved() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push_located(b"\x1b]133;D;history_id=018f;session_id=abcd\x07", |event| { + events.push(event); + }); + + assert_eq!(events.len(), 1); + let event = &events[0]; + assert_eq!(event.event, Event::CommandFinished { exit_code: None }); + assert_eq!(event.params.get("history_id"), Some("018f")); + assert_eq!(event.params.get("session_id"), Some("abcd")); + } + + // -- Default trait -------------------------------------------------------- + + #[test] + fn parser_default() { + let parser = Parser::default(); + assert_eq!(parser.zone(), Zone::Unknown); + } + + #[test] + fn zone_default() { + assert_eq!(Zone::default(), Zone::Unknown); + } + + // -- D with empty exit code field ----------------------------------------- + + #[test] + fn d_with_semicolon_but_empty_code() { + // "133;D;" — semicolon present but no digits. + let data = b"\x1b]133;D;\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { exit_code: None }] + ); + } + + // -- Consecutive OSC sequences without gap -------------------------------- + + #[test] + fn back_to_back_osc_no_gap() { + let data = b"\x1b]133;A\x07\x1b]133;B\x07"; + let events = parse_events(data); + assert_eq!(events, vec![Event::PromptStart, Event::CommandStart]); + } + + // -- CSI sequences interleaved (should not confuse parser) ---------------- + + #[test] + fn csi_sequences_ignored() { + // CSI (ESC [) color codes mixed with OSC 133. + let data = b"\x1b[32m\x1b]133;A\x07\x1b[0m$ \x1b]133;B\x07"; + let events = parse_events(data); + assert_eq!(events, vec![Event::PromptStart, Event::CommandStart]); + } + + // -- Large exit codes ----------------------------------------------------- + + #[test] + fn large_exit_code() { + let data = b"\x1b]133;D;2147483647\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { + exit_code: Some(i32::MAX) + }] + ); + } + + #[test] + fn overflow_exit_code_yields_none() { + let data = b"\x1b]133;D;9999999999999\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { exit_code: None }] + ); + } +} diff --git a/crates/turtle/src/atuin_pty_proxy/pty_proxy.rs b/crates/turtle/src/atuin_pty_proxy/pty_proxy.rs new file mode 100644 index 00000000..8dde6f53 --- /dev/null +++ b/crates/turtle/src/atuin_pty_proxy/pty_proxy.rs @@ -0,0 +1,231 @@ +use clap::{Args, Subcommand, ValueEnum}; + +use crate::atuin_pty_proxy::{CommandCaptureSink, runtime}; + +#[derive(Args, Debug)] +pub struct PtyProxy { + /// Highlight OSC 133 prompt, input, output, and exit-code regions + #[arg(long)] + debug_osc133: bool, + + #[command(subcommand)] + cmd: Option<Cmd>, +} + +#[derive(Subcommand, Debug)] +pub enum Cmd { + /// Print shell code to initialize atuin pty-proxy 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")] +#[expect(clippy::enum_variant_names, clippy::doc_markdown)] +enum Shell { + /// Zsh setup + Zsh, + /// Bash setup + Bash, + /// Fish setup + Fish, + /// Nu setup + Nu, +} + +pub(crate) struct RuntimeOptions { + pub(crate) debug_osc133: bool, + pub(crate) command_capture_sink: Option<CommandCaptureSink>, +} + +impl RuntimeOptions { + fn new(debug_osc133: bool, command_capture_sink: Option<CommandCaptureSink>) -> Self { + Self { + debug_osc133: debug_osc133 || env_flag("ATUIN_PTY_PROXY_DEBUG"), + command_capture_sink, + } + } +} + +impl PtyProxy { + pub fn run(self, command_capture_sink: Option<CommandCaptureSink>) { + match self.cmd { + Some(Cmd::Init(init)) => { + if let Err(err) = init.run() { + eprintln!("atuin pty-proxy: {err}"); + std::process::exit(1); + } + } + None => runtime::main(RuntimeOptions::new(self.debug_osc133, command_capture_sink)), + } + } +} + +impl Init { + fn run(self) -> Result<(), String> { + let shell = detect_shell(self.shell)?; + let script = render_init(shell); + print!("{script}"); + Ok(()) + } +} + +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, fish, or nu" + .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), + "nu" => Some(Shell::Nu), + _ => None, + } +} + +fn env_flag(name: &str) -> bool { + std::env::var(name).is_ok_and(|value| { + matches!( + value.trim().to_ascii_lowercase().as_str(), + "1" | "true" | "yes" | "on" + ) + }) +} + +fn render_init(shell: Shell) -> &'static str { + match shell { + Shell::Bash | Shell::Zsh => { + r#"if [[ "$-" == *i* ]] && [[ -t 0 ]] && [[ -t 1 ]]; then + _atuin_pty_proxy_tmux_current="${TMUX:-}" + _atuin_pty_proxy_tmux_previous="${ATUIN_PTY_PROXY_TMUX:-}" + + if [[ -z "${ATUIN_PTY_PROXY_ACTIVE:-}" ]] || [[ "$_atuin_pty_proxy_tmux_current" != "$_atuin_pty_proxy_tmux_previous" ]]; then + export ATUIN_PTY_PROXY_ACTIVE=1 + export ATUIN_PTY_PROXY_TMUX="$_atuin_pty_proxy_tmux_current" + exec atuin pty-proxy + fi + + unset _atuin_pty_proxy_tmux_current _atuin_pty_proxy_tmux_previous +fi +"# + } + Shell::Fish => { + r#"if status is-interactive; and test -t 0; and test -t 1 + set -l _atuin_pty_proxy_tmux_current "" + if set -q TMUX + set _atuin_pty_proxy_tmux_current "$TMUX" + end + + set -l _atuin_pty_proxy_tmux_previous "" + if set -q ATUIN_PTY_PROXY_TMUX + set _atuin_pty_proxy_tmux_previous "$ATUIN_PTY_PROXY_TMUX" + end + + if not set -q ATUIN_PTY_PROXY_ACTIVE + set -gx ATUIN_PTY_PROXY_ACTIVE 1 + set -gx ATUIN_PTY_PROXY_TMUX "$_atuin_pty_proxy_tmux_current" + exec atuin pty-proxy + else if test "$_atuin_pty_proxy_tmux_current" != "$_atuin_pty_proxy_tmux_previous" + set -gx ATUIN_PTY_PROXY_ACTIVE 1 + set -gx ATUIN_PTY_PROXY_TMUX "$_atuin_pty_proxy_tmux_current" + exec atuin pty-proxy + end +end +"# + } + // Nushell cannot dynamically source the output of `atuin init nu`, + // so we only output the pty-proxy preamble here. Users must also set up + // `atuin init nu` separately. + Shell::Nu => { + r#"if (is-terminal --stdin) and (is-terminal --stdout) { + let tmux_current = ($env.TMUX? | default "") + let tmux_previous = ($env.ATUIN_PTY_PROXY_TMUX? | default "") + + if (($env.ATUIN_PTY_PROXY_ACTIVE? | default "") | is-empty) or ($tmux_current != $tmux_previous) { + $env.ATUIN_PTY_PROXY_ACTIVE = "1" + $env.ATUIN_PTY_PROXY_TMUX = $tmux_current + exec atuin pty-proxy + } +} +"# + } + } +} + +#[cfg(test)] +mod tests { + use super::{Shell, 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)); + assert_eq!(shell_from_name("nu"), Some(Shell::Nu)); + } + + #[test] + fn posix_init_uses_exec_and_tmux_guard() { + let script = render_init(Shell::Bash); + assert!(script.contains("exec atuin pty-proxy")); + assert!(script.contains("ATUIN_PTY_PROXY_TMUX")); + assert!(!script.contains("eval \"$(atuin init bash)\"")); + } + + #[test] + fn posix_init_has_no_double_braces() { + let script = render_init(Shell::Bash); + assert!(!script.contains("${{"), "double braces in bash init script"); + } + + #[test] + fn fish_init_uses_source() { + let script = render_init(Shell::Fish); + assert!(script.contains("exec atuin pty-proxy")); + assert!(!script.contains("atuin init fish | source")); + } + + #[test] + fn nu_init_uses_exec_and_tty_guard() { + let script = render_init(Shell::Nu); + assert!(script.contains("exec atuin pty-proxy")); + assert!(script.contains("ATUIN_PTY_PROXY_TMUX")); + assert!(script.contains("is-terminal --stdin")); + assert!(script.contains("is-terminal --stdout")); + assert!(script.contains("ATUIN_PTY_PROXY_ACTIVE")); + } +} diff --git a/crates/turtle/src/atuin_pty_proxy/runtime.rs b/crates/turtle/src/atuin_pty_proxy/runtime.rs new file mode 100644 index 00000000..37c77eef --- /dev/null +++ b/crates/turtle/src/atuin_pty_proxy/runtime.rs @@ -0,0 +1,184 @@ +use std::io::{Read, Write}; +use std::sync::Arc; +use std::sync::atomic::{AtomicU16, Ordering}; +use std::sync::mpsc; + +use crossterm::terminal; +use portable_pty::{CommandBuilder, PtySize, native_pty_system}; + +use crate::atuin_pty_proxy::capture::CommandCaptureTracker; +use crate::atuin_pty_proxy::debug::{Osc133DebugHighlighter, RESET}; +use crate::atuin_pty_proxy::pty_proxy::RuntimeOptions; +use crate::atuin_pty_proxy::screen::{self, Msg}; + +pub(crate) fn main(options: RuntimeOptions) { + if let Err(e) = run(options) { + let _ = terminal::disable_raw_mode(); + eprintln!("atuin pty-proxy: {e:#}"); + std::process::exit(1); + } +} + +fn run(options: RuntimeOptions) -> 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:#}"))?; + + let sock_path = screen::socket_path(); + let _ = std::fs::remove_file(&sock_path); + + let mut cmd = CommandBuilder::new_default_prog(); + cmd.cwd(std::env::current_dir()?); + cmd.env("ATUIN_PTY_PROXY_SOCKET", sock_path.as_os_str()); + cmd.env("ATUIN_PTY_PROXY_ACTIVE", "1"); + + let mut child = pair + .slave + .spawn_command(cmd) + .map_err(|e| eyre::eyre!("{e:#}"))?; + + 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:#}"))?; + + let (msg_tx, msg_rx) = mpsc::sync_channel::<Msg>(64); + let current_cols = Arc::new(AtomicU16::new(cols.max(1))); + + screen::spawn_parser_thread(rows, cols, msg_rx); + screen::spawn_socket_server(sock_path.clone(), msg_tx.clone()); + spawn_resize_handler(pair.master, msg_tx.clone(), current_cols.clone())?; + + terminal::enable_raw_mode()?; + + let stdout_thread = std::thread::spawn(move || { + let mut stdout = std::io::stdout(); + let mut highlighter = options.debug_osc133.then(Osc133DebugHighlighter::new); + let mut capture_tracker = options + .command_capture_sink + .as_ref() + .map(|_| CommandCaptureTracker::new(current_cols)); + let mut buf = [0u8; 8192]; + + loop { + match pty_reader.read(&mut buf) { + Ok(0) | Err(_) => break, + Ok(n) => { + if let (Some(tracker), Some(sink)) = ( + capture_tracker.as_mut(), + options.command_capture_sink.as_ref(), + ) { + tracker.push(&buf[..n], sink); + } + + if let Some(highlighter) = highlighter.as_mut() { + let rendered = highlighter.render(&buf[..n]); + let _ = msg_tx.try_send(Msg::Data(rendered.clone())); + + if stdout.write_all(&rendered).is_err() { + break; + } + } else { + let _ = msg_tx.try_send(Msg::Data(buf[..n].to_vec())); + + if stdout.write_all(&buf[..n]).is_err() { + break; + } + } + let _ = stdout.flush(); + } + } + } + + if highlighter.is_some() { + let _ = stdout.write_all(RESET); + let _ = stdout.flush(); + } + }); + + 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(); + let _ = std::fs::remove_file(&sock_path); + + std::process::exit(process_exit_code(status.exit_code())); +} + +fn spawn_resize_handler( + master: Box<dyn portable_pty::MasterPty + Send>, + resize_tx: mpsc::SyncSender<Msg>, + current_cols: Arc<AtomicU16>, +) -> eyre::Result<()> { + use signal_hook::consts::SIGWINCH; + use signal_hook::iterator::Signals; + + let mut signals = Signals::new([SIGWINCH])?; + + std::thread::spawn(move || { + for _ in signals.forever() { + if let Ok((cols, rows)) = terminal::size() { + current_cols.store(cols.max(1), Ordering::Relaxed); + let _ = master.resize(PtySize { + rows, + cols, + pixel_width: 0, + pixel_height: 0, + }); + let _ = resize_tx.try_send(Msg::Resize { rows, cols }); + } + } + }); + + Ok(()) +} + +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); + } +} diff --git a/crates/turtle/src/atuin_pty_proxy/screen.rs b/crates/turtle/src/atuin_pty_proxy/screen.rs new file mode 100644 index 00000000..5b892e21 --- /dev/null +++ b/crates/turtle/src/atuin_pty_proxy/screen.rs @@ -0,0 +1,104 @@ +use std::io::Write; +use std::os::unix::net::UnixListener; +use std::path::PathBuf; +use std::sync::mpsc::{self, Receiver, SyncSender}; + +pub(crate) enum Msg { + Data(Vec<u8>), + Resize { rows: u16, cols: u16 }, + ScreenRequest(mpsc::Sender<Vec<u8>>), +} + +pub(crate) fn socket_path() -> PathBuf { + let dir = std::env::temp_dir(); + dir.join(format!("atuin-pty-proxy-{}.sock", std::process::id())) +} + +pub(crate) fn spawn_parser_thread(rows: u16, cols: u16, msg_rx: Receiver<Msg>) { + std::thread::spawn(move || { + let mut parser = vt100::Parser::new(rows, cols, 0); + + loop { + let first = match msg_rx.recv() { + Ok(msg) => msg, + Err(_) => break, + }; + + handle_parser_msg(&mut parser, first); + + while let Ok(msg) = msg_rx.try_recv() { + handle_parser_msg(&mut parser, msg); + } + } + }); +} + +pub(crate) fn spawn_socket_server(sock_path: PathBuf, screen_tx: SyncSender<Msg>) { + std::thread::spawn(move || { + let listener = match UnixListener::bind(&sock_path) { + Ok(l) => l, + Err(e) => { + eprintln!("atuin pty-proxy: 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(Msg::ScreenRequest(reply_tx)).is_err() { + break; + } + if let Ok(data) = reply_rx.recv() { + let _ = stream.write_all(&data); + let _ = stream.flush(); + } + } + }); +} + +/// 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: Msg) { + match msg { + Msg::Data(data) => parser.process(&data), + Msg::Resize { rows, cols } => parser.screen_mut().set_size(rows, cols), + Msg::ScreenRequest(reply_tx) => { + let _ = reply_tx.send(encode_screen(parser)); + } + } +} |
