Restore reconnect status from snapshots

src/bin/dosh-client.rs

@@ -6560,6 +6560,7 @@ async fn run_terminal(
     // Non-destructive disconnect status line. Off in forward-only mode (no TTY to
     // draw on) and respecting the client config (default on, env override).
     let mut disconnect_status = DisconnectStatus::new(resolve_disconnect_status() && !forward_only);
+    let mut status_restore_pending = false;
     let (forward_tx, mut forward_rx) = mpsc::channel::<ForwardEvent>(1024);
     let _forward_keepalive = if local_forwards.is_empty() {
         Some(forward_tx.clone())
@@ -6606,6 +6607,7 @@ async fn run_terminal(
     if let Some(frame) = first_frame {
         if !forward_only {
             render_frame(&frame)?;
+            note_snapshot_rendered(&frame, &mut disconnect_status, &mut status_restore_pending);
             predictor.observe_output(&frame.bytes);
             last_terminal_frame_at = Instant::now();
         }
@@ -6714,6 +6716,11 @@ async fn run_terminal(
                                 );
                                 refresh_live_addr(&mut addr, &cred)?;
                                 render_frame(&frame)?;
+                                note_snapshot_rendered(
+                                    &frame,
+                                    &mut disconnect_status,
+                                    &mut status_restore_pending,
+                                );
                                 predictor.observe_output(&frame.bytes);
                                 last_terminal_frame_at = Instant::now();
                                 last_packet_at = Instant::now();
@@ -6768,6 +6775,11 @@ async fn run_terminal(
                                     &mut stale_terminal_input_suppress_until,
                                 );
                                 render_frame(&frame)?;
+                                note_snapshot_rendered(
+                                    &frame,
+                                    &mut disconnect_status,
+                                    &mut status_restore_pending,
+                                );
                                 predictor.observe_output(&frame.bytes);
                                 last_terminal_frame_at = Instant::now();
                                 last_packet_at = Instant::now();
@@ -6986,6 +6998,11 @@ async fn run_terminal(
                                     );
                                     if !forward_only {
                                         render_frame(&frame)?;
+                                        note_snapshot_rendered(
+                                            &frame,
+                                            &mut disconnect_status,
+                                            &mut status_restore_pending,
+                                        );
                                         predictor.observe_output(&frame.bytes);
                                         last_terminal_frame_at = Instant::now();
                                     }
@@ -7053,6 +7070,11 @@ async fn run_terminal(
                     );
                     if !forward_only {
                         render_frame(&frame)?;
+                        note_snapshot_rendered(
+                            &frame,
+                            &mut disconnect_status,
+                            &mut status_restore_pending,
+                        );
                         predictor.observe_output(&frame.bytes);
                         last_terminal_frame_at = Instant::now();
                         wake_repaint_retry_until = None;
@@ -7082,8 +7104,6 @@ async fn run_terminal(
                 if packet.header.conn_id != [0u8; 16] && packet.header.conn_id != cred.client_id {
                     continue;
                 }
-                // Contact resumed: clear any disconnect status line immediately.
-                disconnect_status.apply(disconnect_status.on_contact())?;
                 match packet.header.kind {
                     PacketKind::Frame | PacketKind::ResumeOk => {
                         let decrypted = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT)
@@ -7112,6 +7132,11 @@ async fn run_terminal(
                                 );
                                 if !forward_only {
                                     render_frame(&frame)?;
+                                    note_snapshot_rendered(
+                                        &frame,
+                                        &mut disconnect_status,
+                                        &mut status_restore_pending,
+                                    );
                                     predictor.observe_output(&frame.bytes);
                                     last_terminal_frame_at = Instant::now();
                                     wake_repaint_retry_until = None;
@@ -7130,6 +7155,10 @@ async fn run_terminal(
                             }
                             continue;
                         };
+                        note_authenticated_contact(
+                            &mut disconnect_status,
+                            &mut status_restore_pending,
+                        )?;
                         let Ok(frame) = protocol::from_body::<Frame>(&plain) else {
                             continue;
                         };
@@ -7139,6 +7168,11 @@ async fn run_terminal(
                             predictor.clear_pending()?;
                             if !forward_only {
                                 render_frame(&frame)?;
+                                note_snapshot_rendered(
+                                    &frame,
+                                    &mut disconnect_status,
+                                    &mut status_restore_pending,
+                                );
                                 predictor.observe_output(&frame.bytes);
                                 last_terminal_frame_at = Instant::now();
                                 wake_repaint_retry_until = None;
@@ -7161,6 +7195,25 @@ async fn run_terminal(
                         send_ack(&socket, addr, &cred, &mut send_seq).await?;
                     }
                     PacketKind::Pong => {
+                        if protocol::decrypt_body(
+                            &packet,
+                            &cred.session_key,
+                            SERVER_TO_CLIENT,
+                        )
+                        .or_else(|err| match previous_session_key {
+                            Some(prev) => {
+                                protocol::decrypt_body(&packet, &prev, SERVER_TO_CLIENT)
+                            }
+                            None => Err(err),
+                        })
+                        .is_err()
+                        {
+                            continue;
+                        }
+                        note_authenticated_contact(
+                            &mut disconnect_status,
+                            &mut status_restore_pending,
+                        )?;
                         last_packet_at = Instant::now();
                         if should_flush_terminal_input_after_contact(
                             forward_only,
@@ -7188,6 +7241,10 @@ async fn run_terminal(
                         let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
                             continue;
                         };
+                        note_authenticated_contact(
+                            &mut disconnect_status,
+                            &mut status_restore_pending,
+                        )?;
                         let Ok(rekey) = protocol::from_body::<Rekey>(&plain) else {
                             continue;
                         };
@@ -7240,6 +7297,11 @@ async fn run_terminal(
                                 );
                                 if !forward_only {
                                     render_frame(&frame)?;
+                                    note_snapshot_rendered(
+                                        &frame,
+                                        &mut disconnect_status,
+                                        &mut status_restore_pending,
+                                    );
                                     predictor.observe_output(&frame.bytes);
                                     last_terminal_frame_at = Instant::now();
                                     wake_repaint_retry_until = None;
@@ -7800,7 +7862,9 @@ async fn run_terminal(
                 }
                 let mut repainted_this_tick = false;
                 let stale = last_packet_at.elapsed();
-                if stale >= Duration::from_secs(reconnect_timeout_secs.max(1)) {
+                if status_restore_pending
+                    || stale >= Duration::from_secs(reconnect_timeout_secs.max(1))
+                {
                     if let Some(frame) = reconnect(
                         &socket,
                         &mut cred,
@@ -7817,6 +7881,11 @@ async fn run_terminal(
                         );
                         if !forward_only {
                             render_frame(&frame)?;
+                            note_snapshot_rendered(
+                                &frame,
+                                &mut disconnect_status,
+                                &mut status_restore_pending,
+                            );
                             predictor.observe_output(&frame.bytes);
                             last_terminal_frame_at = Instant::now();
                             last_idle_repaint_attempt_at = Instant::now();
@@ -7896,6 +7965,11 @@ async fn run_terminal(
                                 &mut stale_terminal_input_suppress_until,
                             );
                             render_frame(&frame)?;
+                            note_snapshot_rendered(
+                                &frame,
+                                &mut disconnect_status,
+                                &mut status_restore_pending,
+                            );
                             predictor.observe_output(&frame.bytes);
                             last_terminal_frame_at = Instant::now();
                             last_packet_at = Instant::now();
@@ -7913,7 +7987,7 @@ async fn run_terminal(
                         }
                     }
                 }
-                // Refresh / clear the non-destructive disconnect status line based
+                // Refresh or retire the reconnect status overlay based
                 // on how long the link has been silent (recomputed after any
                 // reconnect attempt above may have reset `last_packet_at`).
                 if !forward_only {
@@ -7922,6 +7996,9 @@ async fn run_terminal(
                     } else {
                         disconnect_status.on_tick(last_packet_at.elapsed())
                     };
+                    if action == StatusAction::Forget {
+                        status_restore_pending = true;
+                    }
                     disconnect_status.apply(action)?;
                 }
             }
@@ -7972,7 +8049,7 @@ async fn run_terminal(
         }
     }
     // Leave the terminal clean on exit: erase any lingering status line.
-    let _ = disconnect_status.apply(StatusAction::Clear);
+    let _ = disconnect_status.apply(StatusAction::Forget);
     let _ = detach_once(&socket, &cred, send_seq).await;
     Ok(())
 }
@@ -10175,7 +10252,7 @@ fn input_matches_escape(bytes: &[u8], escape_key: Option<&[u8]>) -> bool {
     !escape_key.is_empty() && bytes.windows(escape_key.len()).any(|w| w == escape_key)
 }
 
-/// Resolve whether the non-destructive disconnect status line is enabled.
+/// Resolve whether the snapshot-restored disconnect status line is enabled.
 /// `DOSH_DISCONNECT_STATUS` (0/false/off to disable) overrides for quick
 /// experiments; otherwise the client config's `disconnect_status` applies,
 /// defaulting to on.
@@ -10400,18 +10477,18 @@ const TERMINAL_SNAPSHOT_RESET: &[u8] = concat!(
 /// idle period (the run loop only pings after 2s of quiet) never flashes it.
 const DISCONNECT_STATUS_THRESHOLD_SECS: u64 = 2;
 
-/// Non-destructive, mosh-style disconnect status bar.
+/// Mosh-style disconnect status bar with snapshot-backed restoration.
 ///
 /// When server packets stop arriving we paint one blue line on the terminal's
 /// top row — e.g. `[dosh] reconnecting — last contact 3s ago` — using
 /// save/restore cursor (ESC 7 / ESC 8) so the real application cursor never
-/// moves. The run loop suppresses the bar while the server is in the alternate
+/// moves. Painting row 1 necessarily replaces its cells, so authenticated
-/// screen because a saved cursor does not make writing row 1 non-destructive for
+/// recovery retires the overlay without clearing the row and requests an
-/// full-screen TUIs. The bar is refreshed on the status timer tick and cleared
+/// authoritative server snapshot to restore it. The run loop suppresses the bar
-/// the instant a packet resumes.
+/// while the server is in the alternate screen.
 ///
 /// The decision logic and the rendered text are pure and unit-tested; only
-/// [`DisconnectStatus::emit`] / [`DisconnectStatus::emit_clear`] touch stdout.
+/// [`DisconnectStatus::emit`] touches stdout.
 struct DisconnectStatus {
     enabled: bool,
     /// Whether a status line is currently painted on screen (so we know to clear
@@ -10429,12 +10506,32 @@ enum StatusAction {
     None,
     /// Paint/repaint the status bar with this exact text on the top row.
     Show(String),
-    /// Erase the status line (link recovered or feature disabled).
-    Clear,
     /// Drop tracked status state without touching the terminal.
     Forget,
 }
 
+fn note_authenticated_contact(
+    status: &mut DisconnectStatus,
+    restore_pending: &mut bool,
+) -> Result<()> {
+    let action = status.on_contact();
+    if action == StatusAction::Forget {
+        *restore_pending = true;
+    }
+    status.apply(action)
+}
+
+fn note_snapshot_rendered(
+    frame: &Frame,
+    status: &mut DisconnectStatus,
+    restore_pending: &mut bool,
+) {
+    if frame.snapshot {
+        status.applied(&StatusAction::Forget);
+        *restore_pending = false;
+    }
+}
+
 impl DisconnectStatus {
     fn new(enabled: bool) -> Self {
         Self {
@@ -10462,7 +10559,7 @@ impl DisconnectStatus {
     fn on_tick(&self, stale: Duration) -> StatusAction {
         if !self.enabled {
             return if self.shown {
-                StatusAction::Clear
+                StatusAction::Forget
             } else {
                 StatusAction::None
             };
@@ -10475,17 +10572,17 @@ impl DisconnectStatus {
                 StatusAction::Show(text)
             }
         } else if self.shown {
-            StatusAction::Clear
+            StatusAction::Forget
         } else {
             StatusAction::None
         }
     }
 
-    /// Pure decision for "a packet just arrived": clear any visible status line
+    /// Pure decision for authenticated contact. The caller schedules a snapshot
-    /// the instant contact resumes, otherwise do nothing.
+    /// repaint before forgetting a visible overlay; it never clears row 1.
     fn on_contact(&self) -> StatusAction {
         if self.shown {
-            StatusAction::Clear
+            StatusAction::Forget
         } else {
             StatusAction::None
         }
@@ -10510,10 +10607,6 @@ impl DisconnectStatus {
                 self.shown = true;
                 self.drawn_text = text.clone();
             }
-            StatusAction::Clear => {
-                self.shown = false;
-                self.drawn_text.clear();
-            }
             StatusAction::Forget => {
                 self.shown = false;
                 self.drawn_text.clear();
@@ -10526,7 +10619,6 @@ impl DisconnectStatus {
     fn apply(&mut self, action: StatusAction) -> Result<()> {
         match &action {
             StatusAction::Show(text) => self.emit(text)?,
-            StatusAction::Clear => self.emit_clear()?,
             StatusAction::Forget => {}
             StatusAction::None => {}
         }
@@ -10541,12 +10633,6 @@ impl DisconnectStatus {
         let bytes = render_status_overlay(text, rows);
         emit_status(&bytes)
     }
-
-    /// Erase the top-row status bar, again with save/restore cursor.
-    fn emit_clear(&self) -> Result<()> {
-        let (_, rows) = terminal_size();
-        emit_status(&render_status_clear(rows))
-    }
 }
 
 /// Pure: build the escape sequence that paints `text` on a top blue status bar,
@@ -10563,16 +10649,6 @@ fn render_status_overlay(text: &str, _rows: u16) -> Vec<u8> {
     out
 }
 
-/// Pure: build the escape sequence that clears the top status bar.
-fn render_status_clear(_rows: u16) -> Vec<u8> {
-    let mut out = Vec::with_capacity(12);
-    out.extend_from_slice(b"\x1b7");
-    out.extend_from_slice(b"\x1b[1;1H");
-    out.extend_from_slice(b"\x1b[2K");
-    out.extend_from_slice(b"\x1b8");
-    out
-}
-
 #[cfg(not(test))]
 fn emit_status(bytes: &[u8]) -> Result<()> {
     let mut stdout = std::io::stdout();
@@ -10879,14 +10955,14 @@ mod tests {
         is_local_status_target, is_resume_response_for_client, latest_release_download_url,
         load_first_native_identity_with_prompt, local_symlink_target_is_dir,
         local_username_from_env, native_proxy_udp_warning, newest_client_trace_path_from,
-        parse_dynamic_forward, parse_escape_key, parse_local_forward, parse_remote_forward,
+        note_authenticated_contact, note_snapshot_rendered, parse_dynamic_forward,
-        parse_single_remote_path, parse_ssh_config, parse_trace_line, parse_trace_options,
+        parse_escape_key, parse_local_forward, parse_remote_forward, parse_single_remote_path,
-        parse_trace_report_options, parse_trace_summary, parse_update_options,
+        parse_ssh_config, parse_trace_line, parse_trace_options, parse_trace_report_options,
-        post_submit_hold_duration, queue_or_send_stream_data, queue_pending_user_input,
+        parse_trace_summary, parse_update_options, post_submit_hold_duration,
-        queue_stale_pending_user_input, raw_contains_host_table, recv_response_until,
+        queue_or_send_stream_data, queue_pending_user_input, queue_stale_pending_user_input,
-        refresh_live_addr, release_artifact_name_for, release_tag_download_url,
+        raw_contains_host_table, recv_response_until, refresh_live_addr, release_artifact_name_for,
-        release_tag_from_effective_url, release_version_from_tag, remote_unix_server_update_script,
+        release_tag_download_url, release_tag_from_effective_url, release_version_from_tag,
-        render_frame_bytes, render_status_clear, render_status_overlay, requested_env,
+        remote_unix_server_update_script, render_frame_bytes, render_status_overlay, requested_env,
         resolve_forward_agent_endpoint, resolved_startup_command, retire_stream_state,
         retransmit_stream_closes, retransmit_stream_eofs, retransmit_stream_opens,
         retransmit_stream_window_adjusts, rewrite_forward_command, sanitize_trace_name,
@@ -14374,10 +14450,10 @@ mod tests {
         );
     }
 
-    /// Contact resuming clears a shown line immediately, and a sub-threshold tick
+    /// Contact retires a shown overlay without erasing row 1; the run loop then
-    /// also clears it; clearing when nothing is shown is a no-op.
+    /// restores the covered cells from an authenticated snapshot.
     #[test]
-    fn disconnect_status_clears_on_contact_and_recovery() {
+    fn disconnect_status_forgets_on_contact_and_recovery() {
         let mut status = DisconnectStatus::new(true);
 
         // Nothing shown yet: contact is a no-op.
@@ -14387,27 +14463,29 @@ mod tests {
         status.applied(&show);
         assert!(status.shown);
 
-        // A packet arrives: clear immediately.
+        let forget = status.on_contact();
-        let clear = status.on_contact();
+        assert_eq!(forget, StatusAction::Forget);
-        assert_eq!(clear, StatusAction::Clear);
+        status.applied(&forget);
-        status.applied(&clear);
         assert!(!status.shown);
 
-        // Show again, then a sub-threshold tick (link recovered) clears it too.
+        // A sub-threshold authenticated recovery also requests restoration.
         let show = status.on_tick(Duration::from_secs(5));
         status.applied(&show);
-        assert_eq!(status.on_tick(Duration::from_secs(0)), StatusAction::Clear);
+        assert_eq!(status.on_tick(Duration::from_secs(0)), StatusAction::Forget);
     }
 
-    /// When the feature is disabled the machine never shows, and clears anything
+    /// When disabled the machine never shows, and retires stale state without
-    /// already on screen (e.g. config toggled off at runtime).
+    /// writing destructive row-clearing sequences.
     #[test]
     fn disconnect_status_disabled_never_shows() {
         let mut status = DisconnectStatus::new(false);
         assert_eq!(status.on_tick(Duration::from_secs(10)), StatusAction::None);
-        // Simulate a stale "shown" flag and confirm it gets cleared.
+        // Simulate a stale "shown" flag and confirm it requests restoration.
         status.shown = true;
-        assert_eq!(status.on_tick(Duration::from_secs(10)), StatusAction::Clear);
+        assert_eq!(
+            status.on_tick(Duration::from_secs(10)),
+            StatusAction::Forget
+        );
     }
 
     #[test]
@@ -14424,6 +14502,23 @@ mod tests {
         assert_eq!(status.drawn_text, "");
     }
 
+    #[test]
+    fn authenticated_contact_requires_snapshot_to_restore_covered_row() {
+        let mut status = DisconnectStatus::new(true);
+        let show = status.on_tick(Duration::from_secs(3));
+        status.applied(&show);
+        let mut restore_pending = false;
+
+        note_authenticated_contact(&mut status, &mut restore_pending).unwrap();
+        assert!(!status.shown);
+        assert!(restore_pending);
+
+        note_snapshot_rendered(&test_frame(1, false), &mut status, &mut restore_pending);
+        assert!(restore_pending);
+        note_snapshot_rendered(&test_frame(2, true), &mut status, &mut restore_pending);
+        assert!(!restore_pending);
+    }
+
     /// The rendered overlay must preserve cursor position: save cursor (ESC 7),
     /// park on the top row, paint the blue reconnect bar, then restore cursor
     /// (ESC 8). The run loop suppresses this renderer entirely in full-screen
@@ -14441,11 +14536,6 @@ mod tests {
         assert!(bytes.windows(8).any(|w| w == b"\x1b[44;37m"));
         // Resets attributes before restoring the cursor.
         assert!(bytes.windows(4).any(|w| w == b"\x1b[0m"));
-
-        let clear = render_status_clear(24);
-        assert!(clear.starts_with(b"\x1b7"));
-        assert!(clear.ends_with(b"\x1b8"));
-        assert!(clear.windows(4).any(|w| w == b"\x1b[2K"));
     }
 
     #[test]