From 66e13ea53567919e91c1106833e25d46b43ea537 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20M=C3=BCller?= <sebastian@primeintellect.ai>
Date: Thu, 1 Oct 2026 15:08:21 -0700
Subject: [PATCH] fix(pa-cli): keep partial daemon lines across socket poll
 timeouts (#3284)

read_line polled the daemon socket with a 50ms read timeout and started
a fresh buffer on every poll iteration. BufRead::read_line keeps the
bytes that arrived before a poll times out and returns WouldBlock, so
the next iteration dropped the start of a slow line and the client
parsed only its tail. A truncated tail is unparseable JSON that callers
skip, so the command hung until its 30s timeout (15s for the handshake)
instead of showing a result, e.g. `list` against a busy daemon sending
a large response slowly.

Keep the partial bytes in a byte buffer on the client across polls and
across reads (probe_daemon keeps using a client after a swallowed hello
timeout), fill it with read_until, and decode the line as UTF-8 only
when it is complete: decoding per poll drops the bytes again when a
timeout lands inside a multi-byte character. Deadlines and error
messages are unchanged.

Regression test: read_line_keeps_bytes_across_poll_timeouts replays a
line split by poll timeouts, one inside a multi-byte character, with
no sleeps and no production test hooks.
---
 .changes/fix-daemon-client-partial-line.md |  1 +
 crates/pa-cli/src/daemon_client.rs         | 77 +++++++++++++++++++++-
 2 files changed, 76 insertions(+), 2 deletions(-)
 create mode 100644 .changes/fix-daemon-client-partial-line.md

diff --git a/.changes/fix-daemon-client-partial-line.md b/.changes/fix-daemon-client-partial-line.md
new file mode 100644
index 000000000..045544fc4
--- /dev/null
+++ b/.changes/fix-daemon-client-partial-line.md
@@ -0,0 +1 @@
+- `prime-agent list` and other daemon commands no longer hang until they time out when a busy daemon sends a large response slowly.
diff --git a/crates/pa-cli/src/daemon_client.rs b/crates/pa-cli/src/daemon_client.rs
index d67fbf221..69b8f80a8 100644
--- a/crates/pa-cli/src/daemon_client.rs
+++ b/crates/pa-cli/src/daemon_client.rs
@@ -37,6 +37,9 @@ const READ_POLL: Duration = Duration::from_millis(50);
 pub(crate) struct DaemonClient {
     socket_path: PathBuf,
     reader: BufReader<Box<dyn BlockingTransportStream>>,
+    /// Bytes of a line whose newline has not arrived yet, kept across poll
+    /// timeouts and reads.
+    partial_line: Vec<u8>,
     writer: Box<dyn BlockingTransportStream>,
     hello: Option<serde_json::Value>,
     daemon_closing_reason: Option<String>,
@@ -69,6 +72,7 @@ impl DaemonClient {
         Ok(DaemonClient {
             socket_path: socket_path.to_path_buf(),
             reader: BufReader::new(stream),
+            partial_line: Vec::new(),
             writer,
             hello: None,
             daemon_closing_reason: None,
@@ -183,8 +187,7 @@ impl DaemonClient {
                 .get_mut()
                 .set_read_timeout(READ_POLL)
                 .map_err(|error| anyhow!("daemon socket error: {error}"))?;
-            let mut line = String::new();
-            let read = self.reader.read_line(&mut line);
+            let read = self.reader.read_until(b'\n', &mut self.partial_line);
             match read {
                 Ok(0) => {
                     let reason = self
@@ -198,6 +201,9 @@ impl DaemonClient {
                     ));
                 }
                 Ok(_) => {
+                    let line = String::from_utf8(std::mem::take(&mut self.partial_line)).map_err(
+                        |_| anyhow!("daemon socket error: stream did not contain valid UTF-8"),
+                    )?;
                     if line.trim().is_empty() {
                         continue;
                     }
@@ -367,4 +373,71 @@ mod tests {
         assert!(text.contains("Socket: "), "{text}");
         server.join().unwrap();
     }
+
+    /// Replays scripted reads; `None` is a poll timeout.
+    #[derive(Debug)]
+    struct ScriptedReads(std::collections::VecDeque<Option<&'static [u8]>>);
+
+    impl std::io::Read for ScriptedReads {
+        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+            match self.0.pop_front() {
+                Some(Some(bytes)) => {
+                    buf[..bytes.len()].copy_from_slice(bytes);
+                    Ok(bytes.len())
+                }
+                Some(None) => Err(std::io::ErrorKind::WouldBlock.into()),
+                None => Ok(0),
+            }
+        }
+    }
+
+    impl std::io::Write for ScriptedReads {
+        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+            Ok(buf.len())
+        }
+
+        fn flush(&mut self) -> std::io::Result<()> {
+            Ok(())
+        }
+    }
+
+    impl BlockingTransportStream for ScriptedReads {
+        fn try_clone_box(&self) -> std::io::Result<Box<dyn BlockingTransportStream>> {
+            Err(std::io::ErrorKind::Unsupported.into())
+        }
+
+        fn set_read_timeout(&self, _: Duration) -> std::io::Result<()> {
+            Ok(())
+        }
+    }
+
+    /// A line split by poll timeouts, one of them inside a multi-byte
+    /// character, comes back whole.
+    #[test]
+    fn read_line_keeps_bytes_across_poll_timeouts() {
+        let dir = tempfile::TempDir::new().unwrap();
+        let socket = dir.path().join("daemon.sock");
+        let _listener = UnixListener::bind(&socket).unwrap();
+        let mut client = DaemonClient::connect(&socket).unwrap();
+        let line: &'static [u8] = "{\"type\":\"daemon_hello\",\"name\":\"é\"}\n".as_bytes();
+        let split = line.len() - 4; // between the two bytes of "é"
+        client.reader = BufReader::new(Box::new(ScriptedReads(
+            [
+                Some(&line[..10]),
+                None,
+                Some(&line[10..split]),
+                None,
+                Some(&line[split..]),
+            ]
+            .into(),
+        )));
+        let read = client
+            .read_line(
+                Instant::now() + Duration::from_secs(10),
+                10_000,
+                Operation::Handshake,
+            )
+            .unwrap();
+        assert_eq!(read.as_bytes(), line);
+    }
 }
-- 
knotmirror 0.0.1