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 CLEAR_SCREEN: &[u8] = b"\x1b[2J\x1b[H";
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(CLEAR_SCREEN.len())
65 .enumerate()
66 .filter_map(|(index, window)| (window == CLEAR_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 record(asciicast: &mut Asciicast, started: Instant, code: &str, data: impl Into<String>) {
84 asciicast.events.push(AsciicastEvent {
85 time: elapsed(started),
86 code: code.into(),
87 data: data.into(),
88 });
89}
90
91fn apply_interaction(
92 interaction: &TerminalInteraction,
93 writer: &mut dyn Write,
94 master: &dyn MasterPty,
95 parser: &mut vt100::Parser,
96 asciicast: &mut Asciicast,
97 started: Instant,
98) -> Result<(), TerminalCaptureError> {
99 if let Some(text) = &interaction.text {
100 writer
101 .write_all(text.as_bytes())
102 .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
103 writer
104 .flush()
105 .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
106 record(asciicast, started, "i", text.clone());
107 }
108 if let Some(key) = &interaction.key {
109 writer
110 .write_all(key.sequence().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", key.sequence());
116 }
117 if let Some(resize) = interaction.resize {
118 resize_terminal(master, parser, resize)?;
119 record(
120 asciicast,
121 started,
122 "r",
123 format!("{}x{}", resize.cols, resize.rows),
124 );
125 }
126 Ok(())
127}
128
129fn resize_terminal(
130 master: &dyn MasterPty,
131 parser: &mut vt100::Parser,
132 resize: TerminalResize,
133) -> Result<(), TerminalCaptureError> {
134 master
135 .resize(PtySize {
136 rows: resize.rows,
137 cols: resize.cols,
138 pixel_width: 0,
139 pixel_height: 0,
140 })
141 .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
142 parser.set_size(resize.rows, resize.cols);
143 Ok(())
144}
145
146fn asciicast(options: &TerminalCaptureOptions) -> Asciicast {
147 let mut env = HashMap::new();
148 env.insert("SHELL".into(), options.file.clone());
149 env.insert(
150 "TERM".into(),
151 options
152 .env
153 .get("TERM")
154 .cloned()
155 .unwrap_or_else(|| "xterm-256color".into()),
156 );
157 Asciicast {
158 header: AsciicastHeader {
159 version: 2,
160 width: options.cols,
161 height: options.rows,
162 timestamp: chrono::Utc::now().timestamp(),
163 env,
164 },
165 events: Vec::new(),
166 }
167}
168
169fn spawn_reader(mut reader: Box<dyn Read + Send>) -> mpsc::Receiver<Vec<u8>> {
170 let (sender, receiver) = mpsc::channel();
171 std::thread::spawn(move || {
172 let mut buffer = [0_u8; 8192];
173 loop {
174 match reader.read(&mut buffer) {
175 Ok(0) | Err(_) => break,
176 Ok(length) => {
177 if sender.send(buffer[..length].to_vec()).is_err() {
178 break;
179 }
180 }
181 }
182 }
183 });
184 receiver
185}
186
187fn capture_result(
188 status: portable_pty::ExitStatus,
189 output: String,
190 frames: Vec<TerminalFrame>,
191 interaction_count: usize,
192 asciicast: Asciicast,
193) -> TerminalCapture {
194 TerminalCapture {
195 exit_code: status.exit_code() as i32,
196 signal: status.signal().map(str::to_owned),
197 transcript: unroll_terminal_frames(&frames),
198 output,
199 frames,
200 interaction_count,
201 asciicast,
202 }
203}
204
205pub fn capture_terminal(
206 options: TerminalCaptureOptions,
207) -> Result<TerminalCapture, TerminalCaptureError> {
208 if options.file.is_empty() {
209 return Err(TerminalCaptureError::new(
210 "capture_terminal requires a file",
211 None,
212 ));
213 }
214 let pty = native_pty_system()
215 .openpty(PtySize {
216 rows: options.rows,
217 cols: options.cols,
218 pixel_width: 0,
219 pixel_height: 0,
220 })
221 .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
222 let mut command = CommandBuilder::new(&options.file);
223 command.args(&options.args);
224 if let Some(cwd) = &options.cwd {
225 command.cwd(cwd);
226 }
227 command.env(
228 "TERM",
229 options
230 .env
231 .get("TERM")
232 .map_or("xterm-256color", String::as_str),
233 );
234 for (name, value) in &options.env {
235 command.env(name, value);
236 }
237 let mut child = pty
238 .slave
239 .spawn_command(command)
240 .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
241 drop(pty.slave);
242 let reader = pty
243 .master
244 .try_clone_reader()
245 .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
246 let mut writer = pty
247 .master
248 .take_writer()
249 .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
250 let receiver = spawn_reader(reader);
251 let started = Instant::now();
252 let mut parser = vt100::Parser::new(options.rows, options.cols, 100_000);
253 let mut recording = asciicast(&options);
254 let mut output = String::new();
255 let mut frames = Vec::new();
256 let mut interaction_index = 0;
257 let mut last_output = None;
258 let mut dirty = false;
259 let mut reader_closed = false;
260 let mut status = None;
261 let mut timed_out = false;
262 let mut stop_deadline = None;
263
264 loop {
265 match receiver.recv_timeout(Duration::from_millis(5)) {
266 Ok(data) => {
267 let text = String::from_utf8_lossy(&data);
268 output.push_str(&text);
269 record(&mut recording, started, "o", text.into_owned());
270 let segments = render_segments(&data);
271 let segment_count = segments.len();
272 for (index, segment) in segments.into_iter().enumerate() {
273 parser.process(segment);
274 if index + 1 < segment_count {
275 append_frame(&mut frames, &parser, started);
276 }
277 }
278 last_output = Some(Instant::now());
279 dirty = true;
280 if options
281 .stop_marker
282 .as_ref()
283 .is_some_and(|marker| output.contains(marker))
284 && stop_deadline.is_none()
285 {
286 append_frame(&mut frames, &parser, started);
287 stop_deadline = Some(Instant::now() + options.stop_marker_grace);
288 }
289 }
290 Err(mpsc::RecvTimeoutError::Disconnected) => reader_closed = true,
291 Err(mpsc::RecvTimeoutError::Timeout) => {}
292 }
293
294 while let Some(interaction) = options.interactions.get(interaction_index) {
295 if interaction
296 .after
297 .as_ref()
298 .is_some_and(|marker| !output.contains(marker))
299 {
300 break;
301 }
302 append_frame(&mut frames, &parser, started);
303 apply_interaction(
304 interaction,
305 writer.as_mut(),
306 pty.master.as_ref(),
307 &mut parser,
308 &mut recording,
309 started,
310 )?;
311 interaction_index += 1;
312 }
313
314 if dirty && last_output.is_some_and(|instant| instant.elapsed() >= options.settle_duration)
315 {
316 append_frame(&mut frames, &parser, started);
317 dirty = false;
318 }
319 if status.is_none() {
320 status = child
321 .try_wait()
322 .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?;
323 }
324 if status.is_some() && reader_closed {
325 break;
326 }
327 if status.is_none()
328 && (started.elapsed() >= options.timeout
329 || stop_deadline.is_some_and(|deadline| Instant::now() >= deadline))
330 {
331 timed_out = started.elapsed() >= options.timeout;
332 let _ = child.kill();
333 status = Some(
334 child
335 .wait()
336 .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?,
337 );
338 }
339 }
340
341 append_frame(&mut frames, &parser, started);
342 let capture = capture_result(
343 status.expect("child status is available after capture loop"),
344 output,
345 frames,
346 interaction_index,
347 recording,
348 );
349 if let Some(directory) = &options.artifact_directory {
350 write_terminal_artifacts(
351 directory,
352 &capture.frames,
353 &capture.transcript,
354 &capture.asciicast,
355 )?;
356 }
357 if timed_out {
358 return Err(TerminalCaptureError::new(
359 format!(
360 "terminal command timed out after {} ms",
361 options.timeout.as_millis()
362 ),
363 Some(capture),
364 ));
365 }
366 Ok(capture)
367}
368
369pub async fn capture_terminal_async(
370 options: TerminalCaptureOptions,
371) -> Result<TerminalCapture, TerminalCaptureError> {
372 tokio::task::spawn_blocking(move || capture_terminal(options))
373 .await
374 .map_err(|error| TerminalCaptureError::new(error.to_string(), None))?
375}