use std::collections::HashMap;
use std::time::Instant;
use opendev_runtime::InterruptToken;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BackgroundAgentState {
Running,
Completed,
Failed,
Killed,
}
impl std::fmt::Display for BackgroundAgentState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Running => write!(f, "running"),
Self::Completed => write!(f, "completed"),
Self::Failed => write!(f, "failed"),
Self::Killed => write!(f, "killed"),
}
}
}
#[derive(Debug)]
pub struct BackgroundAgentTask {
pub task_id: String,
pub query: String,
pub session_id: String,
pub started_at: Instant,
pub state: BackgroundAgentState,
pub interrupt_token: InterruptToken,
pub result_summary: Option<String>,
pub tool_call_count: usize,
pub cost_usd: f64,
pub current_tool: Option<String>,
pub activity_log: Vec<String>,
pub pending_spawn_count: usize,
pub hidden: bool,
}
impl BackgroundAgentTask {
pub fn runtime_seconds(&self) -> f64 {
self.started_at.elapsed().as_secs_f64()
}
pub fn is_running(&self) -> bool {
self.state == BackgroundAgentState::Running
}
}
#[derive(Debug)]
pub struct BackgroundAgentManager {
tasks: HashMap<String, BackgroundAgentTask>,
pub max_concurrent: usize,
}
impl BackgroundAgentManager {
pub fn new() -> Self {
Self {
tasks: HashMap::new(),
max_concurrent: 3,
}
}
pub fn add_task(
&mut self,
task_id: String,
query: String,
session_id: String,
interrupt_token: InterruptToken,
) {
self.tasks.insert(
task_id.clone(),
BackgroundAgentTask {
task_id,
query,
session_id,
started_at: Instant::now(),
state: BackgroundAgentState::Running,
interrupt_token,
result_summary: None,
tool_call_count: 0,
cost_usd: 0.0,
current_tool: None,
activity_log: Vec::new(),
pending_spawn_count: 0,
hidden: false,
},
);
}
pub fn mark_completed(
&mut self,
task_id: &str,
success: bool,
result_summary: String,
tool_call_count: usize,
cost_usd: f64,
) {
if let Some(task) = self.tasks.get_mut(task_id) {
if task.state != BackgroundAgentState::Killed {
task.state = if success {
BackgroundAgentState::Completed
} else {
BackgroundAgentState::Failed
};
}
task.result_summary = Some(result_summary);
task.tool_call_count = tool_call_count;
task.cost_usd = cost_usd;
}
}
pub fn update_progress(&mut self, task_id: &str, tool_name: String, tool_count: usize) {
if let Some(task) = self.tasks.get_mut(task_id) {
task.current_tool = Some(tool_name);
task.tool_call_count = tool_count;
}
}
pub fn increment_pending_spawn(&mut self, task_id: &str) {
if let Some(task) = self.tasks.get_mut(task_id) {
task.pending_spawn_count += 1;
}
}
pub fn decrement_pending_spawn(&mut self, task_id: &str) {
if let Some(task) = self.tasks.get_mut(task_id) {
task.pending_spawn_count = task.pending_spawn_count.saturating_sub(1);
}
}
pub fn push_activity(&mut self, task_id: &str, line: String) {
if let Some(task) = self.tasks.get_mut(task_id) {
if line.starts_with('\u{27e1}')
&& task
.activity_log
.last()
.is_some_and(|l| l.starts_with('\u{27e1}'))
{
*task.activity_log.last_mut().unwrap() = line;
return;
}
if task.activity_log.len() >= 50 {
task.activity_log.remove(0);
}
task.activity_log.push(line);
}
}
pub fn hide_task(&mut self, task_id: &str) {
if let Some(task) = self.tasks.get_mut(task_id) {
task.hidden = true;
}
}
pub fn kill_task(&mut self, task_id: &str) -> bool {
if let Some(task) = self.tasks.get_mut(task_id)
&& task.is_running()
{
task.interrupt_token.request();
task.state = BackgroundAgentState::Killed;
return true;
}
false
}
pub fn get_task(&self, task_id: &str) -> Option<&BackgroundAgentTask> {
self.tasks.get(task_id)
}
pub fn all_tasks(&self) -> Vec<&BackgroundAgentTask> {
let mut tasks: Vec<&BackgroundAgentTask> = self.tasks.values().collect();
tasks.sort_by(|a, b| b.started_at.cmp(&a.started_at));
tasks
}
pub fn running_count(&self) -> usize {
self.tasks.values().filter(|t| t.is_running()).count()
}
pub fn can_accept(&self) -> bool {
self.running_count() < self.max_concurrent
}
pub fn len(&self) -> usize {
self.tasks.len()
}
pub fn is_empty(&self) -> bool {
self.tasks.is_empty()
}
pub fn cleanup_old(&mut self, max_age_secs: f64) {
self.tasks
.retain(|_, t| t.is_running() || t.runtime_seconds() < max_age_secs);
}
}
impl Default for BackgroundAgentManager {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
#[path = "background_agents_tests.rs"]
mod tests;