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