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    id: String,
11    command: String,
12    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    id: String,
26    invocation: CommandInvocation,
27    status: BackgroundTaskStatus,
28    result: Option<Result<CommandOutput, String>>,
29    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        // Documented detached: bounded by the command execution itself and
70        // terminated early via `cancel_rx`; the outcome is observable through
71        // the task map (`status`/`result`), which readers poll via
72        // `get_task`/`get_task_output`.
73        let executor = self.executor.clone();
74        let tasks = self.tasks.clone();
75        let id = task_id.clone();
76
77        tokio::spawn(async move {
78            let result = tokio::select! {
79                command_result = execute_command(executor.as_ref(), &invocation) => {
80                    command_result
81                }
82                _ = cancel_rx => {
83                    // Task was cancelled
84                    Err(anyhow::anyhow!("Command was cancelled"))
85                }
86            };
87
88            let mut tasks = tasks.write().await;
89            if let Some(task) = tasks.get_mut(&id) {
90                task.status = match result.is_ok() {
91                    true => BackgroundTaskStatus::Completed,
92                    false => BackgroundTaskStatus::Failed,
93                };
94                task.result = Some(result.map_err(|e| e.to_string()));
95                task.cancel_tx = None; // Clear the cancel sender
96            }
97        });
98
99        Ok(task_id)
100    }
101
102    pub async fn get_task(&self, task_id: &str) -> Option<BackgroundTaskHandle> {
103        let tasks = self.tasks.read().await;
104        tasks.get(task_id).map(|task| BackgroundTaskHandle {
105            id: task.id.clone(),
106            command: task.invocation.command.clone(),
107            status: task.status.clone(),
108        })
109    }
110
111    pub async fn get_task_output(&self, task_id: &str) -> Option<Result<CommandOutput, String>> {
112        let tasks = self.tasks.read().await;
113        tasks.get(task_id).and_then(|task| task.result.clone())
114    }
115
116    pub async fn list_tasks(&self) -> Vec<BackgroundTaskHandle> {
117        let tasks = self.tasks.read().await;
118        tasks
119            .values()
120            .map(|task| BackgroundTaskHandle {
121                id: task.id.clone(),
122                command: task.invocation.command.clone(),
123                status: task.status.clone(),
124            })
125            .collect()
126    }
127
128    pub async fn cancel_task(&self, task_id: &str) -> Result<()> {
129        let mut tasks = self.tasks.write().await;
130        if let Some(task) = tasks.get_mut(task_id)
131            && let Some(cancel_tx) = task.cancel_tx.take()
132            && cancel_tx.send(()).is_ok()
133        {
134            task.status = BackgroundTaskStatus::Failed;
135            return Ok(());
136        }
137        anyhow::bail!("Task not found or already completed: {task_id}");
138    }
139
140    async fn generate_task_id(&self) -> String {
141        let mut next_id = self.next_id.write().await;
142        let id = format!("bg-{next_id}");
143        *next_id += 1;
144        id
145    }
146
147    async fn update_task_status(&self, task_id: &str, status: BackgroundTaskStatus) {
148        let mut tasks = self.tasks.write().await;
149        if let Some(task) = tasks.get_mut(task_id) {
150            task.status = status;
151        }
152    }
153}
154
155async fn execute_command<E: CommandExecutor>(executor: &E, invocation: &CommandInvocation) -> Result<CommandOutput> {
156    executor.execute(invocation)
157}
vtcode_bash_runner/background.rs

vtcode_bash_runner/
background.rs
