command_stream/terminal/
capture.rs

1use super::artifacts::{unroll_terminal_frames, write_terminal_artifacts};
2use super::types::{
3    Asciicast, AsciicastEvent, AsciicastHeader, TerminalCapture, TerminalCaptureError,
4    TerminalCaptureOptions, TerminalCursor, TerminalFrame, TerminalInteraction, TerminalResize,
5};
6use portable_pty::{native_pty_system, CommandBuilder, MasterPty, PtySize};
7use std::collections::HashMap;
8use std::io::{Read, Write};
9use std::sync::mpsc;
10use std::time::{Duration, Instant};
11
12const ERASE_SCREEN: &[u8] = b"\x1b[2J";
13
14fn elapsed(started: Instant) -> f64 {
15    (started.elapsed().as_secs_f64() * 1_000_000.0).round() / 1_000_000.0
16}
17
18fn trim_trailing_blank(mut lines: Vec<String>) -> Vec<String> {
19    while lines.last().is_some_and(String::is_empty) {
20        lines.pop();
21    }
22    lines
23}
24
25fn frame(parser: &vt100::Parser, started: Instant) -> TerminalFrame {
26    let screen = parser.screen();
27    let (rows, cols) = screen.size();
28    let (cursor_y, cursor_x) = screen.cursor_position();
29    let lines = trim_trailing_blank(screen.rows(0, cols).collect());
30    TerminalFrame {
31        time: elapsed(started),
32        cols,
33        rows,
34        cursor: TerminalCursor {
35            x: cursor_x,
36            y: cursor_y,
37        },
38        alternate: screen.alternate_screen(),
39        screen: lines.clone(),
40        lines,
41    }
42}
43
44fn same_frame(left: &TerminalFrame, right: &TerminalFrame) -> bool {
45    left.cols == right.cols
46        && left.rows == right.rows
47        && left.cursor == right.cursor
48        && left.alternate == right.alternate
49        && left.lines == right.lines
50}
51
52fn append_frame(frames: &mut Vec<TerminalFrame>, parser: &vt100::Parser, started: Instant) {
53    let next = frame(parser, started);
54    if frames
55        .last()
56        .is_none_or(|previous| !same_frame(previous, &next))
57    {
58        frames.push(next);
59    }
60}
61
62fn render_segments(data: &[u8]) -> Vec<&[u8]> {
63    let positions = data
64        .windows(ERASE_SCREEN.len())
65        .enumerate()
66        .filter_map(|(index, window)| (window == ERASE_SCREEN).then_some(index))
67        .collect::<Vec<_>>();
68    if positions.is_empty() {
69        return vec![data];
70    }
71
72    let mut segments = Vec::new();
73    if positions[0] > 0 {
74        segments.push(&data[..positions[0]]);
75    }
76    for (index, position) in positions.iter().enumerate() {
77        let end = positions.get(index + 1).copied().unwrap_or(data.len());
78        segments.push(&data[*position..end]);
79    }
80    segments
81}
82
83fn drain_complete_render_data(pending: &mut Vec<u8>) -> Vec<u8> {
84    let maximum = pending.len().min(ERASE_SCREEN.len() - 1);
85    let pending_length = (1..=maximum)
86        .rev()
87        .find(|length| ERASE_SCREEN.starts_with(&pending[pending.len() - length..]))
88        .unwrap_or(0);
89    pending.drain(..pending.len() - pending_length).collect()
90}
91
92fn record(asciicast: &mut Asciicast, started: Instant, code: &str, data: impl Into<String>) {
93    asciicast.events.push(AsciicastEvent {
94        time: elapsed(started),
95        code: code.into(),
96        data: data.into(),
97    });
98}
99
100fn apply_interaction(
101    interaction: &TerminalInteraction,
102    writer: &mut dyn Write,
103    master: &dyn MasterPty,
104    parser: &mut vt100::Parser,
105    asciicast: &mut Asciicast,
106    started: Instant,
107) -> Result<(), TerminalCaptureError> {
108    if let Some(text) = &interaction.text {
109        writer
110            .write_all(text.as_bytes())
111            .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
112        writer
113            .flush()
114            .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
115        record(asciicast, started, "i", text.clone());
116    }
117    if let Some(key) = &interaction.key {
118        writer
119            .write_all(key.sequence().as_bytes())
120            .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
121        writer
122            .flush()
123            .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
124        record(asciicast, started, "i", key.sequence());
125    }
126    if let Some(resize) = interaction.resize {
127        resize_terminal(master, parser, resize)?;
128        record(
129            asciicast,
130            started,
131            "r",
132            format!("{}x{}", resize.cols, resize.rows),
133        );
134    }
135    Ok(())
136}
137
138fn resize_terminal(
139    master: &dyn MasterPty,
140    parser: &mut vt100::Parser,
141    resize: TerminalResize,
142) -> Result<(), TerminalCaptureError> {
143    master
144        .resize(PtySize {
145            rows: resize.rows,
146            cols: resize.cols,
147            pixel_width: 0,
148            pixel_height: 0,
149        })
150        .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
151    parser.set_size(resize.rows, resize.cols);
152    Ok(())
153}
154
155fn asciicast(options: &TerminalCaptureOptions) -> Asciicast {
156    let mut env = HashMap::new();
157    env.insert("SHELL".into(), options.file.clone());
158    env.insert(
159        "TERM".into(),
160        options
161            .env
162            .get("TERM")
163            .cloned()
164            .unwrap_or_else(|| "xterm-256color".into()),
165    );
166    Asciicast {
167        header: AsciicastHeader {
168            version: 2,
169            width: options.cols,
170            height: options.rows,
171            timestamp: chrono::Utc::now().timestamp(),
172            env,
173        },
174        events: Vec::new(),
175    }
176}
177
178fn spawn_reader(mut reader: Box<dyn Read + Send>) -> mpsc::Receiver<Vec<u8>> {
179    let (sender, receiver) = mpsc::channel();
180    std::thread::spawn(move || {
181        let mut buffer = [0_u8; 8192];
182        loop {
183            match reader.read(&mut buffer) {
184                Ok(0) | Err(_) => break,
185                Ok(length) => {
186                    if sender.send(buffer[..length].to_vec()).is_err() {
187                        break;
188                    }
189                }
190            }
191        }
192    });
193    receiver
194}
195
196fn capture_result(
197    status: portable_pty::ExitStatus,
198    output: String,
199    frames: Vec<TerminalFrame>,
200    interaction_count: usize,
201    asciicast: Asciicast,
202) -> TerminalCapture {
203    TerminalCapture {
204        exit_code: status.exit_code() as i32,
205        signal: status.signal().map(str::to_owned),
206        transcript: unroll_terminal_frames(&frames),
207        output,
208        frames,
209        interaction_count,
210        asciicast,
211    }
212}
213
214pub fn capture_terminal(
215    options: TerminalCaptureOptions,
216) -> Result<TerminalCapture, TerminalCaptureError> {
217    if options.file.is_empty() {
218        return Err(TerminalCaptureError::new(
219            "capture_terminal requires a file",
220            None,
221        ));
222    }
223    let pty = native_pty_system()
224        .openpty(PtySize {
225            rows: options.rows,
226            cols: options.cols,
227            pixel_width: 0,
228            pixel_height: 0,
229        })
230        .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
231    let mut command = CommandBuilder::new(&options.file);
232    command.args(&options.args);
233    if let Some(cwd) = &options.cwd {
234        command.cwd(cwd);
235    }
236    command.env(
237        "TERM",
238        options
239            .env
240            .get("TERM")
241            .map_or("xterm-256color", String::as_str),
242    );
243    for (name, value) in &options.env {
244        command.env(name, value);
245    }
246    let mut child = pty
247        .slave
248        .spawn_command(command)
249        .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
250    drop(pty.slave);
251    let reader = pty
252        .master
253        .try_clone_reader()
254        .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
255    let mut writer = pty
256        .master
257        .take_writer()
258        .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
259    let receiver = spawn_reader(reader);
260    let started = Instant::now();
261    let mut parser = vt100::Parser::new(options.rows, options.cols, 100_000);
262    let mut recording = asciicast(&options);
263    let mut output = String::new();
264    let mut frames = Vec::new();
265    let mut pending_render = Vec::new();
266    let mut terminal_has_output = false;
267    let mut interaction_index = 0;
268    let mut last_output = None;
269    let mut dirty = false;
270    let mut reader_closed = false;
271    let mut status = None;
272    let mut timed_out = false;
273    let mut stop_deadline = None;
274
275    loop {
276        match receiver.recv_timeout(Duration::from_millis(5)) {
277            Ok(data) => {
278                let text = String::from_utf8_lossy(&data);
279                output.push_str(&text);
280                record(&mut recording, started, "o", text.into_owned());
281                pending_render.extend_from_slice(&data);
282                let render_data = drain_complete_render_data(&mut pending_render);
283                let segments = render_segments(&render_data);
284                let segment_count = segments.len();
285                if terminal_has_output && render_data.starts_with(ERASE_SCREEN) {
286                    append_frame(&mut frames, &parser, started);
287                }
288                for (index, segment) in segments.into_iter().enumerate() {
289                    parser.process(segment);
290                    terminal_has_output |= !segment.is_empty();
291                    if index + 1 < segment_count {
292                        append_frame(&mut frames, &parser, started);
293                    }
294                }
295                last_output = Some(Instant::now());
296                dirty = true;
297                if options
298                    .stop_marker
299                    .as_ref()
300                    .is_some_and(|marker| output.contains(marker))
301                    && stop_deadline.is_none()
302                {
303                    append_frame(&mut frames, &parser, started);
304                    stop_deadline = Some(Instant::now() + options.stop_marker_grace);
305                }
306            }
307            Err(mpsc::RecvTimeoutError::Disconnected) => reader_closed = true,
308            Err(mpsc::RecvTimeoutError::Timeout) => {}
309        }
310
311        while let Some(interaction) = options.interactions.get(interaction_index) {
312            if interaction
313                .after
314                .as_ref()
315                .is_some_and(|marker| !output.contains(marker))
316            {
317                break;
318            }
319            append_frame(&mut frames, &parser, started);
320            apply_interaction(
321                interaction,
322                writer.as_mut(),
323                pty.master.as_ref(),
324                &mut parser,
325                &mut recording,
326                started,
327            )?;
328            interaction_index += 1;
329        }
330
331        if dirty && last_output.is_some_and(|instant| instant.elapsed() >= options.settle_duration)
332        {
333            append_frame(&mut frames, &parser, started);
334            dirty = false;
335        }
336        if status.is_none() {
337            status = child
338                .try_wait()
339                .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
340        }
341        if status.is_some() && reader_closed {
342            break;
343        }
344        if status.is_none()
345            && (started.elapsed() >= options.timeout
346                || stop_deadline.is_some_and(|deadline| Instant::now() >= deadline))
347        {
348            timed_out = started.elapsed() >= options.timeout;
349            let _ = child.kill();
350            status = Some(
351                child
352                    .wait()
353                    .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?,
354            );
355        }
356    }
357
358    parser.process(&pending_render);
359    append_frame(&mut frames, &parser, started);
360    let capture = capture_result(
361        status.expect("child status is available after capture loop"),
362        output,
363        frames,
364        interaction_index,
365        recording,
366    );
367    if let Some(directory) = &options.artifact_directory {
368        write_terminal_artifacts(
369            directory,
370            &capture.frames,
371            &capture.transcript,
372            &capture.asciicast,
373        )?;
374    }
375    if timed_out {
376        return Err(TerminalCaptureError::new(
377            format!(
378                "terminal command timed out after {} ms",
379                options.timeout.as_millis()
380            ),
381            Some(capture),
382        ));
383    }
384    Ok(capture)
385}
386
387pub async fn capture_terminal_async(
388    options: TerminalCaptureOptions,
389) -> Result<TerminalCapture, TerminalCaptureError> {
390    tokio::task::spawn_blocking(move || capture_terminal(options))
391        .await
392        .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?
393}
command_stream/terminal/capture.rs

command_stream/terminal/
capture.rs
