acp runtime: bound the agent's output queue between turns

hive-runtime/src/acp/rpc.rs

@@ -4,7 +4,7 @@
 use std::collections::HashMap;
 use std::path::Path;
 use std::process::Stdio;
-use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::sync::{Arc, Mutex, PoisonError};
 
 use serde_json::{Value, json};
@@ -43,7 +43,7 @@ pub(super) struct Connection {
     stdin: Arc<tokio::sync::Mutex<ChildStdin>>,
     next_id: AtomicU64,
     pending: Pending,
-    pub(super) incoming: mpsc::UnboundedReceiver<Incoming>,
+    pub(super) incoming: mpsc::Receiver<Incoming>,
 }
 
 impl Connection {
@@ -80,13 +80,17 @@ impl Connection {
         };
         let stdin = Arc::new(tokio::sync::Mutex::new(stdin));
         let pending: Pending = Arc::default();
-        let (tx, incoming) = mpsc::unbounded_channel();
+        let (tx, incoming) = mpsc::channel(INCOMING_CAP);
+        let tx = Feed {
+            tx,
+            full: Arc::default(),
+        };
 
         let stderr_tx = tx.clone();
         tokio::spawn(async move {
             let mut lines = BufReader::new(stderr).lines();
             while let Some(line) = next_line(&mut lines, "stderr").await {
-                if stderr_tx.send(Incoming::Stderr(line)).is_err() {
+                if !stderr_tx.push(Incoming::Stderr(line)) {
                     break;
                 }
             }
@@ -110,7 +114,7 @@ impl Connection {
                 .lock()
                 .unwrap_or_else(PoisonError::into_inner)
                 .clear();
-            let _ = reader.tx.send(Incoming::Closed);
+            reader.tx.push(Incoming::Closed);
         });
 
         Ok(Self {
@@ -203,12 +207,45 @@ impl Drop for Connection {
     }
 }
 
+/// How many items wait for the harness before the oldest news is dropped. A
+/// turn drains the queue as it goes, so only an agent that talks while no turn
+/// runs gets near this; its items are thrown away at the next prompt anyway.
+const INCOMING_CAP: usize = 4096;
+
+/// The sending half of the agent's output queue. A full queue drops the new
+/// item, and says so once.
+#[derive(Clone)]
+struct Feed {
+    tx: mpsc::Sender<Incoming>,
+    full: Arc<AtomicBool>,
+}
+
+impl Feed {
+    /// Queues `item`. `false` once the receiver is gone, which ends the
+    /// caller's loop; a full queue drops the item and answers `true`.
+    fn push(&self, item: Incoming) -> bool {
+        match self.tx.try_send(item) {
+            Ok(()) => true,
+            Err(mpsc::error::TrySendError::Full(_)) => {
+                if !self.full.swap(true, Ordering::Relaxed) {
+                    tracing::warn!(
+                        cap = INCOMING_CAP,
+                        "ACP agent output is not being read; dropping what it sends until a turn drains it"
+                    );
+                }
+                true
+            }
+            Err(mpsc::error::TrySendError::Closed(_)) => false,
+        }
+    }
+}
+
 /// The stdout side: routes responses to their waiting request, answers the
 /// agent's own requests, and queues its notifications.
 struct Reader {
     stdin: Arc<tokio::sync::Mutex<ChildStdin>>,
     pending: Pending,
-    tx: mpsc::UnboundedSender<Incoming>,
+    tx: Feed,
     permit: PermissionPolicy,
     servers: Vec<String>,
 }
@@ -217,7 +254,7 @@ impl Reader {
     async fn dispatch(&self, line: &str) {
         let Ok(message) = serde_json::from_str::<Value>(line) else {
             if !line.trim().is_empty() {
-                let _ = self.tx.send(Incoming::Stdout(line.to_owned()));
+                self.tx.push(Incoming::Stdout(line.to_owned()));
             }
             return;
         };
@@ -242,7 +279,7 @@ impl Reader {
             }
             (Some("session/update"), None) => {
                 let params = message.get("params").cloned().unwrap_or(Value::Null);
-                let _ = self.tx.send(Incoming::Update(params));
+                self.tx.push(Incoming::Update(params));
             }
             _ => {}
         }
@@ -353,6 +390,46 @@ pub(super) fn permission_outcome(
 
 #[cfg(test)]
 mod tests {
+    use super::{Feed, Incoming};
+    use tokio::sync::mpsc;
+
+    fn feed(cap: usize) -> (Feed, mpsc::Receiver<Incoming>) {
+        let (tx, rx) = mpsc::channel(cap);
+        (
+            Feed {
+                tx,
+                full: std::sync::Arc::default(),
+            },
+            rx,
+        )
+    }
+
+    #[test]
+    fn a_full_queue_drops_the_new_item_and_keeps_what_waits() {
+        let (feed, mut rx) = feed(2);
+        assert!(feed.push(Incoming::Stderr("one".into())));
+        assert!(feed.push(Incoming::Stderr("two".into())));
+        assert!(
+            feed.push(Incoming::Stderr("three".into())),
+            "a drop is not an error"
+        );
+        let mut got = Vec::new();
+        while let Ok(Incoming::Stderr(line)) = rx.try_recv() {
+            got.push(line);
+        }
+        assert_eq!(got, ["one", "two"]);
+        // Draining makes room again.
+        assert!(feed.push(Incoming::Stderr("four".into())));
+        assert!(matches!(rx.try_recv(), Ok(Incoming::Stderr(l)) if l == "four"));
+    }
+
+    #[test]
+    fn a_gone_receiver_ends_the_pushing_loop() {
+        let (feed, rx) = feed(2);
+        drop(rx);
+        assert!(!feed.push(Incoming::Stderr("x".into())));
+    }
+
     use super::{group_id, next_line, permission_outcome};
     use crate::acp::{PermissionAsk, PermissionPolicy};
     use serde_json::json;