From b88aa1301e324c80d46cac767f87109f38cf4fe2 Mon Sep 17 00:00:00 2001 From: Ogulcan Celik Date: Tue, 14 Jul 2026 22:07:06 +0300 Subject: fix: flush live handoff responses before exit --- src/api/mod.rs | 1 + src/api/server.rs | 66 ++++++++++++++++++++++++++++++++++++++++++++++---- src/server/headless.rs | 43 ++++++++++++++++++++++++++++++-- 3 files changed, 103 insertions(+), 7 deletions(-) (limited to 'src') diff --git a/src/api/mod.rs b/src/api/mod.rs index 6f8fa9df..dad4cfd8 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -67,6 +67,7 @@ pub(crate) fn request_changes_ui(request: &Request) -> bool { pub struct ApiRequestMessage { pub request: Request, pub respond_to: std::sync::mpsc::Sender, + pub response_write_complete: Option>, } pub type ApiRequestSender = mpsc::UnboundedSender; diff --git a/src/api/server.rs b/src/api/server.rs index 9cef5ec1..5dfad989 100644 --- a/src/api/server.rs +++ b/src/api/server.rs @@ -257,6 +257,7 @@ fn handle_connection( result } method_body => { + let (response_write_tx, response_write_rx) = std::sync::mpsc::channel(); let response = handle_request( Request { id: request_id.clone(), @@ -264,8 +265,10 @@ fn handle_connection( }, api_tx, capabilities, + Some(response_write_rx), ); let result = write_text_line_allow_disconnect(&mut stream, &response); + let _ = response_write_tx.send(()); match &result { Ok(()) => crate::logging::api_request_completed( &request_id, @@ -286,6 +289,7 @@ fn handle_request( request: Request, api_tx: &ApiRequestSender, capabilities: Option, + response_write_complete: Option>, ) -> String { match request.method { Method::Ping(_) => serde_json::to_string(&SuccessResponse { @@ -300,7 +304,12 @@ fn handle_request( r#"{"id":"","error":{"code":"internal_error","message":"failed to encode response"}}"# .to_string() }), - _ => dispatch_to_app(request, api_tx), + _ => dispatch_to_app_with_timeout_and_write_completion( + request, + api_tx, + None, + response_write_complete, + ), } } @@ -693,20 +702,26 @@ pub(super) fn should_stop_connection( local_stream_peer_closed(stream) } -fn dispatch_to_app(request: Request, api_tx: &ApiRequestSender) -> String { - dispatch_to_app_with_timeout(request, api_tx, None) +pub(super) fn dispatch_to_app_with_timeout( + request: Request, + api_tx: &ApiRequestSender, + timeout: Option, +) -> String { + dispatch_to_app_with_timeout_and_write_completion(request, api_tx, timeout, None) } -pub(super) fn dispatch_to_app_with_timeout( +fn dispatch_to_app_with_timeout_and_write_completion( request: Request, api_tx: &ApiRequestSender, timeout: Option, + response_write_complete: Option>, ) -> String { let request_id = request.id.clone(); let (respond_to, response_rx) = std::sync::mpsc::channel(); if let Err(err) = api_tx.send(ApiRequestMessage { request, respond_to, + response_write_complete, }) { return error_response_json( request_id, @@ -952,6 +967,7 @@ mod tests { live_handoff: true, detached_server_daemon: true, }), + None, ); let parsed: SuccessResponse = serde_json::from_str(&response).unwrap(); @@ -968,7 +984,8 @@ mod tests { }; let request_for_thread = request.clone(); - let thread = std::thread::spawn(move || handle_request(request_for_thread, &tx, None)); + let thread = + std::thread::spawn(move || handle_request(request_for_thread, &tx, None, None)); let msg = rx.blocking_recv().unwrap(); assert_eq!(msg.request.id, "req_2"); @@ -987,6 +1004,45 @@ mod tests { assert_eq!(parsed.id, "req_2"); } + #[test] + fn dispatched_request_reports_response_write_completion() { + let (api_tx, mut api_rx) = mpsc::unbounded_channel(); + let (mut client, server, _path) = local_stream_pair("api-response-write-complete"); + client + .write_all(br#"{"id":"req_write","method":"workspace.list","params":{}}"#) + .unwrap(); + client.write_all(b"\n").unwrap(); + client.flush().unwrap(); + + let running = Arc::new(AtomicBool::new(true)); + let server_running = Arc::clone(&running); + let event_hub = EventHub::default(); + let server_thread = std::thread::spawn(move || { + handle_connection(server, &api_tx, &event_hub, &server_running, None) + }); + + let msg = api_rx.blocking_recv().unwrap(); + let response_write_complete = msg + .response_write_complete + .expect("socket-dispatched requests include write completion"); + msg.respond_to + .send( + serde_json::to_string(&SuccessResponse { + id: msg.request.id, + result: ResponseResult::Ok {}, + }) + .unwrap(), + ) + .unwrap(); + + response_write_complete + .recv_timeout(Duration::from_secs(1)) + .expect("response write completion"); + let response: SuccessResponse = serde_json::from_str(&read_line(&mut client)).unwrap(); + assert_eq!(response.id, "req_write"); + server_thread.join().unwrap().unwrap(); + } + #[test] fn events_wait_agent_status_returns_initial_match() { let (api_tx, responder) = diff --git a/src/server/headless.rs b/src/server/headless.rs index bc64f0a8..7725954c 100644 --- a/src/server/headless.rs +++ b/src/server/headless.rs @@ -71,6 +71,26 @@ use crate::server::client_transport::ClientWriter; #[cfg(test)] use std::fs; +const LIVE_HANDOFF_RESPONSE_WRITE_TIMEOUT: Duration = Duration::from_secs(6); + +fn wait_for_live_handoff_response_write( + response_write_complete: Option>, +) { + let Some(response_write_complete) = response_write_complete else { + return; + }; + + match response_write_complete.recv_timeout(LIVE_HANDOFF_RESPONSE_WRITE_TIMEOUT) { + Ok(()) => {} + Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { + warn!("timed out waiting for live handoff response write; old server exiting"); + } + Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => { + warn!("live handoff response writer disconnected; old server exiting"); + } + } +} + fn sound_notify_message(sound: crate::sound::Sound) -> &'static str { match sound { crate::sound::Sound::Done => "agent done", @@ -771,6 +791,7 @@ impl HeadlessServer { method, }, respond_to, + response_write_complete: None, }, true, ); @@ -1101,11 +1122,14 @@ impl HeadlessServer { } crate::server::handoff::wait_owned_ack(&mut stream); + Ok(()) + } + + fn finish_live_handoff_shutdown(&mut self) { self.shutting_down = true; self.app.state.should_quit = true; self.app.no_session = true; info!("live handoff completed; old server exiting"); - Ok(()) } #[cfg(not(unix))] @@ -2812,7 +2836,9 @@ impl HeadlessServer { let metadata_expired = self.app.expire_due_metadata(Instant::now()); if let api::schema::Method::ServerLiveHandoff(params) = &msg.request.method { - let response = match self.perform_live_handoff(params.clone()) { + let handoff_result = self.perform_live_handoff(params.clone()); + let handoff_succeeded = handoff_result.is_ok(); + let response = match handoff_result { Ok(()) => serde_json::to_string(&api::schema::SuccessResponse { id: msg.request.id, result: api::schema::ResponseResult::Ok {}, @@ -2827,6 +2853,10 @@ impl HeadlessServer { } .unwrap_or_else(|_| "{}".to_string()); let _ = msg.respond_to.send(response); + if handoff_succeeded { + wait_for_live_handoff_response_write(msg.response_write_complete); + self.finish_live_handoff_shutdown(); + } return true; } @@ -4374,6 +4404,7 @@ mod tests { method: api::schema::Method::PaneList(api::schema::PaneListParams::default()), }, respond_to, + response_write_complete: None, }); let response: api::schema::SuccessResponse = serde_json::from_str(&response_rx.recv().unwrap()).unwrap(); @@ -4413,6 +4444,7 @@ mod tests { method: api::schema::Method::ServerStop(api::schema::EmptyParams::default()), }, respond_to, + response_write_complete: None, }) ); let response = response_rx @@ -8395,6 +8427,7 @@ next_tab = "" ), }, respond_to, + response_write_complete: None, }); assert!(changed); @@ -8480,6 +8513,7 @@ next_tab = "" ), }, respond_to, + response_write_complete: None, }); assert!(changed); @@ -8531,6 +8565,7 @@ next_tab = "" ), }, respond_to, + response_write_complete: None, }); assert!(changed); @@ -8562,6 +8597,7 @@ next_tab = "" ), }, respond_to, + response_write_complete: None, }); assert!(changed); @@ -8598,6 +8634,7 @@ next_tab = "" ), }, respond_to, + response_write_complete: None, }) ); @@ -8653,6 +8690,7 @@ next_tab = "" ), }, respond_to, + response_write_complete: None, }) ); @@ -8924,6 +8962,7 @@ next_tab = "" }), }, respond_to, + response_write_complete: None, }); assert!(changed); -- cgit v1.2.3