Arm wake repaint retry from input paths

src/bin/dosh-client.rs

@@ -5914,6 +5914,13 @@ async fn run_terminal(
                             forward_only,
                             input_status_tick_gap,
                         ) {
+                            let reconnect_started_at = Instant::now();
+                            if let Some(deadline) =
+                                wake_repaint_retry_deadline(reconnect_started_at, input_status_tick_gap)
+                            {
+                                wake_repaint_retry_until = Some(deadline);
+                                trace_wake_repaint_retry_arm("input_sleep_gap", input_status_tick_gap);
+                            }
                             dosh::trace::event(
                                 "client.local_sleep_input_reconnect_start",
                                 &[(
@@ -5921,7 +5928,7 @@ async fn run_terminal(
                                     input_status_tick_gap.as_millis().to_string(),
                                 )],
                             );
-                            last_status_tick_at = Instant::now();
+                            last_status_tick_at = reconnect_started_at;
                             arm_stale_terminal_input_suppression(
                                 &mut stale_terminal_input_suppress_until,
                             );
@@ -5973,6 +5980,9 @@ async fn run_terminal(
                                 &[("silent_ms", last_packet_at.elapsed().as_millis().to_string())],
                             );
                             last_focus_repaint_at = Instant::now();
+                            wake_repaint_retry_until =
+                                Some(last_focus_repaint_at + LOCAL_SLEEP_REPAINT_RETRY_WINDOW);
+                            trace_wake_repaint_retry_arm("focus", last_packet_at.elapsed());
                             if let Some(frame) = reconnect(
                                 &socket,
                                 &mut cred,
@@ -6369,6 +6379,7 @@ async fn run_terminal(
                                 render_frame(&frame)?;
                                 predictor.observe_output(&frame.bytes);
                                 last_terminal_frame_at = Instant::now();
+                                wake_repaint_retry_until = None;
                                 flush_startup_input_if_ready(
                                     &socket,
                                     addr,
@@ -6885,19 +6896,11 @@ async fn run_terminal(
                 let status_tick_at = Instant::now();
                 let status_tick_gap = status_tick_at.duration_since(last_status_tick_at);
                 last_status_tick_at = status_tick_at;
-                if status_tick_gap >= LOCAL_SLEEP_REPAINT_AFTER {
+                if let Some(deadline) =
-                    wake_repaint_retry_until =
+                    wake_repaint_retry_deadline(status_tick_at, status_tick_gap)
-                        Some(status_tick_at + LOCAL_SLEEP_REPAINT_RETRY_WINDOW);
+                {
-                    dosh::trace::event(
+                    wake_repaint_retry_until = Some(deadline);
-                        "client.wake_repaint_retry_arm",
+                    trace_wake_repaint_retry_arm("status_tick_gap", status_tick_gap);
-                        &[
-                            ("tick_gap_ms", status_tick_gap.as_millis().to_string()),
-                            (
-                                "window_ms",
-                                LOCAL_SLEEP_REPAINT_RETRY_WINDOW.as_millis().to_string(),
-                            ),
-                        ],
-                    );
                 }
                 let mut repainted_this_tick = false;
                 let stale = last_packet_at.elapsed();
@@ -7457,6 +7460,24 @@ fn should_reconnect_before_input_for_local_sleep(
     !forward_only && status_tick_gap >= LOCAL_SLEEP_REPAINT_AFTER
 }
 
+fn wake_repaint_retry_deadline(now: Instant, status_tick_gap: Duration) -> Option<Instant> {
+    (status_tick_gap >= LOCAL_SLEEP_REPAINT_AFTER).then_some(now + LOCAL_SLEEP_REPAINT_RETRY_WINDOW)
+}
+
+fn trace_wake_repaint_retry_arm(trigger: &str, elapsed: Duration) {
+    dosh::trace::event(
+        "client.wake_repaint_retry_arm",
+        &[
+            ("trigger", trigger.to_string()),
+            ("elapsed_ms", elapsed.as_millis().to_string()),
+            (
+                "window_ms",
+                LOCAL_SLEEP_REPAINT_RETRY_WINDOW.as_millis().to_string(),
+            ),
+        ],
+    );
+}
+
 fn should_strip_unowned_terminal_reports(alternate_screen: bool, mouse_tracking: bool) -> bool {
     // Only a full-screen app that explicitly enabled mouse reporting owns these
     // bytes. In every other state, local terminal mouse reports are stale UI
@@ -9496,7 +9517,8 @@ mod tests {
         summarize_trace_file_with_mode, terminal_private_mode_transition, toml_bare_key_or_quoted,
         top_trace_events, trace_report_warnings, unix_update_script,
         update_binary_version_for_installer, update_installer_url, update_version_status,
-        upsert_managed_block, valid_forward_host, vscode_safe_alias, windows_update_script,
+        upsert_managed_block, valid_forward_host, vscode_safe_alias, wake_repaint_retry_deadline,
+        windows_update_script,
     };
     use dosh::config::{ClientConfig, CommandExtension, HostConfig};
     use dosh::native::EnvVar;
@@ -11816,6 +11838,23 @@ mod tests {
         ));
     }
 
+    #[test]
+    fn local_sleep_gap_arms_wake_repaint_retry_window() {
+        let now = Instant::now();
+        assert_eq!(
+            wake_repaint_retry_deadline(now, LOCAL_SLEEP_REPAINT_AFTER - Duration::from_millis(1)),
+            None
+        );
+        assert_eq!(
+            wake_repaint_retry_deadline(now, LOCAL_SLEEP_REPAINT_AFTER),
+            Some(now + LOCAL_SLEEP_REPAINT_RETRY_WINDOW)
+        );
+        assert_eq!(
+            wake_repaint_retry_deadline(now, LOCAL_SLEEP_REPAINT_AFTER + Duration::from_secs(30)),
+            Some(now + LOCAL_SLEEP_REPAINT_RETRY_WINDOW)
+        );
+    }
+
     #[test]
     fn reconnect_mouse_quarantine_is_long_enough_for_sleep_wake_noise() {
         assert!(POST_RECONNECT_STALE_INPUT_GRACE >= LOCAL_SLEEP_REPAINT_AFTER);