Count Codex background terminals

server/src/session/codex.rs

@@ -9,7 +9,7 @@
 
 mod translate;
 
-use std::collections::VecDeque;
+use std::collections::{HashSet, VecDeque};
 use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex};
@@ -35,6 +35,22 @@ const STDERR_LOG: &str = "codex-stderr.log";
 const THREAD_FILE: &str = "codex-thread.json";
 const STATE_FILE: &str = "codex-state.json";
 const POLL: std::time::Duration = std::time::Duration::from_millis(50);
+const BACKGROUND_POLL: std::time::Duration = std::time::Duration::from_secs(1);
+
+type BackgroundProcesses = Arc<Mutex<HashSet<String>>>;
+
+#[derive(Default)]
+struct BackgroundQuery {
+    pending: Option<PendingBackgroundQuery>,
+    unsupported: bool,
+}
+
+struct PendingBackgroundQuery {
+    request: String,
+    thread: String,
+    found: HashSet<String>,
+    dirty: bool,
+}
 
 #[derive(Clone, Default, Serialize, Deserialize)]
 struct Waiting {
@@ -108,6 +124,8 @@ struct Inner {
     transport: Transport,
     session_dir: PathBuf,
     subagents: Arc<Subagents>,
+    background_processes: BackgroundProcesses,
+    background_query: Mutex<BackgroundQuery>,
     reading: AtomicBool,
 }
 
@@ -175,6 +193,7 @@ impl CodexDriver {
             }
         });
 
