ansi/
marks.rs

1//! Shell-integration marks (OSC 133, and VS Code's OSC 633), found without a
2//! full parse.
3//!
4//! Shells that support "semantic prompts" bracket every command with four
5//! sequences: prompt start (`A`), command start (`B`, where the user's input
6//! begins), output start (`C`, the command ran) and command end (`D`, with an
7//! optional exit code), each `ESC ] 133 ; <letter> [; params] <BEL or ST>`.
8//! A terminal that knows where they fall can jump by command rather than by
9//! line, and show which commands failed.
10//!
11//! [`MarkScanner`] finds them in a raw byte stream, chunk by chunk, with no
12//! other state: it does not need a parser or a screen, so it can run over
13//! output nobody is looking at. A sequence split across chunks is found.
14
15/// Which boundary a mark records.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
17pub enum MarkKind {
18    /// `A`: the prompt starts.
19    PromptStart,
20    /// `B`: the prompt ends and the command line begins.
21    CommandStart,
22    /// `C`: the command was submitted; its output follows.
23    OutputStart,
24    /// `D`: the command finished.
25    CommandEnd,
26}
27
28impl MarkKind {
29    /// The protocol letter.
30    pub fn letter(self) -> u8 {
31        match self {
32            MarkKind::PromptStart => b'A',
33            MarkKind::CommandStart => b'B',
34            MarkKind::OutputStart => b'C',
35            MarkKind::CommandEnd => b'D',
36        }
37    }
38
39    fn from_letter(b: u8) -> Option<MarkKind> {
40        match b {
41            b'A' => Some(MarkKind::PromptStart),
42            b'B' => Some(MarkKind::CommandStart),
43            b'C' => Some(MarkKind::OutputStart),
44            b'D' => Some(MarkKind::CommandEnd),
45            _ => None,
46        }
47    }
48}
49
50/// One mark found in the stream.
51#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52pub struct ShellMark {
53    /// Stream offset of the `ESC` that opens the sequence.
54    pub offset: u64,
55    /// Which boundary.
56    pub kind: MarkKind,
57    /// For [`MarkKind::CommandEnd`], the exit code, when the shell sent one.
58    pub exit_code: Option<i32>,
59}
60
61/// Longest sequence kept while waiting for its terminator; longer ones are
62/// not marks we understand and are dropped.
63const MAX_SEQ: usize = 128;
64
65/// Finds shell-integration marks in a byte stream fed in chunks.
66#[derive(Debug, Default, Clone)]
67pub struct MarkScanner {
68    /// Bytes of a candidate sequence, from its `ESC`; empty when none is open.
69    pending: Vec<u8>,
70    /// Stream offset of the candidate's `ESC`.
71    start: u64,
72    /// Stream offset of the next byte fed.
73    pos: u64,
74}
75
76impl MarkScanner {
77    /// A scanner at stream offset 0.
78    pub fn new() -> Self {
79        Self::default()
80    }
81
82    /// A scanner that treats the next byte fed as stream offset `pos`.
83    pub fn at(pos: u64) -> Self {
84        MarkScanner {
85            pos,
86            ..Self::default()
87        }
88    }
89
90    /// Stream offset of the next byte.
91    pub fn position(&self) -> u64 {
92        self.pos
93    }
94
95    /// Scan `bytes`, appending every complete mark to `out`.
96    pub fn feed(&mut self, bytes: &[u8], out: &mut Vec<ShellMark>) {
97        let mut i = 0;
98        while i < bytes.len() {
99            if self.pending.is_empty() {
100                // Fast path: nothing open, skip to the next ESC.
101                match bytes[i..].iter().position(|&b| b == 0x1B) {
102                    None => {
103                        self.pos += (bytes.len() - i) as u64;
104                        return;
105                    }
106                    Some(p) => {
107                        self.pos += p as u64;
108                        i += p;
109                        self.start = self.pos;
110                        self.pending.push(0x1B);
111                        self.pos += 1;
112                        i += 1;
113                    }
114                }
115                continue;
116            }
117            let b = bytes[i];
118            self.step(b, out);
119            self.pos += 1;
120            i += 1;
121        }
122    }
123
124    fn step(&mut self, b: u8, out: &mut Vec<ShellMark>) {
125        let p = &self.pending;
126        let n = p.len();
127        // ESC \ ends a string: check before anything else.
128        if n >= 2 && p[n - 1] == 0x1B {
129            if b == b'\\' {
130                self.pending.pop();
131                self.finish(out);
132            } else {
133                self.restart_at_last_esc(b);
134            }
135            return;
136        }
137        match b {
138            0x07 if n >= 6 => self.finish(out),
139            0x1B if n >= 6 => self.pending.push(0x1B), // maybe ST
140            0x1B => {
141                // A fresh ESC restarts the candidate.
142                self.pending.clear();
143                self.pending.push(0x1B);
144                self.start = self.pos;
145            }
146            _ => {
147                let ok = match n {
148                    1 => b == b']',
149                    2 | 3 => b == b'1' || b == b'6' || b == b'3',
150                    4 => b == b'3',
151                    5 => b == b';',
152                    6 => MarkKind::from_letter(b).is_some(),
153                    _ => n < MAX_SEQ && b >= 0x20,
154                };
155                if ok {
156                    self.pending.push(b);
157                    if n == 4 && !matches!(&self.pending[2..5], b"133" | b"633") {
158                        self.pending.clear();
159                    }
160                } else {
161                    self.pending.clear();
162                }
163            }
164        }
165    }
166
167    /// The candidate had an ESC that turned out not to start ST; treat that
168    /// ESC as the start of a new candidate and feed `b` to it.
169    fn restart_at_last_esc(&mut self, b: u8) {
170        self.pending.clear();
171        self.pending.push(0x1B);
172        self.start = self.pos - 1;
173        let mut sink = Vec::new();
174        self.step(b, &mut sink);
175    }
176
177    fn finish(&mut self, out: &mut Vec<ShellMark>) {
178        // pending: ESC ] 1 3 3 ; X [; params]
179        let p = std::mem::take(&mut self.pending);
180        if p.len() < 7 {
181            return;
182        }
183        let Some(kind) = MarkKind::from_letter(p[6]) else {
184            return;
185        };
186        let params = &p[7..];
187        let exit_code = if kind == MarkKind::CommandEnd {
188            params
189                .strip_prefix(b";")
190                .and_then(|rest| rest.split(|&c| c == b';').next())
191                .and_then(|code| std::str::from_utf8(code).ok())
192                .and_then(|code| code.parse::<i32>().ok())
193        } else {
194            None
195        };
196        out.push(ShellMark {
197            offset: self.start,
198            kind,
199            exit_code,
200        });
201    }
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207
208    fn scan(bytes: &[u8]) -> Vec<ShellMark> {
209        let mut s = MarkScanner::new();
210        let mut out = Vec::new();
211        s.feed(bytes, &mut out);
212        out
213    }
214
215    fn cycle() -> Vec<u8> {
216        b"\x1b]133;A\x07$ \x1b]133;B\x07cargo test\r\n\x1b]133;C\x07running...\r\nFAILED\r\n\x1b]133;D;101\x1b\\".to_vec()
217    }
218
219    #[test]
220    fn finds_a_whole_command_cycle() {
221        let data = cycle();
222        let marks = scan(&data);
223        let kinds: Vec<_> = marks.iter().map(|m| m.kind).collect();
224        assert_eq!(
225            kinds,
226            [
227                MarkKind::PromptStart,
228                MarkKind::CommandStart,
229                MarkKind::OutputStart,
230                MarkKind::CommandEnd
231            ]
232        );
233        assert_eq!(marks[3].exit_code, Some(101));
234        for m in &marks {
235            assert_eq!(data[m.offset as usize], 0x1B, "{m:?}");
236            assert_eq!(data[m.offset as usize + 6], m.kind.letter());
237        }
238    }
239
240    #[test]
241    fn every_chunk_split_finds_the_same_marks() {
242        let mut data = b"noise \x1b[1mbold\x1b[0m ".to_vec();
243        data.extend(cycle());
244        data.extend_from_slice(b"\x1b]633;D\x07\x1b]0;title\x07tail");
245        let whole = scan(&data);
246        assert_eq!(whole.len(), 5);
247        for split in 0..=data.len() {
248            let mut s = MarkScanner::new();
249            let mut out = Vec::new();
250            s.feed(&data[..split], &mut out);
251            s.feed(&data[split..], &mut out);
252            assert_eq!(out, whole, "split at {split}");
253            assert_eq!(s.position(), data.len() as u64);
254        }
255        // One byte at a time, too.
256        let mut s = MarkScanner::new();
257        let mut out = Vec::new();
258        for b in &data {
259            s.feed(std::slice::from_ref(b), &mut out);
260        }
261        assert_eq!(out, whole);
262    }
263
264    #[test]
265    fn exit_codes_are_optional_and_parsed_when_present() {
266        let m = scan(b"\x1b]133;D\x07\x1b]133;D;0\x07\x1b]133;D;-1;aid=7\x07\x1b]133;D;x\x07");
267        let codes: Vec<_> = m.iter().map(|m| m.exit_code).collect();
268        assert_eq!(codes, [None, Some(0), Some(-1), None]);
269    }
270
271    #[test]
272    fn look_alikes_are_not_marks() {
273        let m = scan(b"\x1b]0;title\x07\x1b]1330;A\x07\x1b]133;Z\x07\x1b[133;A\x07\x1b]133A\x07");
274        assert!(m.is_empty(), "{m:?}");
275    }
276
277    #[test]
278    fn an_esc_inside_a_candidate_restarts_it() {
279        let data = b"\x1b]133\x1b]133;A\x07";
280        let m = scan(data);
281        assert_eq!(m.len(), 1);
282        assert_eq!(m[0].offset, 5);
283    }
284
285    #[test]
286    fn offsets_continue_from_a_starting_position() {
287        let mut s = MarkScanner::at(1000);
288        let mut out = Vec::new();
289        s.feed(b"xx\x1b]133;C\x07", &mut out);
290        assert_eq!(out[0].offset, 1002);
291    }
292}
ansi/marks.rs

ansi/
marks.rs
