2 files changed, 46 insertions, 24 deletions

diff --git a/niri-ipc/src/socket.rs b/niri-ipc/src/socket.rs
index d629f1a4..790097cb 100644
--- a/niri-ipc/src/socket.rs
+++ b/niri-ipc/src/socket.rs
@@ -16,7 +16,7 @@ pub const SOCKET_PATH_ENV: &str = "NIRI_SOCKET";
 /// This struct is used to communicate with the niri IPC server. It handles the socket connection
 /// and serialization/deserialization of messages.
 pub struct Socket {
-    stream: UnixStream,
+    stream: BufReader<UnixStream>,
 }
 
 impl Socket {
@@ -37,6 +37,7 @@ impl Socket {
     /// Connects to the niri IPC socket at the given path.
     pub fn connect_to(path: impl AsRef<Path>) -> io::Result<Self> {
         let stream = UnixStream::connect(path.as_ref())?;
+        let stream = BufReader::new(stream);
         Ok(Self { stream })
     }
 
@@ -47,31 +48,54 @@ impl Socket {
     /// * `Ok(Ok(response))`: successful [`Response`](crate::Response) from niri
     /// * `Ok(Err(message))`: error message from niri
     /// * `Err(error)`: error communicating with niri
-    ///
-    /// This method also returns a blocking function that you can call to keep reading [`Event`]s
-    /// after requesting an [`EventStream`][Request::EventStream]. This function is not useful
-    /// otherwise.
-    pub fn send(self, request: Request) -> io::Result<(Reply, impl FnMut() -> io::Result<Event>)> {
-        let Self { mut stream } = self;
-
+    pub fn send(&mut self, request: Request) -> io::Result<Reply> {
         let mut buf = serde_json::to_string(&request).unwrap();
-        stream.write_all(buf.as_bytes())?;
-        stream.shutdown(Shutdown::Write)?;
-
-        let mut reader = BufReader::new(stream);
+        buf.push('\n');
+        self.stream.get_mut().write_all(buf.as_bytes())?;
 
         buf.clear();
-        reader.read_line(&mut buf)?;
+        self.stream.read_line(&mut buf)?;
 
         let reply = serde_json::from_str(&buf)?;
+        Ok(reply)
+    }
 
-        let events = move || {
+    /// Starts reading event stream [`Event`]s from the socket.
+    ///
+    /// The returned function will block until the next [`Event`] arrives, then return it.
+    ///
+    /// Use this only after requesting an [`EventStream`][Request::EventStream].
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// use niri_ipc::{Request, Response};
+    /// use niri_ipc::socket::Socket;
+    ///
+    /// fn main() -> std::io::Result<()> {
+    ///     let mut socket = Socket::connect()?;
+    ///
+    ///     let reply = socket.send(Request::EventStream)?;
+    ///     if matches!(reply, Ok(Response::Handled)) {
+    ///         let mut read_event = socket.read_events();
+    ///         while let Ok(event) = read_event() {
+    ///             println!("Received event: {event:?}");
+    ///         }
+    ///     }
+    ///
+    ///     Ok(())
+    /// }
+    /// ```
+    pub fn read_events(self) -> impl FnMut() -> io::Result<Event> {
+        let Self { mut stream } = self;
+        let _ = stream.get_mut().shutdown(Shutdown::Write);
+
+        let mut buf = String::new();
+        move || {
             buf.clear();
-            reader.read_line(&mut buf)?;
+            stream.read_line(&mut buf)?;
             let event = serde_json::from_str(&buf)?;
             Ok(event)
-        };
-
-        Ok((reply, events))
+        }
     }
 }
diff --git a/src/ipc/client.rs b/src/ipc/client.rs
index 2eab7e15..4e0357d4 100644
--- a/src/ipc/client.rs
+++ b/src/ipc/client.rs
@@ -35,9 +35,9 @@ pub fn handle_msg(msg: Msg, json: bool) -> anyhow::Result<()> {
         Msg::OverviewState => Request::OverviewState,
     };
 
-    let socket = Socket::connect().context("error connecting to the niri socket")?;
+    let mut socket = Socket::connect().context("error connecting to the niri socket")?;
 
-    let (reply, mut read_event) = socket
+    let reply = socket
         .send(request)
         .context("error communicating with niri")?;
 
@@ -45,10 +45,7 @@ pub fn handle_msg(msg: Msg, json: bool) -> anyhow::Result<()> {
         Err(_) if !matches!(msg, Msg::Version) => {
             // If we got an error, it might be that the CLI is a different version from the running
             // niri instance. Request the running instance version to compare and print a message.
-            Socket::connect()
-                .and_then(|socket| socket.send(Request::Version))
-                .ok()
-                .map(|(reply, _read_event)| reply)
+            socket.send(Request::Version).ok()
         }
         _ => None,
     };
@@ -392,6 +389,7 @@ pub fn handle_msg(msg: Msg, json: bool) -> anyhow::Result<()> {
                 println!("Started reading events.");
             }
 
+            let mut read_event = socket.read_events();
             loop {
                 let event = read_event().context("error reading event from niri")?;