Show Codex background task counts

server/src/session/subagent.rs

@@ -12,7 +12,7 @@
 //! side uses. Only ids matching [`is_subagent_id`] are ever turned into a
 //! path.
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::fs;
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex};
@@ -97,7 +97,7 @@ impl Subagent {
         *self.status.lock().unwrap() != SessionStatus::Exited
     }
 
-    fn append(&self, event: Event) {
+    fn append(&self, event: Event) -> bool {
         let mut transcript = self.transcript.lock().unwrap();
         match transcript.append(event, super::now()) {
             Ok(entry) => {
@@ -107,8 +107,12 @@ impl Subagent {
                 // No subscribers is fine; the transcript already has it,
                 // same as a session's pump.
                 let _ = self.events.send(entry);
+                true
+            }
+            Err(err) => {
+                tracing::error!("subagent transcript append failed: {err:#}");
+                false
             }
-            Err(err) => tracing::error!("subagent transcript append failed: {err:#}"),
         }
     }
 }
@@ -125,14 +129,26 @@ pub struct Subagents {
     /// The session's own directory; subagents live under `<dir>/subagents`.
     dir: PathBuf,
     live: Mutex<HashMap<String, Arc<Subagent>>>,
+    /// Open ids, seeded from disk so an adopted session starts with the measured count rather than
+    /// waiting to see lifecycle edges which are already behind its stdout cursor.
+    open: Mutex<HashSet<String>>,
 }
 
 impl Subagents {
     pub fn new(session_dir: PathBuf) -> Self {
-        Self {
+        let subagents = Self {
             dir: session_dir,
             live: Mutex::new(HashMap::new()),
-        }
+            open: Mutex::new(HashSet::new()),
+        };
+        subagents.open.lock().unwrap().extend(
+            subagents
+                .list(true)
+                .into_iter()
+                .filter(|info| info.status == SessionStatus::Running)
+                .map(|info| info.id),
+        );
+        subagents
     }
 
     fn subagents_dir(&self) -> PathBuf {
@@ -222,6 +238,9 @@ impl Subagents {
         }
         match self.open_or_create(id, title, prompt) {
             Ok(subagent) => {
+                if subagent.is_open() {
+                    self.open.lock().unwrap().insert(id.to_string());
+                }
                 live.insert(id.to_string(), subagent);
             }
             Err(err) => tracing::error!("couldn't start subagent {id}: {err:#}"),
@@ -266,26 +285,21 @@ impl Subagents {
         self.get(id).is_some_and(|subagent| subagent.is_open())
     }
 
-    /// Whether this session has any subagent still working, read from the
-    /// directory rather than from what this process has seen.
+    /// Whether this session has any subagent still working.
     ///
     /// That is the whole point of it. A backend restart adopts a session's
     /// process and picks its stdout back up from a recorded offset, so the
     /// `task_started` lines for subagents launched before the restart are
     /// already behind that offset and the translator never sees them -- it
     /// starts with an empty set and reports the session `Idle` at the end of
-    /// a turn it should have called [`SessionStatus::Waiting`]. Asking the
-    /// registry is a measurement instead of bookkeeping, so it is right for
-    /// a session this process did not start.
-    ///
-    /// Measured rather than cached because the wrong answer has to be able
-    /// to correct itself: a subagent left `Running` by a previous run is
-    /// finished by the session's own exit (see `finish_all`), and the next
-    /// turn to end then reads the truth.
+    /// a turn it should have called [`SessionStatus::Waiting`].
     pub fn any_open(&self, session_running: bool) -> bool {
-        self.list(session_running)
-            .iter()
-            .any(|info| info.status == SessionStatus::Running)
+        session_running && self.open_count() > 0
+    }
+
+    /// Latest measured number of live subagents, including ones found on disk at construction.
+    pub fn open_count(&self) -> usize {
+        self.open.lock().unwrap().len()
     }
 
     /// Appends one event to a subagent's own transcript. A no-op, with a
@@ -294,7 +308,9 @@ impl Subagents {
     /// at.
     pub fn record(&self, id: &str, event: Event) {
         match self.live.lock().unwrap().get(id).cloned() {
-            Some(subagent) => subagent.append(event),
+            Some(subagent) => {
+                subagent.append(event);
+            }
             None => tracing::debug!("dropping an event for unknown subagent {id}"),
         }
     }
@@ -308,10 +324,11 @@ impl Subagents {
     pub fn finish(&self, id: &str) {
         if let Some(subagent) = self.live.lock().unwrap().get(id).cloned()
             && subagent.is_open()
-        {
-            subagent.append(Event::Status {
+            && subagent.append(Event::Status {
                 state: SessionStatus::Exited,
-            });
+            })
+        {
+            self.open.lock().unwrap().remove(id);
         }
     }
 
@@ -323,10 +340,11 @@ impl Subagents {
     pub fn reopen(&self, id: &str) {
         if let Some(subagent) = self.live.lock().unwrap().get(id).cloned()
             && !subagent.is_open()
-        {
-            subagent.append(Event::Status {
+            && subagent.append(Event::Status {
                 state: SessionStatus::Running,
-            });
+            })
+        {
+            self.open.lock().unwrap().insert(id.to_string());
         }
     }
 
@@ -348,13 +366,8 @@ impl Subagents {
             if info.status != SessionStatus::Running {
                 continue;
             }
-            if let Some(subagent) = self.get(&info.id)
-                && subagent.is_open()
-            {
-                subagent.append(Event::Status {
-                    state: SessionStatus::Exited,
-                });
-            }
+            let _ = self.get(&info.id);
+            self.finish(&info.id);
         }
     }
 
@@ -399,6 +412,7 @@ impl Subagents {
             // line for this id starts a new subagent rather than appending
             // to an unlinked file nothing can read.
             live.remove(id);
+            self.open.lock().unwrap().remove(id);
         }
         Ok(())
     }
@@ -501,6 +515,7 @@ mod tests {
                 .count(),
             1
         );
+        assert_eq!(subagents.open_count(), 1);
     }
 
     #[test]
@@ -518,6 +533,7 @@ mod tests {
         }
         // A fresh registry, the way a backend restart builds one.
         let subagents = Subagents::new(dir.path().to_path_buf());
+        assert_eq!(subagents.open_count(), 1);
         let subagent = subagents.get("toolu_2").expect("reopened");
         assert!(subagent.is_open());
         subagents.record(
@@ -539,6 +555,7 @@ mod tests {
         let subagents = Subagents::new(dir.path().to_path_buf());
         subagents.start("toolu_3", "helper", None);
         subagents.finish("toolu_3");
+        assert_eq!(subagents.open_count(), 0);
         let subagent = subagents.get("toolu_3").unwrap();
         assert!(!subagent.is_open());
         // On disk too, not only in the live cache `is_open` reads.
@@ -551,6 +568,11 @@ mod tests {
 
         // Finishing an id that was never a subagent is a no-op, not a panic.
         subagents.finish("never-started");
+
+        subagents.reopen("toolu_3");
+        assert_eq!(subagents.open_count(), 1);
+        subagents.finish("toolu_3");
+        assert_eq!(subagents.open_count(), 0);
     }
 
     #[test]