diff --git a/crates/maestro-cli/src/attach.rs b/crates/maestro-cli/src/attach.rs index 3f511f8..6c9df3f 100644 --- a/crates/maestro-cli/src/attach.rs +++ b/crates/maestro-cli/src/attach.rs @@ -134,7 +134,7 @@ pub async fn run_attach( let (ws, _resp) = tokio_tungstenite::connect_async(&url).await?; let (mut sink, mut stream) = ws.split(); - let _raw = RawTty::enable()?; + let raw = RawTty::enable()?; // Send initial window size so the remote PTY matches our terminal. if let Some((cols, rows)) = current_winsize() { @@ -156,12 +156,17 @@ pub async fn run_attach( "\r\n[maestro] detached (key 0x{detach_byte:02x})\r\n" )); + let mut session_ended = false; + loop { tokio::select! { // Server → local stdout msg = stream.next() => { match msg { - None => break, + None => { + session_ended = true; + break; + } Some(Err(e)) => return Err(AttachError::WebSocket(e)), Some(Ok(Message::Binary(bytes))) => { stdout.write_all(&bytes).await?; @@ -177,7 +182,10 @@ pub async fn run_attach( Some(Ok(Message::Ping(p))) => { sink.send(Message::Pong(p)).await?; } - Some(Ok(Message::Close(_))) => break, + Some(Ok(Message::Close(_))) => { + session_ended = true; + break; + } Some(Ok(_)) => {} } } @@ -212,6 +220,13 @@ pub async fn run_attach( } } + // Restore the TTY before printing so the message lands in cooked mode. + drop(raw); + if session_ended { + let _ = stdout.write_all(b"[maestro] session ended\n").await; + let _ = stdout.flush().await; + } + Ok(()) } diff --git a/crates/maestro-daemon/src/attach.rs b/crates/maestro-daemon/src/attach.rs index ef5b8fd..63b7750 100644 --- a/crates/maestro-daemon/src/attach.rs +++ b/crates/maestro-daemon/src/attach.rs @@ -74,12 +74,22 @@ async fn serve_attach( manager: SessionManager, id: Uuid, ) -> Result<(), axum::Error> { + let (mut sink, mut stream) = socket.split(); + let mut rx = match manager.subscribe_pty(id).await { Some(rx) => rx, - None => return Ok(()), // session was killed between the auth check and now + None => { + // session was killed between the auth check and now — let the + // client see a clean close instead of a dropped socket. + let _ = sink.send(Message::Close(None)).await; + return Ok(()); + } }; - let (mut sink, mut stream) = socket.split(); + // Set to true when the session itself terminates (broadcast closes); we + // owe the client a Close frame in that case. Local breakouts (e.g. the + // client sent its own Close, or the socket died) skip it. + let mut session_ended = false; loop { tokio::select! { @@ -94,7 +104,10 @@ async fn serve_attach( // dropped some frames under back-pressure; keep going continue; } - Err(RecvError::Closed) => break, + Err(RecvError::Closed) => { + session_ended = true; + break; + } }, // Client input → tmux @@ -126,5 +139,9 @@ async fn serve_attach( }, } } + + if session_ended { + let _ = sink.send(Message::Close(None)).await; + } Ok(()) } diff --git a/crates/maestro-daemon/tests/attach.rs b/crates/maestro-daemon/tests/attach.rs index 3d44bdd..0079083 100644 --- a/crates/maestro-daemon/tests/attach.rs +++ b/crates/maestro-daemon/tests/attach.rs @@ -210,6 +210,80 @@ async fn multi_attach_broadcasts_to_every_client() { let _ = h.manager.kill(id).await; } +#[tokio::test] +async fn sends_close_frame_when_tmux_session_dies_under_liveness_watcher() { + // Bug NAP-226: when bash exits, tmux destroys the session, the liveness + // watcher reaps it, and the broadcast channel closes — but the WS used to + // be dropped silently, leaving the client hung in raw mode. + let h = spawn_daemon_ws("laptop").await; + let id = create_bash_session(&h).await; + let token = token_for(id, "laptop", 300); + let mut ws = connect_attach(h.addr, id, &token).await.unwrap(); + + tokio::time::sleep(Duration::from_millis(100)).await; + + h.manager.start_liveness_watcher(Duration::from_millis(50)); + + // Kill the tmux session out-of-band; the watcher should reap and the + // broadcast channel should close, triggering serve_attach to send Close. + h._fx.tmux.tmux.kill_session(id).expect("kill tmux"); + + let mut saw_close = false; + let deadline = std::time::Instant::now() + Duration::from_secs(3); + while std::time::Instant::now() < deadline { + match tokio::time::timeout(Duration::from_millis(100), ws.next()).await { + Ok(Some(Ok(WsMessage::Close(_)))) => { + saw_close = true; + break; + } + Ok(None) => break, + Ok(Some(Err(_))) => break, + Ok(_) => continue, + Err(_) => continue, + } + } + assert!( + saw_close, + "daemon must send a Close frame when the liveness watcher reaps the session" + ); +} + +#[tokio::test] +async fn sends_close_frame_when_session_is_killed() { + let h = spawn_daemon_ws("laptop").await; + let id = create_bash_session(&h).await; + let token = token_for(id, "laptop", 300); + let mut ws = connect_attach(h.addr, id, &token).await.unwrap(); + + // give the broadcast subscriber a moment to register before killing. + tokio::time::sleep(Duration::from_millis(100)).await; + + h.manager.kill(id).await.unwrap(); + + let mut saw_close = false; + let deadline = std::time::Instant::now() + Duration::from_secs(2); + while std::time::Instant::now() < deadline { + match tokio::time::timeout(Duration::from_millis(100), ws.next()).await { + Ok(Some(Ok(WsMessage::Close(_)))) => { + saw_close = true; + break; + } + Ok(None) => { + // socket dropped without a Close frame — that's exactly the + // bug NAP-226 is fixing. + break; + } + Ok(Some(Err(_))) => break, + Ok(_) => continue, + Err(_) => continue, + } + } + assert!( + saw_close, + "daemon must send a Close frame when the session dies so the client can exit cleanly" + ); +} + #[tokio::test] async fn ping_text_frame_gets_pong() { let h = spawn_daemon_ws("laptop").await;