2 files changed, 62 insertions, 5 deletions

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<String>,
+    pub response_write_complete: Option<std::sync::mpsc::Receiver<()>>,
 }
 
 pub type ApiRequestSender = mpsc::UnboundedSender<ApiRequestMessage>;
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<ServerCapabilities>,
+    response_write_complete: Option<std::sync::mpsc::Receiver<()>>,
 ) -> 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<Duration>,
+) -> 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<Duration>,
+    response_write_complete: Option<std::sync::mpsc::Receiver<()>>,
 ) -> 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");
@@ -988,6 +1005,45 @@ mod tests {
     }
 
     #[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) =
             spawn_pane_get_responder(crate::api::schema::AgentStatus::Blocked);