use everruns_core::session_task::TASK_KIND_MONITOR;
use everruns_core::{SessionId, SessionStore, SessionTask, SessionTaskRegistry, TokenUsage};
use std::collections::HashSet;
const MAX_TREE_SESSIONS: usize = 256;
const MAX_TREE_ROWS: usize = 2_048;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct BackgroundCounts {
pub(crate) running: usize,
pub(crate) scheduled: usize,
pub(crate) total: usize,
}
#[derive(Clone, Debug)]
pub(crate) struct TaskTreeRow {
pub(crate) task: SessionTask,
pub(crate) prefix: String,
pub(crate) branch_usage: Option<TokenUsage>,
parent: Option<usize>,
}
#[derive(Clone, Debug, Default)]
pub(crate) struct TaskTree {
pub(crate) rows: Vec<TaskTreeRow>,
pub(crate) errors: Vec<String>,
}
impl TaskTree {
pub(crate) fn counts(&self) -> Option<BackgroundCounts> {
counts(self.rows.iter().map(|row| &row.task))
}
pub(crate) fn selected(&self, selected: usize) -> Option<&SessionTask> {
self.rows.get(selected).map(|row| &row.task)
}
}
enum Visit {
Session {
session_id: SessionId,
parent: Option<usize>,
ancestors_last: Vec<bool>,
},
Task {
task: Box<SessionTask>,
parent: Option<usize>,
ancestors_last: Vec<bool>,
is_last: bool,
},
}
pub(crate) async fn load_task_tree(
root_session_id: SessionId,
registry: &dyn SessionTaskRegistry,
sessions: &dyn SessionStore,
) -> TaskTree {
let mut tree = TaskTree::default();
let mut stack = vec![Visit::Session {
session_id: root_session_id,
parent: None,
ancestors_last: Vec::new(),
}];
let mut visited = HashSet::new();
let mut sessions_truncated = false;
let mut rows_truncated = false;
while let Some(visit) = stack.pop() {
match visit {
Visit::Session {
session_id,
parent,
ancestors_last,
} => {
if visited.len() >= MAX_TREE_SESSIONS {
if !sessions_truncated {
tree.errors.push(format!(
"task tree stopped after {MAX_TREE_SESSIONS} sessions"
));
sessions_truncated = true;
}
continue;
}
if !visited.insert(session_id) {
tree.errors
.push(format!("task tree cycle at session {session_id}"));
continue;
}
match registry.list(session_id, None).await {
Ok(mut tasks) => {
tasks.sort_by(|a, b| {
a.created_at
.cmp(&b.created_at)
.then_with(|| a.id.cmp(&b.id))
});
let count = tasks.len();
for (index, task) in tasks.into_iter().enumerate().rev() {
stack.push(Visit::Task {
task: Box::new(task),
parent,
ancestors_last: ancestors_last.clone(),
is_last: index + 1 == count,
});
}
}
Err(error) => tree.errors.push(format!("session {session_id}: {error}")),
}
}
Visit::Task {
task,
parent,
ancestors_last,
is_last,
} => {
if tree.rows.len() >= MAX_TREE_ROWS {
if !rows_truncated {
tree.errors
.push(format!("task tree stopped after {MAX_TREE_ROWS} tasks"));
rows_truncated = true;
}
continue;
}
let task = *task;
let mut prefix = String::new();
for ancestor_last in &ancestors_last {
prefix.push_str(if *ancestor_last { " " } else { "│ " });
}
prefix.push_str(if is_last { "└─ " } else { "├─ " });
let (child_session_id, branch_usage) =
if let Some(child_id) = task.links.child_session_id {
match sessions.get_session(child_id).await {
Ok(Some(session))
if session.parent_session_id == Some(task.session_id)
|| session.forked_from_session_id == Some(task.session_id) =>
{
(Some(child_id), session.usage)
}
Ok(Some(_)) => {
tree.errors.push(format!(
"linked session {child_id} is not a child of {}",
task.session_id
));
(None, None)
}
Ok(None) => {
tree.errors
.push(format!("child session {child_id} not found"));
(None, None)
}
Err(error) => {
tree.errors
.push(format!("child session {child_id}: {error}"));
(None, None)
}
}
} else {
(None, None)
};
let row_index = tree.rows.len();
tree.rows.push(TaskTreeRow {
task,
prefix,
branch_usage,
parent,
});
if let Some(child_id) = child_session_id {
let mut child_ancestors = ancestors_last;
child_ancestors.push(is_last);
stack.push(Visit::Session {
session_id: child_id,
parent: Some(row_index),
ancestors_last: child_ancestors,
});
}
}
}
}
for index in (0..tree.rows.len()).rev() {
let Some(parent) = tree.rows[index].parent else {
continue;
};
let Some(usage) = tree.rows[index].branch_usage.clone() else {
continue;
};
match tree.rows[parent].branch_usage.as_mut() {
Some(parent_usage) => parent_usage.add(&usage),
None => tree.rows[parent].branch_usage = Some(usage),
}
}
tree
}
fn counts<'a>(tasks: impl Iterator<Item = &'a SessionTask>) -> Option<BackgroundCounts> {
let mut counts = BackgroundCounts {
running: 0,
scheduled: 0,
total: 0,
};
for task in tasks {
counts.total += 1;
if !task.state.is_terminal() {
if task.kind == TASK_KIND_MONITOR {
counts.scheduled += 1;
} else {
counts.running += 1;
}
}
}
(counts.total > 0).then_some(counts)
}
pub(crate) fn render_task_tree(tree: &TaskTree, selected: Option<usize>) -> String {
if tree.rows.is_empty() {
if tree.errors.is_empty() {
return "No background tasks in this session.".to_string();
}
return format!(
"No background tasks in this session.\n\nSession task registry: {}",
tree.errors.join("; ")
);
}
let counts = tree.counts().expect("non-empty task tree has counts");
let mut out = format!(
"{} task(s): {} running, {} scheduled\n",
counts.total, counts.running, counts.scheduled
);
for (index, row) in tree.rows.iter().enumerate() {
let cursor = if selected == Some(index) { "> " } else { " " };
let task = &row.task;
out.push_str(cursor);
out.push_str(&row.prefix);
out.push_str(&format!(
"[{}] {} {}: {}",
task.id, task.kind, task.state, task.display_name
));
if let Some(usage) = &row.branch_usage {
out.push_str(&format!(
" · {} tok",
format_tokens(usage.total_tokens() as u64)
));
if let Some(cost) = usage.effective_cost_usd() {
out.push_str(&format!(" · ${cost:.4}"));
}
}
if task.cancel_requested_at.is_some() && !task.state.is_terminal() {
out.push_str(" · canceling");
}
if let Some(detail) = task_detail(task) {
out.push_str(" -- ");
out.push_str(detail);
}
if let Some(path) = task.result_path.as_deref() {
out.push_str(&format!(" (result: {path})"));
}
out.push('\n');
}
if !tree.errors.is_empty() {
out.push_str("\nWarnings: ");
out.push_str(&tree.errors.join("; "));
out.push('\n');
}
out.push_str("\nInspect tasks with list_tasks/get_task; cancel with cancel_task.");
out
}
fn task_detail(task: &SessionTask) -> Option<&str> {
if task.state.is_terminal() {
task.summary
.as_deref()
.or_else(|| task.error.as_ref().map(|error| error.message.as_str()))
} else {
task.state_detail
.as_deref()
.or(task.summary.as_deref())
.or_else(|| task.error.as_ref().map(|error| error.message.as_str()))
}
}
fn format_tokens(tokens: u64) -> String {
if tokens >= 1_000_000 {
format!("{:.1}m", tokens as f64 / 1_000_000.0)
} else if tokens >= 1_000 {
format!("{:.1}k", tokens as f64 / 1_000.0)
} else {
tokens.to_string()
}
}
#[cfg(test)]
mod tests {
use super::*;
use async_trait::async_trait;
use chrono::Utc;
use everruns_core::session_task::{
SessionTaskState, TASK_KIND_BACKGROUND_TOOL, TASK_KIND_MONITOR, TASK_KIND_SUBAGENT,
TaskLinks, new_session_task,
};
use everruns_core::{CreateSessionTask, HarnessId, Session, SessionId};
use everruns_local::{LocalSessionTaskRegistry, SqliteDb};
use everruns_runtime::SessionBuilder;
use serde_json::json;
use std::collections::HashMap;
fn task(id: &str, state: SessionTaskState, name: &str) -> SessionTask {
task_with_kind(id, TASK_KIND_BACKGROUND_TOOL, state, name)
}
fn task_with_kind(id: &str, kind: &str, state: SessionTaskState, name: &str) -> SessionTask {
let mut task = new_session_task(
CreateSessionTask {
session_id: SessionId::from_seed(42),
id: Some(id.to_string()),
kind: kind.to_string(),
display_name: name.to_string(),
spec: json!({}),
state,
links: Default::default(),
wake_policy: Default::default(),
},
Utc::now(),
);
task.summary = Some("done".to_string());
task
}
fn tree(tasks: Vec<SessionTask>) -> TaskTree {
TaskTree {
rows: tasks
.into_iter()
.enumerate()
.map(|(index, task)| TaskTreeRow {
task,
prefix: if index == 0 { "├─ " } else { "└─ " }.into(),
branch_usage: None,
parent: None,
})
.collect(),
errors: Vec::new(),
}
}
#[test]
fn counts_reflect_active_and_total_session_tasks() {
let tree = tree(vec![
task("task_a", SessionTaskState::Running, "run"),
task_with_kind(
"task_monitor",
TASK_KIND_MONITOR,
SessionTaskState::Running,
"scheduled review",
),
task("task_b", SessionTaskState::Succeeded, "done"),
]);
assert_eq!(
tree.counts(),
Some(BackgroundCounts {
running: 1,
scheduled: 1,
total: 3,
})
);
assert_eq!(TaskTree::default().counts(), None);
}
#[test]
fn render_lists_tree_rows_and_selection() {
let tree = tree(vec![task(
"task_a",
SessionTaskState::Succeeded,
"write marker",
)]);
let rendered = render_task_tree(&tree, Some(0));
assert!(rendered.contains("1 task(s): 0 running, 0 scheduled"));
assert!(rendered.contains("> ├─ [task_a] background_tool succeeded: write marker"));
}
#[test]
fn render_formats_branch_usage() {
let mut tree = tree(vec![task(
"task_agent",
SessionTaskState::Running,
"research",
)]);
tree.rows[0].branch_usage =
Some(TokenUsage::new(12_000, 345).with_effective_cost(Some(0.12345)));
let rendered = render_task_tree(&tree, None);
assert!(rendered.contains("12.3k tok · $0.1235"));
}
#[test]
fn render_empty_reports_no_tasks() {
assert_eq!(
render_task_tree(&TaskTree::default(), None),
"No background tasks in this session."
);
}
struct Sessions(HashMap<SessionId, Session>);
#[async_trait]
impl SessionStore for Sessions {
async fn get_session(
&self,
session_id: SessionId,
) -> everruns_core::Result<Option<Session>> {
Ok(self.0.get(&session_id).cloned())
}
}
#[tokio::test]
async fn load_nests_child_tasks_and_rolls_usage_to_branch_root() {
let root_id = SessionId::from_seed(700);
let child_id = SessionId::from_seed(701);
let grandchild_id = SessionId::from_seed(702);
let registry = LocalSessionTaskRegistry::new(
SqliteDb::open_in_memory().expect("in-memory task database"),
)
.expect("task registry");
registry
.create(CreateSessionTask {
session_id: root_id,
id: Some("task_parent".into()),
kind: TASK_KIND_SUBAGENT.into(),
display_name: "research".into(),
spec: json!({}),
state: SessionTaskState::Running,
links: TaskLinks {
child_session_id: Some(child_id),
..Default::default()
},
wake_policy: Default::default(),
})
.await
.unwrap();
registry
.create(CreateSessionTask {
session_id: child_id,
id: Some("task_child".into()),
kind: TASK_KIND_SUBAGENT.into(),
display_name: "verify".into(),
spec: json!({}),
state: SessionTaskState::Running,
links: TaskLinks {
child_session_id: Some(grandchild_id),
..Default::default()
},
wake_policy: Default::default(),
})
.await
.unwrap();
let mut child = SessionBuilder::new(HarnessId::from_seed(1))
.id(child_id)
.build();
child.parent_session_id = Some(root_id);
child.usage = Some(TokenUsage::new(100, 20).with_effective_cost(Some(0.01)));
let mut grandchild = SessionBuilder::new(HarnessId::from_seed(1))
.id(grandchild_id)
.build();
grandchild.parent_session_id = Some(child_id);
grandchild.usage = Some(TokenUsage::new(50, 5).with_effective_cost(Some(0.02)));
let sessions = Sessions(HashMap::from([
(child_id, child),
(grandchild_id, grandchild),
]));
let tree = load_task_tree(root_id, ®istry, &sessions).await;
assert_eq!(tree.rows.len(), 2);
assert_eq!(tree.rows[0].task.id, "task_parent");
assert_eq!(tree.rows[1].task.id, "task_child");
assert!(tree.rows[1].prefix.starts_with(" "));
let root_usage = tree.rows[0].branch_usage.as_ref().unwrap();
assert_eq!(root_usage.total_tokens(), 175);
assert_eq!(root_usage.effective_cost_usd(), Some(0.03));
}
#[tokio::test]
async fn load_rejects_linked_sessions_outside_the_task_lineage() {
let root_id = SessionId::from_seed(710);
let unrelated_id = SessionId::from_seed(711);
let registry = LocalSessionTaskRegistry::new(
SqliteDb::open_in_memory().expect("in-memory task database"),
)
.expect("task registry");
registry
.create(CreateSessionTask {
session_id: root_id,
id: Some("task_parent".into()),
kind: TASK_KIND_SUBAGENT.into(),
display_name: "research".into(),
spec: json!({}),
state: SessionTaskState::Running,
links: TaskLinks {
child_session_id: Some(unrelated_id),
..Default::default()
},
wake_policy: Default::default(),
})
.await
.unwrap();
let mut unrelated = SessionBuilder::new(HarnessId::from_seed(1))
.id(unrelated_id)
.build();
unrelated.usage = Some(TokenUsage::new(1_000, 0));
let sessions = Sessions(HashMap::from([(unrelated_id, unrelated)]));
let tree = load_task_tree(root_id, ®istry, &sessions).await;
assert_eq!(tree.rows.len(), 1);
assert!(tree.rows[0].branch_usage.is_none());
assert!(
tree.errors
.iter()
.any(|error| error.contains("is not a child"))
);
}
}