Gate startup command input before TUI readiness

src/bin/dosh-client.rs

@@ -47,6 +47,7 @@ use tokio::sync::mpsc;
 
 const STREAM_INITIAL_WINDOW: usize = 1024 * 1024;
 const MAX_PENDING_USER_INPUT_BYTES: usize = 1024 * 1024;
+const STARTUP_INPUT_HOLD: Duration = Duration::from_millis(750);
 
 /// Sentinel `target_host` the server uses on a server-initiated `StreamOpen` that
 /// represents an SSH-agent connection (rather than a TCP target). The client
@@ -2842,6 +2843,7 @@ async fn run_terminal(
     let mut stream_recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>> = HashMap::new();
     let mut pending_user_input: VecDeque<Vec<u8>> = VecDeque::new();
     let mut pending_user_input_bytes = 0usize;
+    let mut startup_input_hold_until: Option<Instant> = None;
     if let Some(frame) = first_frame {
         if !forward_only {
             render_frame(&frame)?;
@@ -2858,6 +2860,7 @@ async fn run_terminal(
         && !forward_only
     {
         send_input(&socket, addr, &cred, &mut send_seq, bytes).await?;
+        startup_input_hold_until = Some(Instant::now() + STARTUP_INPUT_HOLD);
     }
 
     let (stdin_tx, mut stdin_rx) = mpsc::unbounded_channel::<Vec<u8>>();
@@ -2896,6 +2899,27 @@ async fn run_terminal(
                             break;
                         }
                         if cred.mode != "view-only" && cred.mode != "forward-only" {
+                            if let Some(deadline) = startup_input_hold_until {
+                                if !predictor.alternate_screen && Instant::now() < deadline {
+                                    queue_pending_user_input(
+                                        &mut pending_user_input,
+                                        &mut pending_user_input_bytes,
+                                        bytes,
+                                    )?;
+                                    continue;
+                                }
+                                startup_input_hold_until = None;
+                                flush_pending_user_input(
+                                    &socket,
+                                    addr,
+                                    &cred,
+                                    &mut send_seq,
+                                    &mut predictor,
+                                    &mut pending_user_input,
+                                    &mut pending_user_input_bytes,
+                                )
+                                .await?;
+                            }
                             if last_packet_at.elapsed() >= Duration::from_secs(2) {
                                 queue_pending_user_input(
                                     &mut pending_user_input,
@@ -3011,6 +3035,16 @@ async fn run_terminal(
                             if !forward_only {
                                 render_frame(&frame)?;
                                 predictor.observe_output(&frame.bytes);
+                                flush_startup_input_if_ready(
+                                    &socket,
+                                    addr,
+                                    &cred,
+                                    &mut send_seq,
+                                    &mut predictor,
+                                    (&mut pending_user_input, &mut pending_user_input_bytes),
+                                    &mut startup_input_hold_until,
+                                )
+                                .await?;
                             }
                             if frame.closed {
                                 send_ack(&socket, addr, &cred, &mut send_seq).await?;
@@ -3481,6 +3515,16 @@ async fn run_terminal(
                 // without waiting for the next keystroke to drive `redraw`.
                 if !forward_only {
                     predictor.refresh_policy()?;
+                    flush_startup_input_if_ready(
+                        &socket,
+                        addr,
+                        &cred,
+                        &mut send_seq,
+                        &mut predictor,
+                        (&mut pending_user_input, &mut pending_user_input_bytes),
+                        &mut startup_input_hold_until,
+                    )
+                    .await?;
                 }
                 // Refresh / clear the non-destructive disconnect status line based
                 // on how long the link has been silent (recomputed after any
@@ -3763,6 +3807,35 @@ async fn flush_pending_user_input(
     Ok(())
 }
 
+async fn flush_startup_input_if_ready(
+    socket: &UdpSocket,
+    addr: SocketAddr,
+    cred: &CachedCredential,
+    send_seq: &mut u64,
+    predictor: &mut Predictor,
+    pending: (&mut VecDeque<Vec<u8>>, &mut usize),
+    startup_hold_until: &mut Option<Instant>,
+) -> Result<()> {
+    let Some(deadline) = *startup_hold_until else {
+        return Ok(());
+    };
+    if !predictor.alternate_screen && Instant::now() < deadline {
+        return Ok(());
+    }
+    *startup_hold_until = None;
+    let (pending, pending_bytes) = pending;
+    flush_pending_user_input(
+        socket,
+        addr,
+        cred,
+        send_seq,
+        predictor,
+        pending,
+        pending_bytes,
+    )
+    .await
+}
+
 async fn send_input(
     socket: &UdpSocket,
     addr: SocketAddr,
@@ -5105,11 +5178,49 @@ struct RawMode;
 
 impl RawMode {
     fn enter() -> Result<Self> {
+        flush_local_terminal_input();
         enable_raw_mode()?;
+        flush_local_terminal_input();
+        drain_local_terminal_input();
         Ok(Self)
     }
 }
 
+#[cfg(unix)]
+fn flush_local_terminal_input() {
+    unsafe {
+        libc::tcflush(libc::STDIN_FILENO, libc::TCIFLUSH);
+    }
+}
+
+#[cfg(not(unix))]
+fn flush_local_terminal_input() {}
+
+#[cfg(unix)]
+fn drain_local_terminal_input() {
+    unsafe {
+        let fd = libc::STDIN_FILENO;
+        let flags = libc::fcntl(fd, libc::F_GETFL);
+        if flags < 0 {
+            return;
+        }
+        if libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) < 0 {
+            return;
+        }
+        let mut buf = [0u8; 1024];
+        loop {
+            let n = libc::read(fd, buf.as_mut_ptr().cast(), buf.len());
+            if n <= 0 {
+                break;
+            }
+        }
+        let _ = libc::fcntl(fd, libc::F_SETFL, flags);
+    }
+}
+
+#[cfg(not(unix))]
+fn drain_local_terminal_input() {}
+
 impl Drop for RawMode {
     fn drop(&mut self) {
         let mut stdout = std::io::stdout();