fix: disconnect stalled terminal observers (#4039)

3 files changed, 91 insertions, 2 deletions

diff --git a/src/platform/linux.rs b/src/platform/linux.rs
index 29534dbe..23393653 100644
--- a/src/platform/linux.rs
+++ b/src/platform/linux.rs
@@ -17,7 +17,7 @@ pub(crate) use super::unix_common::{
     create_remote_ssh_config_file, hostname, local_datetime, remote_bridge_endpoint_path,
     remote_private_temp_base, remote_reattach_argument, remote_reattach_program,
     remote_ssh_config_paths, set_default_plugin_pane_pwd, status_commands_supported,
-    wait_client_stream_readable, StatusCommandGuard,
+    wait_client_stream_readable, write_client_stream, ClientStreamReader, StatusCommandGuard,
 };
 
 const WSL_MARKER_ENV_VARS: &[&str] = &["WSL_DISTRO_NAME", "WSL_INTEROP"];
diff --git a/src/platform/macos.rs b/src/platform/macos.rs
index 8840c7ed..453d3fb3 100644
--- a/src/platform/macos.rs
+++ b/src/platform/macos.rs
@@ -17,7 +17,7 @@ pub(crate) use super::unix_common::{
     create_remote_ssh_config_file, hostname, local_datetime, remote_bridge_endpoint_path,
     remote_private_temp_base, remote_reattach_argument, remote_reattach_program,
     remote_ssh_config_paths, set_default_plugin_pane_pwd, status_commands_supported,
-    wait_client_stream_readable, StatusCommandGuard,
+    wait_client_stream_readable, write_client_stream, ClientStreamReader, StatusCommandGuard,
 };
 
 const PROC_PGRP_ONLY: u32 = 2;
diff --git a/src/platform/unix_common.rs b/src/platform/unix_common.rs
index 756abc17..ccb068f8 100644
--- a/src/platform/unix_common.rs
+++ b/src/platform/unix_common.rs
@@ -8,6 +8,95 @@ pub(crate) fn classify_child_exit(status: &portable_pty::ExitStatus) -> super::C
     }
 }
 
+pub(crate) struct ClientStreamReader<'a>(pub(crate) &'a mut crate::ipc::LocalStream);
+
+impl std::io::Read for ClientStreamReader<'_> {
+    fn read(&mut self, data: &mut [u8]) -> std::io::Result<usize> {
+        use std::os::fd::AsRawFd as _;
+
+        loop {
+            match self.0.read(data) {
+                Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
+                    let crate::ipc::LocalStream::UdSocket(stream) = &*self.0;
+                    let mut descriptor = libc::pollfd {
+                        fd: stream.inner().as_raw_fd(),
+                        events: libc::POLLIN,
+                        revents: 0,
+                    };
+                    // Sleep until input or shutdown, without polling quiet observers.
+                    if unsafe { libc::poll(&mut descriptor, 1, -1) } < 0 {
+                        let error = std::io::Error::last_os_error();
+                        if error.kind() != std::io::ErrorKind::Interrupted {
+                            return Err(error);
+                        }
+                    }
+                }
+                result => return result,
+            }
+        }
+    }
+}
+
+pub(crate) fn write_client_stream(
+    stream: &crate::ipc::LocalStream,
+    mut data: &[u8],
+) -> std::io::Result<()> {
+    use std::io::{self, Write as _};
+    use std::os::fd::AsRawFd as _;
+    use std::time::Instant;
+
+    let crate::ipc::LocalStream::UdSocket(stream) = stream;
+    let mut socket = stream.inner();
+    let Some(timeout) = socket.write_timeout()? else {
+        return socket.write_all(data);
+    };
+    let timed_out = || {
+        // Dropping the writer clone alone would leave the reader blocked.
+        let _ = stream.inner().shutdown(std::net::Shutdown::Both);
+        io::Error::new(
+            io::ErrorKind::TimedOut,
+            "terminal observer stopped receiving output",
+        )
+    };
+    let mut progress = Instant::now();
+    while !data.is_empty() {
+        match socket.write(data) {
+            Ok(0) => return Err(io::ErrorKind::WriteZero.into()),
+            Ok(written) => {
+                data = &data[written..];
+                progress = Instant::now();
+                continue;
+            }
+            Err(error)
+                if matches!(
+                    error.kind(),
+                    io::ErrorKind::WouldBlock | io::ErrorKind::Interrupted
+                ) => {}
+            Err(error) => return Err(error),
+        }
+        let remaining = timeout
+            .checked_sub(progress.elapsed())
+            .ok_or_else(timed_out)?;
+        let mut descriptor = libc::pollfd {
+            fd: socket.as_raw_fd(),
+            events: libc::POLLOUT,
+            revents: 0,
+        };
+        let wait_ms = remaining.as_millis().clamp(1, i32::MAX as u128) as i32;
+        let ready = unsafe { libc::poll(&mut descriptor, 1, wait_ms) };
+        if ready == 0 {
+            return Err(timed_out());
+        }
+        if ready < 0 {
+            let error = io::Error::last_os_error();
+            if error.kind() != io::ErrorKind::Interrupted {
+                return Err(error);
+            }
+        }
+    }
+    Ok(())
+}
+
 pub(crate) fn wait_client_stream_readable(stream: &crate::ipc::LocalStream) -> std::io::Result<()> {
     use std::os::fd::{AsFd as _, AsRawFd as _};
     let crate::ipc::LocalStream::UdSocket(stream) = stream;