1use snafu::ResultExt as _;
5use tracing::Instrument as _;
6
7#[derive(Debug)]
9struct WeztermConfig {
10 scrollback: usize,
12}
13
14impl wezterm_term::TerminalConfiguration for WeztermConfig {
15 fn scrollback_size(&self) -> usize {
16 self.scrollback
17 }
18
19 fn color_palette(&self) -> wezterm_term::color::ColorPalette {
20 wezterm_term::color::ColorPalette::default()
21 }
22}
23
24#[expect(
26 clippy::exhaustive_structs,
27 reason = "
28 I just really like the ability to specify config in a struct. As if it were JSON.
29 I know that means projects depending on this struct run the risk of unexpected
30 breakage when I add a new field. But maybe we can manage those expectations by
31 making sure that all example code is based off `Config::default()`?
32 "
33)]
34#[derive(Clone)]
35pub struct Config {
36 pub width: u16,
38 pub height: u16,
40 pub command: Vec<std::ffi::OsString>,
42 pub scrollback_size: usize,
44 pub scrollback_step: usize,
46}
47
48impl Default for Config {
49 #[inline]
50 fn default() -> Self {
51 Self {
52 width: 100,
53 height: 30,
54 command: vec!["bash".into()],
55 scrollback_size: 1000,
56 scrollback_step: 5,
57 }
58 }
59}
60
61#[non_exhaustive]
64pub struct Channels {
65 pub control_tx: tokio::sync::broadcast::Sender<crate::Protocol>,
67 pub output_tx: tokio::sync::mpsc::Sender<crate::pty::BytesFromPTY>,
69 pub output_rx: tokio::sync::mpsc::Receiver<crate::pty::BytesFromPTY>,
71 pub internal_input_tx: Option<tokio::sync::mpsc::Sender<crate::pty::BytesFromSTDIN>>,
73 shadow_output: tokio::sync::mpsc::Sender<crate::output::native::Output>,
75}
76
77#[non_exhaustive]
79pub struct LastSent {
80 pub pty_sequence: usize,
82 pub pty_size: (usize, usize),
84}
85
86const CURSOR_POSITION_REQUEST: &str = "\x1b[6n";
88
89const APPLICATION_MODE_START: &str = "\x1b[?1h";
91
92const APPLICATION_MODE_END: &str = "\x1b[?1l";
94
95const TIME_TO_WAIT_FOR_MORE_PTY_OUTPUT: u64 = 1000;
97
98#[non_exhaustive]
107pub struct ShadowTerminal {
108 pub terminal: wezterm_term::Terminal,
110 pub config: Config,
112 pub channels: Channels,
114 pub accumulated_pty_output: Vec<u8>,
116 pub wait_for_output_until: Option<tokio::time::Instant>,
118 pub scroll_position: usize,
120 pub last_sent: LastSent,
122}
123
124impl ShadowTerminal {
125 #[inline]
127 pub fn new(
128 config: Config,
129 shadow_output: tokio::sync::mpsc::Sender<crate::output::native::Output>,
130 ) -> Self {
131 let (control_tx, _) = tokio::sync::broadcast::channel(64);
132 let (output_tx, output_rx) = tokio::sync::mpsc::channel(1);
133
134 tracing::debug!("Creating the in-memory Wezterm terminal");
135 let terminal = wezterm_term::Terminal::new(
136 Self::wezterm_size(config.width.into(), config.height.into()),
137 std::sync::Arc::new(WeztermConfig {
138 scrollback: config.scrollback_size,
139 }),
140 "Tattoy",
141 "O_o",
142 Box::<Vec<u8>>::default(),
143 );
144
145 let pty_size = (config.width.into(), config.height.into());
146 Self {
147 terminal,
148 config,
149 channels: Channels {
150 control_tx,
151 output_tx,
152 output_rx,
153 internal_input_tx: None,
154 shadow_output,
155 },
156 accumulated_pty_output: Vec::new(),
157 wait_for_output_until: None,
158 scroll_position: 0,
159 last_sent: LastSent {
160 pty_sequence: 0,
161 pty_size,
162 },
163 }
164 }
165
166 #[inline]
168 pub fn start(
169 &mut self,
170 user_input_rx: tokio::sync::mpsc::Receiver<crate::pty::BytesFromSTDIN>,
171 ) -> tokio::task::JoinHandle<Result<(), crate::errors::PTYError>> {
172 let (internal_input_tx, internal_input_rx) = tokio::sync::mpsc::channel(1);
173 self.channels.internal_input_tx = Some(internal_input_tx);
174
175 let pty = crate::pty::PTY {
176 command: self.config.command.clone(),
177 width: self.config.width,
178 height: self.config.height,
179 control_tx: self.channels.control_tx.clone(),
180 output_tx: self.channels.output_tx.clone(),
181 };
182
183 let current_span = tracing::Span::current();
187 tokio::spawn(async move {
188 pty.run(user_input_rx, internal_input_rx)
189 .instrument(current_span)
190 .await
191 })
192 }
193
194 #[inline]
196 pub async fn run(
197 &mut self,
198 user_input_rx: tokio::sync::mpsc::Receiver<crate::pty::BytesFromSTDIN>,
199 ) {
200 tracing::debug!("Starting Shadow Terminal loop...");
201
202 let mut control_rx = self.channels.control_tx.subscribe();
203 self.start(user_input_rx);
204
205 tracing::debug!("Starting Shadow Terminal main loop");
206 #[expect(
207 clippy::integer_division_remainder_used,
208 reason = "`tokio::select!` generates this."
209 )]
210 loop {
211 let is_wait = self.wait_for_output_until.is_some();
212 let wait_until = self.wait_for_output_until;
213 tokio::select! {
214 Some(bytes) = self.channels.output_rx.recv() => {
215 self.accumulate_pty_output(&bytes);
216 },
217 () = Self::wait_for_more_pty_output(wait_until), if is_wait => {
218 let result = self.handle_pty_output().await;
219 if let Err(error) = result {
220 tracing::error!("Handling PTY output: {error:?}");
221 }
222 }
223 Ok(message) = control_rx.recv() => {
224 self.handle_protocol_message(&message).await;
225 if matches!(message, crate::Protocol::End) {
226 break;
227 }
228 }
229 }
230 }
231
232 tracing::debug!("Shadow Terminal loop finished");
233 }
234
235 async fn wait_for_more_pty_output(maybe_wait_until: Option<tokio::time::Instant>) {
240 if let Some(wait_until) = maybe_wait_until {
241 tokio::time::sleep_until(wait_until).await;
242 }
243 }
244
245 fn accumulate_pty_output(&mut self, bytes: &crate::pty::BytesFromPTY) {
247 for byte in bytes {
249 if byte == &0 {
250 break;
251 }
252 self.accumulated_pty_output.push(*byte);
253 }
254
255 let next_output_broadcast = tokio::time::Instant::now()
256 + tokio::time::Duration::from_micros(TIME_TO_WAIT_FOR_MORE_PTY_OUTPUT);
257 self.wait_for_output_until = Some(next_output_broadcast);
258 }
259
260 fn find_subsequence(haystack: &[u8], needle: &[u8]) -> Option<usize> {
262 haystack
263 .windows(needle.len())
264 .position(|window| window == needle)
265 }
266
267 pub(crate) async fn handle_pty_output(
269 &mut self,
270 ) -> Result<(), crate::errors::ShadowTerminalError> {
271 let bytes_copy = self.accumulated_pty_output.clone();
272 let bytes = bytes_copy.as_slice();
273
274 if Self::find_subsequence(bytes, APPLICATION_MODE_START.as_bytes()).is_some() {
275 tracing::trace!("Starting terminal 'application mode'");
276 crate::output::native::raw_string_direct_to_terminal(APPLICATION_MODE_START)
277 .with_whatever_context(|err| {
278 format!("Sending 'application mode start' ANSI code: {err:?}")
279 })?;
280 }
281
282 if Self::find_subsequence(bytes, APPLICATION_MODE_END.as_bytes()).is_some() {
283 tracing::trace!("APPLICATION_MODE_END");
284 crate::output::native::raw_string_direct_to_terminal(APPLICATION_MODE_END)
285 .with_whatever_context(|err| {
286 format!("Sending 'application mode end' ANSI code: {err:?}")
287 })?;
288 }
289
290 self.handle_cursor_position_request(bytes).await?;
291 self.terminal.advance_bytes(bytes);
292 tracing::trace!("Wezterm shadow terminal advanced {} bytes", bytes.len());
293 let result = self.send_outputs().await;
294 if let Err(error) = result {
295 tracing::error!("{error:?}");
296 }
297 self.accumulated_pty_output.clear();
298 self.wait_for_output_until = None;
299 Ok(())
300 }
301
302 #[expect(
307 clippy::needless_pass_by_ref_mut,
308 reason = "
309 When I set this to `&self` then we get an actual compiler error that the `send()` method
310 on the channel is not safe because it's not `Send`. I don't understand this.
311 "
312 )]
313 async fn handle_cursor_position_request(
314 &mut self,
315 bytes: &[u8],
316 ) -> Result<(), crate::errors::ShadowTerminalError> {
317 if Self::find_subsequence(bytes, CURSOR_POSITION_REQUEST.as_bytes()).is_none() {
318 return Ok(());
319 }
320
321 let mut payload: crate::pty::BytesFromSTDIN = [0; 128];
322 let cursor_position = self.terminal.cursor_pos();
323 let response_string = format!("\x1b[{};{}R", cursor_position.y, cursor_position.x);
324 let response_bytes = response_string.as_bytes();
325
326 for chunk in response_bytes.chunks(128) {
327 crate::pty::PTY::add_bytes_to_buffer(&mut payload, chunk).with_whatever_context(
328 |error| format!("Couldn't add response to payload buffer: {error:?}"),
329 )?;
330
331 if let Some(sender) = self.channels.internal_input_tx.as_ref() {
332 tracing::debug!(
333 "Responding to cursor position request with: {}",
334 response_string.replace('\x1b', "^")
335 );
336 let result = sender.send(payload).await;
337 if let Err(error) = result {
338 snafu::whatever!("Couldn't send internal input: {error:?}");
339 }
340 }
341 }
342
343 Ok(())
344 }
345
346 async fn send_outputs(&mut self) -> Result<(), crate::errors::ShadowTerminalError> {
349 let screen_output =
350 self.build_current_output(&crate::output::native::SurfaceKind::Screen)?;
351 self.send_output(screen_output).await?;
352
353 if !self.terminal.is_alt_screen_active() {
354 let scrollback_output =
355 self.build_current_output(&crate::output::native::SurfaceKind::Scrollback)?;
356 self.send_output(scrollback_output).await?;
357 }
358
359 self.last_sent = LastSent {
360 pty_sequence: self.terminal.current_seqno(),
361 pty_size: (self.terminal.get_size().cols, self.terminal.get_size().rows),
362 };
363
364 Ok(())
365 }
366
367 #[expect(
369 clippy::needless_pass_by_ref_mut,
370 reason = "
371 Weirdly, we get the following error when `mut` is not used:
372 rustc: future cannot be sent between threads safely
373 within `shadow_terminal::ShadowTerminal`, the trait `std::marker::Sync` is not implemented for `std::cell::RefCell<termwiz::escape::parser::ParseState>`
374 if you want to do aliasing and mutation between multiple threads, use `std::sync::RwLock` instead
375 "
376 )]
377 async fn send_output(
378 &mut self,
379 output: crate::output::native::Output,
380 ) -> Result<(), crate::errors::ShadowTerminalError> {
381 let result = self.channels.shadow_output.send(output).await;
382 if let Err(error) = result {
383 tracing::error!("Sending shadow output: {error:?}");
384 return Ok(());
385 }
386
387 Ok(())
388 }
389
390 #[inline]
396 pub fn kill(&self) -> Result<(), crate::errors::ShadowTerminalError> {
397 tracing::debug!("`ShadowTerminal.kill()` called");
398
399 self.channels
400 .control_tx
401 .send(crate::Protocol::End)
402 .with_whatever_context(|err| {
403 format!("`ShadowTerminal.kill()`: Killing ShadowCouldn't write bytes into PTY's STDIN: {err:?}")
404 })?;
405
406 Ok(())
407 }
408
409 async fn handle_protocol_message(&mut self, message: &crate::Protocol) {
411 tracing::debug!("Shadow Terminal received protocol message: {message:?}");
412
413 #[expect(clippy::wildcard_enum_match_arm, reason = "It's our internal protocol")]
414 match message {
415 crate::Protocol::Resize { width, height } => {
416 self.terminal.resize(Self::wezterm_size(
417 usize::from(*width),
418 usize::from(*height),
419 ));
420 tracing::trace!("Wezterm terminal resized to: {width}x{height}");
421 }
422 crate::Protocol::Scroll(scroll) => {
423 match scroll {
424 crate::Scroll::Up => {
425 let size = self.terminal.get_size();
426 let total_lines = self.terminal.screen().scrollback_rows() - size.rows;
427
428 self.scroll_position += self.config.scrollback_step;
429 self.scroll_position = self.scroll_position.min(total_lines);
430 }
431 crate::Scroll::Down => {
432 if self.scroll_position < self.config.scrollback_step {
433 self.scroll_position = 0;
434 } else {
435 self.scroll_position -= self.config.scrollback_step;
436 }
437 }
438 crate::Scroll::Cancel => {
439 self.scroll_position = 0;
440 }
441 }
442
443 let result = self.send_outputs().await;
444 if let Err(error) = result {
445 tracing::error!("Couldn't send PTY output from shadow terminal: {error:?}");
446 }
447 }
448
449 _ => (),
450 }
451 }
452
453 const fn wezterm_size(width: usize, height: usize) -> wezterm_term::TerminalSize {
455 wezterm_term::TerminalSize {
456 cols: width,
457 rows: height,
458 pixel_width: 0,
459 pixel_height: 0,
460 dpi: 0,
461 }
462 }
463
464 #[inline]
469 pub fn resize(
470 &mut self,
471 width: u16,
472 height: u16,
473 ) -> Result<(), tokio::sync::broadcast::error::SendError<crate::Protocol>> {
474 self.channels
475 .control_tx
476 .send(crate::Protocol::Resize { width, height })?;
477 self.terminal
478 .resize(Self::wezterm_size(width.into(), height.into()));
479 Ok(())
480 }
481}
482
483impl Drop for ShadowTerminal {
484 #[inline]
485 fn drop(&mut self) {
486 tracing::trace!("Running ShadowTerminal.drop()");
487 self.kill().unwrap_or_else(|error| {
488 tracing::debug!("`ShadowTerminal.drop()`: {error:?}");
489 });
490 }
491}