Preserve restored screen through update reattach

src/bin/dosh-server.rs

@@ -804,7 +804,7 @@ fn attach_snapshot(session: &Session, cols: u16, rows: u16) -> Vec<u8> {
     let cols = cols.max(1);
     let rows = rows.max(1);
     if let Some(restored) = &session.restored_screen
-        && restored.output_seq == session.output_seq
+        && (restored.output_seq == session.output_seq || session.clients.is_empty())
         && restored.cols == cols
         && restored.rows == rows
     {
@@ -3691,8 +3691,11 @@ async fn broadcast_output(
             .sessions
             .get_mut(&output.session)
             .ok_or_else(|| anyhow!("unknown session"))?;
+        let has_clients = !session.clients.is_empty();
         session.parser.process(&output.bytes);
-        session.restored_screen = None;
+        if has_clients {
+            session.restored_screen = None;
+        }
         session.output_seq += 1;
         let output_seq = session.output_seq;
         session.recent.push_back(output.bytes.clone());
@@ -3709,7 +3712,8 @@ async fn broadcast_output(
                 .bytes_since_persist
                 .saturating_add(output.bytes.len());
             let now = Instant::now();
-            if session.bytes_since_persist >= SCREEN_PERSIST_BYTES
+            if should_persist_screen_immediately(&output.bytes)
+                || session.bytes_since_persist >= SCREEN_PERSIST_BYTES
                 || now.duration_since(session.last_screen_persist_at) >= SCREEN_PERSIST_MAX_AGE
             {
                 session.bytes_since_persist = 0;
@@ -3847,6 +3851,10 @@ fn screen_snapshot(screen: &vt100::Screen) -> Vec<u8> {
     bytes
 }
 
+fn should_persist_screen_immediately(bytes: &[u8]) -> bool {
+    bytes.len() <= 1024 && bytes.contains(&b'\n')
+}
+
 async fn retransmit_pending(
     state: &Arc<Mutex<ServerState>>,
     socket: &Arc<UdpSocket>,
@@ -4316,6 +4324,104 @@ mod tests {
         assert_ne!(attach_snapshot(&session, 100, 30), restored);
     }
 
+    #[test]
+    fn first_reattach_prefers_restored_screen_even_after_pre_attach_pty_noise() {
+        let restored = b"\x1b[?1049lRESTORED_AFTER_RESTART".to_vec();
+        let mut session = Session {
+            pty: None,
+            parser: vt100::Parser::new(24, 80, 100),
+            restored_screen: Some(RestoredScreenSnapshot {
+                cols: 80,
+                rows: 24,
+                output_seq: 7,
+                snapshot: restored.clone(),
+            }),
+            clients: HashMap::new(),
+            output_seq: 8,
+            recent: VecDeque::new(),
+            empty_since: Some(Instant::now()),
+            holder_control: None,
+            persistent: true,
+            bytes_since_persist: 0,
+            last_persisted_seq: 7,
+            last_screen_persist_at: Instant::now(),
+        };
+        session.parser.process(b"$ ");
+
+        assert_eq!(attach_snapshot(&session, 80, 24), restored);
+
+        session
+            .clients
+            .insert([1u8; 16], test_client_state([2u8; 32]));
+        assert_ne!(attach_snapshot(&session, 80, 24), restored);
+    }
+
+    #[tokio::test]
+    async fn no_client_pty_output_before_first_reattach_keeps_restored_screen() {
+        let restored = b"\x1b[?1049lRESTORED_AFTER_RESTART".to_vec();
+        let (pty_tx, _rx) = mpsc::unbounded_channel();
+        let state = Arc::new(Mutex::new(ServerState::new(
+            ServerConfig {
+                persist_sessions: true,
+                ..ServerConfig::default()
+            },
+            [0u8; 32],
+            pty_tx,
+        )));
+        {
+            let mut locked = state.lock().unwrap();
+            locked.sessions.insert(
+                "work".to_string(),
+                Session {
+                    pty: None,
+                    parser: vt100::Parser::new(24, 80, 100),
+                    restored_screen: Some(RestoredScreenSnapshot {
+                        cols: 80,
+                        rows: 24,
+                        output_seq: 7,
+                        snapshot: restored.clone(),
+                    }),
+                    clients: HashMap::new(),
+                    output_seq: 7,
+                    recent: VecDeque::new(),
+                    empty_since: Some(Instant::now()),
+                    holder_control: None,
+                    persistent: true,
+                    bytes_since_persist: 0,
+                    last_persisted_seq: 7,
+                    last_screen_persist_at: Instant::now(),
+                },
+            );
+        }
+
+        let socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
+        broadcast_output(
+            &state,
+            &socket,
+            PtyOutput {
+                session: "work".to_string(),
+                bytes: b"$ ".to_vec(),
+                exited: false,
+            },
+        )
+        .await
+        .unwrap();
+
+        let locked = state.lock().unwrap();
+        let session = locked.sessions.get("work").unwrap();
+        assert!(session.restored_screen.is_some());
+        assert_eq!(attach_snapshot(session, 80, 24), restored);
+    }
+
+    #[test]
+    fn line_oriented_shell_output_persists_screen_immediately() {
+        assert!(should_persist_screen_immediately(
+            b"IMPLICIT_SCREEN_MARKER\r\n$ "
+        ));
+        assert!(!should_persist_screen_immediately(b"\x1b[2K\x1b[H"));
+        assert!(!should_persist_screen_immediately(&vec![b'x'; 2048]));
+    }
+
     #[tokio::test]
     async fn unknown_resume_reject_keeps_client_id() {
         let (pty_tx, _rx) = mpsc::unbounded_channel();