diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index c73bbd6..8fe21d6 100644 --- a/src/bin/dosh-server.rs +++ b/src/bin/dosh-server.rs @@ -804,7 +804,7 @@ fn attach_snapshot(session: &Session, cols: u16, rows: u16) -> Vec { 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 { bytes } +fn should_persist_screen_immediately(bytes: &[u8]) -> bool { + bytes.len() <= 1024 && bytes.contains(&b'\n') +} + async fn retransmit_pending( state: &Arc>, socket: &Arc, @@ -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();