vtcode_bash_runner/
background.rs

1use anyhow::Result;
2use hashbrown::HashMap;
3use std::sync::Arc;
4use tokio::sync::{RwLock, oneshot};
5
6use crate::executor::{CommandExecutor, CommandInvocation, CommandOutput};
7
8#[derive(Debug, Clone)]
9pub struct BackgroundTaskHandle {
10    pub id: String,
11    pub command: String,
12    pub status: BackgroundTaskStatus,
13}
14
15#[derive(Debug, Clone, PartialEq)]
16pub enum BackgroundTaskStatus {
17    Pending,
18    Running,
19    Completed,
20    Failed,
21}
22
23#[derive(Debug)]
24pub struct BackgroundTask {
25    pub id: String,
26    pub invocation: CommandInvocation,
27    pub status: BackgroundTaskStatus,
28    pub result: Option<Result<CommandOutput, String>>,
29    pub cancel_tx: Option<oneshot::Sender<()>>,
30}
31
32pub struct BackgroundCommandManager<E: CommandExecutor> {
33    executor: Arc<E>,
34    tasks: Arc<RwLock<HashMap<String, BackgroundTask>>>,
35    next_id: Arc<RwLock<u64>>,
36}
37
38impl<E: CommandExecutor + 'static> BackgroundCommandManager<E> {
39    pub fn new(executor: E) -> Self {
40        Self {
41            executor: Arc::new(executor),
42            tasks: Arc::new(RwLock::new(HashMap::new())),
43            next_id: Arc::new(RwLock::new(1)),
44        }
45    }
46
47    pub async fn run_command(&self, invocation: CommandInvocation) -> Result<String> {
48        let task_id = self.generate_task_id().await;
49
50        let (cancel_tx, cancel_rx) = oneshot::channel();
51
52        let task = BackgroundTask {
53            id: task_id.clone(),
54            invocation: invocation.clone(),
55            status: BackgroundTaskStatus::Pending,
56            result: None,
57            cancel_tx: Some(cancel_tx),
58        };
59
60        {
61            let mut tasks = self.tasks.write().await;
62            tasks.insert(task_id.clone(), task);
63        }
64
65        // Update status to running
66        self.update_task_status(&task_id, BackgroundTaskStatus::Running).await;
67
68        // Spawn the background task
69        let executor = self.executor.clone();
70        let tasks = self.tasks.clone();
71        let id = task_id.clone();
72
73        tokio::spawn(async move {
74            let result = tokio::select! {
75                command_result = execute_command(executor.as_ref(), &invocation) => {
76                    command_result
77                }
78                _ = cancel_rx => {
79                    // Task was cancelled
80                    Err(anyhow::anyhow!("Command was cancelled"))
81                }
82            };
83
84            let mut tasks = tasks.write().await;
85            if let Some(task) = tasks.get_mut(&id) {
86                task.status = match result.is_ok() {
87                    true => BackgroundTaskStatus::Completed,
88                    false => BackgroundTaskStatus::Failed,
89                };
90                task.result = Some(result.map_err(|e| e.to_string()));
91                task.cancel_tx = None; // Clear the cancel sender
92            }
93        });
94
95        Ok(task_id)
96    }
97
98    pub async fn get_task(&self, task_id: &str) -> Option<BackgroundTaskHandle> {
99        let tasks = self.tasks.read().await;
100        tasks.get(task_id).map(|task| BackgroundTaskHandle {
101            id: task.id.clone(),
102            command: task.invocation.command.clone(),
103            status: task.status.clone(),
104        })
105    }
106
107    pub async fn get_task_output(&self, task_id: &str) -> Option<Result<CommandOutput, String>> {
108        let tasks = self.tasks.read().await;
109        tasks.get(task_id).and_then(|task| task.result.clone())
110    }
111
112    pub async fn list_tasks(&self) -> Vec<BackgroundTaskHandle> {
113        let tasks = self.tasks.read().await;
114        tasks
115            .values()
116            .map(|task| BackgroundTaskHandle {
117                id: task.id.clone(),
118                command: task.invocation.command.clone(),
119                status: task.status.clone(),
120            })
121            .collect()
122    }
123
124    pub async fn cancel_task(&self, task_id: &str) -> Result<()> {
125        let mut tasks = self.tasks.write().await;
126        if let Some(task) = tasks.get_mut(task_id)
127            && let Some(cancel_tx) = task.cancel_tx.take()
128            && cancel_tx.send(()).is_ok()
129        {
130            task.status = BackgroundTaskStatus::Failed;
131            return Ok(());
132        }
133        anyhow::bail!("Task not found or already completed: {task_id}");
134    }
135
136    async fn generate_task_id(&self) -> String {
137        let mut next_id = self.next_id.write().await;
138        let id = format!("bg-{next_id}");
139        *next_id += 1;
140        id
141    }
142
143    async fn update_task_status(&self, task_id: &str, status: BackgroundTaskStatus) {
144        let mut tasks = self.tasks.write().await;
145        if let Some(task) = tasks.get_mut(task_id) {
146            task.status = status;
147        }
148    }
149}
150
151async fn execute_command<E: CommandExecutor>(
152    executor: &E,
153    invocation: &CommandInvocation,
154) -> Result<CommandOutput> {
155    executor.execute(invocation)
156}
vtcode_bash_runner/background.rs

vtcode_bash_runner/
background.rs
