Show Codex background task counts
This commit is contained in:
1 parent
947ea8ecf2
commit
a9cfea89e5
5 files changed
+155
-52
No files matched your search
@@ -137,6 +137,9 @@ impl CodexDriver {
|
||||
}
|
||||
Some((_, process::Liveness::Dead)) | None => {
|
||||
process::clear(session_dir);
|
||||
// Every child thread lived in the app-server which is gone. Clear adopted open
|
||||
// ids before the replacement process reports its background count.
|
||||
subagents.finish_all();
|
||||
// Requests written to a dead process have no recipient. Put their messages back
|
||||
// in front of the unsent queue so restarting cannot silently lose them.
|
||||
while let Some(message) = state.sent.pop_back() {
|
||||
@@ -200,6 +203,10 @@ impl CodexDriver {
|
||||
}
|
||||
|
||||
impl Driver for CodexDriver {
|
||||
fn background_tasks(&self) -> Option<usize> {
|
||||
Some(self.inner.subagents.open_count())
|
||||
}
|
||||
|
||||
fn send_user_message(&self, text: String, attachments: Vec<AttachmentRef>) {
|
||||
let mut state = self.inner.state.lock().unwrap();
|
||||
if state.closed {
|
||||
|
||||
@@ -64,11 +64,19 @@ impl Translator {
|
||||
/// Translates one multiplexed app-server record. `prefix` contains image
|
||||
/// events extracted by the driver and must precede the item's ToolEnd in
|
||||
/// whichever transcript owns the record.
|
||||
pub(super) fn translate_with_prefix(
|
||||
&mut self,
|
||||
line: &Value,
|
||||
mut prefix: Vec<Event>,
|
||||
) -> Vec<Event> {
|
||||
pub(super) fn translate_with_prefix(&mut self, line: &Value, prefix: Vec<Event>) -> Vec<Event> {
|
||||
let before = self.background_task_count();
|
||||
let mut events = self.translate_inner(line, prefix);
|
||||
let after = self.background_task_count();
|
||||
if after != before
|
||||
&& let Some(count) = after
|
||||
{
|
||||
events.push(Event::BackgroundTasks { count });
|
||||
}
|
||||
events
|
||||
}
|
||||
|
||||
fn translate_inner(&mut self, line: &Value, mut prefix: Vec<Event>) -> Vec<Event> {
|
||||
let method = line.get("method").and_then(Value::as_str);
|
||||
let body = method.and_then(|_| line.get("params")).unwrap_or(line);
|
||||
let kind = method.or_else(|| line.get("type").and_then(Value::as_str));
|
||||
@@ -109,6 +117,12 @@ impl Translator {
|
||||
prefix
|
||||
}
|
||||
|
||||
fn background_task_count(&self) -> Option<usize> {
|
||||
self.subagents
|
||||
.as_ref()
|
||||
.map(|subagents| subagents.open_count())
|
||||
}
|
||||
|
||||
fn translate_line(&mut self, kind: Option<&str>, body: &Value, line: &Value) -> Vec<Event> {
|
||||
match kind {
|
||||
Some("thread.started") => {
|
||||
@@ -1222,10 +1236,13 @@ mod tests {
|
||||
translator.translate(&line(
|
||||
r#"{"method":"item/completed","params":{"threadId":"parent-thread","turnId":"turn-1","item":{"id":"spawn-1","type":"subAgentActivity","kind":"started","agentThreadId":"child-thread","agentPath":"/root/history_boundaries"}}}"#
|
||||
)),
|
||||
vec![Event::ToolEnd {
|
||||
id: "spawn-1".to_string(),
|
||||
output: String::new()
|
||||
}]
|
||||
vec![
|
||||
Event::ToolEnd {
|
||||
id: "spawn-1".to_string(),
|
||||
output: String::new()
|
||||
},
|
||||
Event::BackgroundTasks { count: 1 }
|
||||
]
|
||||
);
|
||||
let rows = subagents.list(true);
|
||||
assert_eq!(rows.len(), 1);
|
||||
@@ -1265,9 +1282,12 @@ mod tests {
|
||||
translator.translate(&line(
|
||||
r#"{"method":"item/completed","params":{"threadId":"parent-thread","turnId":"turn-1","item":{"id":"subagent-completed-1","type":"subAgentActivity","kind":"completed","agentThreadId":"child-thread","agentPath":"/root/history_boundaries"}}}"#
|
||||
)),
|
||||
vec![Event::Status {
|
||||
state: SessionStatus::Idle
|
||||
}]
|
||||
vec![
|
||||
Event::Status {
|
||||
state: SessionStatus::Idle
|
||||
},
|
||||
Event::BackgroundTasks { count: 0 }
|
||||
]
|
||||
);
|
||||
|
||||
let child = subagents.get("child-thread").expect("child");
|
||||
@@ -1289,6 +1309,50 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
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 mut translator = Translator::new(
|
||||
Arc::clone(&subagents),
|
||||
Some("parent-thread".to_string()),
|
||||
false,
|
||||
);
|
||||
|
||||
for (call, child, count) in [("spawn-a", "child-a", 1), ("spawn-b", "child-b", 2)] {
|
||||
let events = translator.translate(&json!({
|
||||
"method": "item/completed",
|
||||
"params": {
|
||||
"threadId": "parent-thread",
|
||||
"item": {
|
||||
"id": call,
|
||||
"type": "subAgentActivity",
|
||||
"kind": "started",
|
||||
"agentThreadId": child,
|
||||
"agentPath": format!("/root/{child}")
|
||||
}
|
||||
}
|
||||
}));
|
||||
assert_eq!(events.last(), Some(&Event::BackgroundTasks { count }));
|
||||
}
|
||||
for (child, count) in [("child-a", 1), ("child-b", 0)] {
|
||||
let events = translator.translate(&json!({
|
||||
"method": "item/completed",
|
||||
"params": {
|
||||
"threadId": "parent-thread",
|
||||
"item": {
|
||||
"id": format!("completed-{child}"),
|
||||
"type": "subAgentActivity",
|
||||
"kind": "completed",
|
||||
"agentThreadId": child,
|
||||
"agentPath": format!("/root/{child}")
|
||||
}
|
||||
}
|
||||
}));
|
||||
assert_eq!(events.last(), Some(&Event::BackgroundTasks { count }));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_collaboration_coordination_has_clean_parent_tool_cards() {
|
||||
let mut translator = Translator::default();
|
||||
@@ -1430,12 +1494,11 @@ mod tests {
|
||||
true,
|
||||
);
|
||||
|
||||
assert!(
|
||||
translator
|
||||
.translate(&line(
|
||||
r#"{"method":"item/completed","params":{"threadId":"parent-thread","turnId":"turn-1","item":{"id":"subagent-completed-1","type":"subAgentActivity","kind":"completed","agentThreadId":"child-thread","agentPath":"/root/child"}}}"#
|
||||
))
|
||||
.is_empty()
|
||||
assert_eq!(
|
||||
translator.translate(&line(
|
||||
r#"{"method":"item/completed","params":{"threadId":"parent-thread","turnId":"turn-1","item":{"id":"subagent-completed-1","type":"subAgentActivity","kind":"completed","agentThreadId":"child-thread","agentPath":"/root/child"}}}"#
|
||||
)),
|
||||
vec![Event::BackgroundTasks { count: 0 }]
|
||||
);
|
||||
assert_eq!(subagents.list(true)[0].status, SessionStatus::Exited);
|
||||
}
|
||||
|
||||
@@ -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]
|
||||
|
||||
Reference in new issue
Block a user