+        let background_processes = Arc::new(Mutex::new(HashSet::new()));
         let inner = Arc::new(Inner {
             sink,
             state: Mutex::new(state),
@@ -187,24 +206,30 @@ impl CodexDriver {
             transport,
             session_dir: session_dir.to_path_buf(),
             subagents,
+            background_processes,
+            background_query: Mutex::new(BackgroundQuery::default()),
             reading: AtomicBool::new(true),
         });
 
         if started_here {
             begin_initialization(&inner);
+            let _ = inner.sink.send(Event::BackgroundTasks { count: 0 });
             let _ = inner.sink.send(Event::Status {
                 state: SessionStatus::Idle,
             });
         }
         save_state(&inner);
         spawn_follower(Arc::clone(&inner), record);
+        if !started_here {
+            refresh_background_terminals(&inner);
+        }
         Ok(Self { inner })
     }
 }
 
 impl Driver for CodexDriver {
     fn background_tasks(&self) -> Option<usize> {
-        Some(self.inner.subagents.open_count())
+        Some(background_task_count(&self.inner))
     }
 
     fn send_user_message(&self, text: String, attachments: Vec<AttachmentRef>) {
@@ -441,11 +466,14 @@ fn begin_initialization(inner: &Arc<Inner>) {
         json!({
             "id": id,
             "method": "initialize",
-            "params": {"clientInfo": {
-                "name": "ai-app",
-                "title": "AI Sessions",
-                "version": env!("CARGO_PKG_VERSION")
-            }}
+            "params": {
+                "clientInfo": {
+                    "name": "ai-app",
+                    "title": "AI Sessions",
+                    "version": env!("CARGO_PKG_VERSION")
+                },
+                "capabilities": {"experimentalApi": true}
+            }
         }),
     );
 }
@@ -624,11 +652,150 @@ fn spawn_follower(inner: Arc<Inner>, record: process::Record) {
     tokio::spawn(follow(inner, record, offset));
 }
 
+fn background_task_count(inner: &Inner) -> usize {
+    inner.subagents.open_count() + inner.background_processes.lock().unwrap().len()
+}
+
+fn refresh_background_terminals(inner: &Inner) {
+    let Some(thread) = read_thread(&inner.session_dir) else {
+        return;
+    };
+    let request = {
+        let mut query = inner.background_query.lock().unwrap();
+        if query.unsupported {
+            return;
+        }
+        if let Some(pending) = &mut query.pending {
+            pending.dirty = true;
+            return;
+        }
+        let request = request_id();
+        query.pending = Some(PendingBackgroundQuery {
+            request: request.clone(),
+            thread: thread.clone(),
+            found: HashSet::new(),
+            dirty: false,
+        });
+        request
+    };
+    send_json(
+        inner,
+        json!({
+            "id": request,
+            "method": "thread/backgroundTerminals/list",
+            "params": {"threadId": thread, "limit": 100}
+        }),
+    );
+}
+
+fn handle_background_response(inner: &Inner, line: &Value) -> bool {
+    let Some(id) = line.get("id").and_then(Value::as_str) else {
+        return false;
+    };
+    let mut query = inner.background_query.lock().unwrap();
+    let Some(pending) = &mut query.pending else {
+        return false;
+    };
+    if pending.request != id {
+        return false;
+    }
+    if Some(&pending.thread) != read_thread(&inner.session_dir).as_ref() {
+        query.pending = None;
+        drop(query);
+        refresh_background_terminals(inner);
+        return true;
+    }
+    if line.get("error").is_some() {
+        query.unsupported = line.pointer("/error/code").and_then(Value::as_i64) == Some(-32601);
+        query.pending = None;
+        return true;
+    }
+    let Some(terminals) = line.pointer("/result/data").and_then(Value::as_array) else {
+        query.pending = None;
+        return true;
+    };
+    if !terminals
+        .iter()
+        .all(|terminal| terminal.get("processId").and_then(Value::as_str).is_some())
+    {
+        query.pending = None;
+        return true;
+    }
+    for process in terminals
+        .iter()
+        .filter_map(|terminal| terminal.get("processId").and_then(Value::as_str))
+    {
+        pending.found.insert(process.to_string());
+    }
+    if let Some(cursor) = line.pointer("/result/nextCursor").and_then(Value::as_str) {
+        let thread = pending.thread.clone();
+        let request = request_id();
+        pending.request.clone_from(&request);
+        drop(query);
+        send_json(
+            inner,
+            json!({
+                "id": request,
+                "method": "thread/backgroundTerminals/list",
+                "params": {"threadId": thread, "limit": 100, "cursor": cursor}
+            }),
+        );
+        return true;
+    }
+
+    let pending = query.pending.take().expect("matched pending query");
+    drop(query);
+
+    if pending.dirty {
+        refresh_background_terminals(inner);
+        return true;
+    }
+    let found = pending.found;
+
+    let changed = {
+        let mut processes = inner.background_processes.lock().unwrap();
+        if *processes == found {
+            false
+        } else {
+            *processes = found;
+            true
+        }
+    };
+    if changed {
+        let count = background_task_count(inner);
+        let _ = inner.sink.send(Event::BackgroundTasks { count });
+        let state = inner.state.lock().unwrap();
+        if !state.running && !state.closed {
+            let _ = inner.sink.send(Event::Status {
+                state: if count > 0 {
+                    SessionStatus::Waiting
+                } else {
+                    SessionStatus::Idle
+                },
+            });
+        }
+    }
+    true
+}
+
+fn background_terminals_may_have_changed(method: Option<&str>, params: &Value) -> bool {
+    matches!(
+        method,
+        Some("thread/started" | "turn/completed" | "item/commandExecution/terminalInteraction")
+    ) || matches!(method, Some("item/completed" | "item.completed"))
+        && params
+            .get("item")
+            .and_then(|item| item.get("type"))
+            .and_then(Value::as_str)
+            == Some("commandExecution")
+}
+
 async fn follow(inner: Arc<Inner>, mut record: process::Record, mut offset: u64) {
     let stdout = inner.session_dir.join(STDOUT_LOG);
     let stderr = inner.session_dir.join(STDERR_LOG);
     let mut translator = Translator::new(
         Arc::clone(&inner.subagents),
+        Arc::clone(&inner.background_processes),
         read_thread(&inner.session_dir),
         inner.state.lock().unwrap().active_turn.is_some(),
     );
@@ -638,6 +805,7 @@ async fn follow(inner: Arc<Inner>, mut record: process::Record, mut offset: u64)
     // quadratic in time and allocation.
     let mut read_at = offset;
     let mut pending = Vec::new();
+    let mut last_background_poll = std::time::Instant::now();
     while inner.reading.load(Ordering::SeqCst) {
         let (bytes, next) = match process::read_from(&stdout, read_at) {
             Ok(read) => read,
@@ -673,6 +841,12 @@ async fn follow(inner: Arc<Inner>, mut record: process::Record, mut offset: u64)
             };
             process::write(&inner.session_dir, &record);
         }
+        if !inner.background_processes.lock().unwrap().is_empty()
+            && last_background_poll.elapsed() >= BACKGROUND_POLL
+        {
+            refresh_background_terminals(&inner);
+            last_background_poll = std::time::Instant::now();
+        }
         match record.liveness() {
             process::Liveness::Alive | process::Liveness::Unknown => {}
             process::Liveness::Dead if complete > 0 => {}
@@ -713,7 +887,7 @@ async fn follow(inner: Arc<Inner>, mut record: process::Record, mut offset: u64)
 }
 
 fn handle_line(inner: &Arc<Inner>, translator: &mut Translator, line: &Value) {
-    if line.get("id").is_some() {
+    if line.get("id").is_some() && !handle_background_response(inner, line) {
         handle_response(inner, line);
     }
     let method = line.get("method").and_then(Value::as_str);
@@ -781,6 +955,9 @@ fn handle_line(inner: &Arc<Inner>, translator: &mut Translator, line: &Value) {
     if parent && method == Some("turn/completed") {
         dispatch_waiting(inner);
     }
+    if parent && background_terminals_may_have_changed(method, params) {
+        refresh_background_terminals(inner);
+    }
 }
 
 /// Images embedded in a structured tool result, copied into the session before the translator's
@@ -1364,6 +1541,112 @@ mod tests {
         assert!(!missing_thread("thread is already running"));
     }
 
+    #[test]
+    fn background_terminal_snapshots_are_id_keyed_and_ignore_unknown_responses() {
+        let dir = tempfile::tempdir().expect("tempdir");
+        write_thread(dir.path(), "parent-thread");
+        let (sink, mut events) = mpsc::unbounded_channel();
+        let (to_child, mut requests) = mpsc::unbounded_channel();
+        let inner = Inner {
+            sink,
+            state: Mutex::new(ProtocolState::default()),
+            settings: Mutex::new(Settings {
+                model: None,
+                permission_mode: None,
+                effort: None,
+            }),
+            to_child,
+            transport: Transport::Here,
+            session_dir: dir.path().to_path_buf(),
+            subagents: Arc::new(Subagents::new(dir.path().to_path_buf())),
+            background_processes: Arc::new(Mutex::new(HashSet::new())),
+            background_query: Mutex::new(BackgroundQuery::default()),
+            reading: AtomicBool::new(true),
+        };
+
+        refresh_background_terminals(&inner);
+        let request: Value = serde_json::from_str(&requests.try_recv().expect("list request"))
+            .expect("request json");
+        assert_eq!(request["method"], "thread/backgroundTerminals/list");
+        let request_id = request["id"].as_str().expect("request id");
+        assert!(handle_background_response(
+            &inner,
+            &json!({
+                "id": request_id,
+                "result": {
+                    "data": [
+                        {"processId": "process-a"},
+                        {"processId": "process-b"}
+                    ],
+                    "nextCursor": null
+                }
+            })
+        ));
+        assert_eq!(background_task_count(&inner), 2);
+        assert_eq!(
+            events.try_recv().expect("count"),
+            Event::BackgroundTasks { count: 2 }
+        );
+        assert_eq!(
+            events.try_recv().expect("waiting status"),
+            Event::Status {
+                state: SessionStatus::Waiting
+            }
+        );
+
+        assert!(!handle_background_response(
+            &inner,
+            &json!({"id": "unknown-completion", "result": {"data": []}})
+        ));
+        assert_eq!(background_task_count(&inner), 2);
+        assert!(events.try_recv().is_err());
+
+        refresh_background_terminals(&inner);
+        let request: Value = serde_json::from_str(&requests.try_recv().expect("malformed list"))
+            .expect("request json");
+        assert!(handle_background_response(
+            &inner,
+            &json!({"id": request["id"], "result": {}})
+        ));
+        assert_eq!(background_task_count(&inner), 2);
+        assert!(events.try_recv().is_err());
+
+        refresh_background_terminals(&inner);
+        let request: Value =
+            serde_json::from_str(&requests.try_recv().expect("second list")).expect("request json");
+        refresh_background_terminals(&inner);
+        assert!(handle_background_response(
+            &inner,
+            &json!({
+                "id": request["id"],
+                "result": {
+                    "data": [{"processId": "process-b"}],
+                    "nextCursor": null
+                }
+            })
+        ));
+        assert_eq!(background_task_count(&inner), 2);
+        assert!(events.try_recv().is_err());
+
+        let request: Value = serde_json::from_str(&requests.try_recv().expect("coalesced list"))
+            .expect("request json");
+        assert!(handle_background_response(
+            &inner,
+            &json!({
+                "id": request["id"],
+                "result": {
+                    "data": [{"processId": "process-b"}],
+                    "nextCursor": null
+                }
+            })
+        ));
+        assert_eq!(background_task_count(&inner), 1);
+        assert_eq!(
+            events.try_recv().expect("smaller count"),
+            Event::BackgroundTasks { count: 1 }
+        );
+    }
+
     #[test]
     fn a_missing_thread_is_reported_then_replaced() {
         let dir = tempfile::tempdir().expect("tempdir");
@@ -1385,6 +1668,8 @@ mod tests {
             transport: Transport::Here,
             session_dir: dir.path().to_path_buf(),
             subagents: Arc::new(Subagents::new(dir.path().to_path_buf())),
+            background_processes: Arc::new(Mutex::new(HashSet::new())),
+            background_query: Mutex::new(BackgroundQuery::default()),
             reading: AtomicBool::new(true),
         });
 
@@ -1440,6 +1725,8 @@ mod tests {
             transport: Transport::Here,
             session_dir: dir.path().to_path_buf(),
             subagents: Arc::new(Subagents::new(dir.path().to_path_buf())),
+            background_processes: Arc::new(Mutex::new(HashSet::new())),
+            background_query: Mutex::new(BackgroundQuery::default()),
             reading: AtomicBool::new(true),
         });
         let driver = CodexDriver {
@@ -1505,6 +1792,8 @@ mod tests {
             transport: Transport::Here,
             session_dir: dir.path().to_path_buf(),
             subagents: Arc::new(Subagents::new(dir.path().to_path_buf())),
+            background_processes: Arc::new(Mutex::new(HashSet::new())),
+            background_query: Mutex::new(BackgroundQuery::default()),
             reading: AtomicBool::new(true),
         });
         let driver = CodexDriver {
@@ -1568,10 +1857,13 @@ mod tests {
             transport: Transport::Here,
             session_dir: dir.path().to_path_buf(),
             subagents: Arc::new(Subagents::new(dir.path().to_path_buf())),
+            background_processes: Arc::new(Mutex::new(HashSet::new())),
+            background_query: Mutex::new(BackgroundQuery::default()),
             reading: AtomicBool::new(true),
         });
         let mut translator = Translator::new(
             Arc::clone(&inner.subagents),
+            Arc::clone(&inner.background_processes),
             Some("missing-thread".to_string()),
             false,
         );
@@ -1636,6 +1928,8 @@ mod tests {
             transport: Transport::Here,
             session_dir: dir.path().to_path_buf(),
             subagents: Arc::new(Subagents::new(dir.path().to_path_buf())),
+            background_processes: Arc::new(Mutex::new(HashSet::new())),
+            background_query: Mutex::new(BackgroundQuery::default()),
             reading: AtomicBool::new(true),
         });
 

server/src/session/codex/translate.rs

@@ -5,7 +5,7 @@
 //! ignore the rest; an added Codex item must not make a live session go deaf.
 
 use std::collections::{HashMap, HashSet};
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
 use serde_json::{Value, json};
 
@@ -18,6 +18,7 @@ pub(super) struct Translator {
     completed: bool,
     limited: bool,
     subagents: Option<Arc<Subagents>>,
+    background_processes: Option<Arc<Mutex<HashSet<String>>>>,
     children: HashMap<String, Translator>,
     prompts: HashMap<String, String>,
     async_messages: HashSet<String>,
@@ -25,10 +26,16 @@ pub(super) struct Translator {
 }
 
 impl Translator {
-    pub(super) fn new(subagents: Arc<Subagents>, thread_id: Option<String>, in_turn: bool) -> Self {
+    pub(super) fn new(
+        subagents: Arc<Subagents>,
+        background_processes: Arc<Mutex<HashSet<String>>>,
+        thread_id: Option<String>,
+        in_turn: bool,
+    ) -> Self {
         Self {
             thread_id,
             subagents: Some(subagents),
+            background_processes: Some(background_processes),
             in_turn,
             ..Self::default()
         }
@@ -118,9 +125,14 @@ impl Translator {
     }
 
     fn background_task_count(&self) -> Option<usize> {
-        self.subagents
-            .as_ref()
-            .map(|subagents| subagents.open_count())
+        self.subagents.as_ref().map(|subagents| {
+            subagents.open_count()
+                + self
+                    .background_processes
+                    .as_ref()
+                    .map(|processes| processes.lock().unwrap().len())
+                    .unwrap_or(0)
+        })
     }
 
     fn translate_line(&mut self, kind: Option<&str>, body: &Value, line: &Value) -> Vec<Event> {
@@ -242,11 +254,7 @@ impl Translator {
                     events.extend(self.failure(error));
                 }
                 events.push(Event::Status {
-                    state: if self
-                        .subagents
-                        .as_ref()
-                        .is_some_and(|subagents| subagents.any_open(true))
-                    {
+                    state: if self.background_task_count().is_some_and(|count| count > 0) {
                         SessionStatus::Waiting
                     } else {
                         SessionStatus::Idle
@@ -390,13 +398,7 @@ impl Translator {
                     subagents.finish(id);
                 }
                 self.prompts.remove(id);
-                if was_open
-                    && !self.in_turn
-                    && self
-                        .subagents
-                        .as_ref()
-                        .is_some_and(|subagents| !subagents.any_open(true))
-                {
+                if was_open && !self.in_turn && self.background_task_count() == Some(0) {
                     vec![Event::Status {
                         state: SessionStatus::Idle,
                     }]
@@ -1216,8 +1218,10 @@ mod tests {
     fn codex_subagents_get_their_own_transcripts_and_hold_the_parent_waiting() {
         let dir = tempfile::tempdir().expect("tempdir");
         let subagents = Arc::new(Subagents::new(dir.path().to_path_buf()));
+        let background_processes = Arc::new(Mutex::new(HashSet::new()));
         let mut translator = Translator::new(
             Arc::clone(&subagents),
+            Arc::clone(&background_processes),
             Some("parent-thread".to_string()),
             false,
         );
@@ -1313,8 +1317,10 @@ mod tests {
     fn codex_reports_each_change_to_its_live_background_count() {
         let dir = tempfile::tempdir().expect("tempdir");
         let subagents = Arc::new(Subagents::new(dir.path().to_path_buf()));
+        let background_processes = Arc::new(Mutex::new(HashSet::new()));
         let mut translator = Translator::new(
             Arc::clone(&subagents),
+            Arc::clone(&background_processes),
             Some("parent-thread".to_string()),
             false,
         );
@@ -1335,7 +1341,11 @@ mod tests {
             }));
             assert_eq!(events.last(), Some(&Event::BackgroundTasks { count }));
         }
-        for (child, count) in [("child-a", 1), ("child-b", 0)] {
+        background_processes
+            .lock()
+            .unwrap()
+            .insert("command-a".to_string());
+        for (child, count) in [("child-a", 2), ("child-b", 1)] {
             let events = translator.translate(&json!({
                 "method": "item/completed",
                 "params": {
@@ -1350,6 +1360,14 @@ mod tests {
                 }
             }));
             assert_eq!(events.last(), Some(&Event::BackgroundTasks { count }));
+            assert!(!events.iter().any(|event| {
+                matches!(
+                    event,
+                    Event::Status {
+                        state: SessionStatus::Idle
+                    }
+                )
+            }));
         }
     }
 
@@ -1490,6 +1508,7 @@ mod tests {
         subagents.start("child-thread", "child", Some("work"));
         let mut translator = Translator::new(
             Arc::clone(&subagents),
+            Arc::new(Mutex::new(HashSet::new())),
             Some("parent-thread".to_string()),
             true,
         );