ubiquity_core/
command.rs

1//! Command execution system with local, WASM, and cloud parity
2
3use async_trait::async_trait;
4use futures::Stream;
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7use std::pin::Pin;
8use std::sync::Arc;
9use std::time::Duration;
10use tokio::sync::{mpsc, oneshot, RwLock};
11use uuid::Uuid;
12
13use crate::error::UbiquityError;
14
15/// Command execution event types
16#[derive(Debug, Clone, Serialize, Deserialize)]
17#[serde(tag = "type")]
18pub enum CommandEvent {
19    /// Command has started execution
20    Started {
21        command_id: Uuid,
22        command: String,
23        args: Vec<String>,
24        timestamp: chrono::DateTime<chrono::Utc>,
25    },
26    /// Standard output data
27    Stdout {
28        command_id: Uuid,
29        data: String,
30        timestamp: chrono::DateTime<chrono::Utc>,
31    },
32    /// Standard error data
33    Stderr {
34        command_id: Uuid,
35        data: String,
36        timestamp: chrono::DateTime<chrono::Utc>,
37    },
38    /// Command progress update
39    Progress {
40        command_id: Uuid,
41        percentage: f32,
42        message: String,
43        timestamp: chrono::DateTime<chrono::Utc>,
44    },
45    /// Command completed successfully
46    Completed {
47        command_id: Uuid,
48        exit_code: i32,
49        duration_ms: u64,
50        timestamp: chrono::DateTime<chrono::Utc>,
51    },
52    /// Command failed with error
53    Failed {
54        command_id: Uuid,
55        error: String,
56        duration_ms: u64,
57        timestamp: chrono::DateTime<chrono::Utc>,
58    },
59    /// Command was cancelled
60    Cancelled {
61        command_id: Uuid,
62        duration_ms: u64,
63        timestamp: chrono::DateTime<chrono::Utc>,
64    },
65}
66
67/// Command execution request
68#[derive(Debug, Clone, Serialize, Deserialize)]
69pub struct CommandRequest {
70    pub id: Uuid,
71    pub command: String,
72    pub args: Vec<String>,
73    pub env: HashMap<String, String>,
74    pub working_dir: Option<String>,
75    pub timeout: Option<Duration>,
76    pub stdin: Option<String>,
77}
78
79impl CommandRequest {
80    pub fn new(command: impl Into<String>) -> Self {
81        Self {
82            id: Uuid::new_v4(),
83            command: command.into(),
84            args: Vec::new(),
85            env: HashMap::new(),
86            working_dir: None,
87            timeout: None,
88            stdin: None,
89        }
90    }
91
92    pub fn with_args(mut self, args: Vec<String>) -> Self {
93        self.args = args;
94        self
95    }
96
97    pub fn with_env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
98        self.env.insert(key.into(), value.into());
99        self
100    }
101
102    pub fn with_working_dir(mut self, dir: impl Into<String>) -> Self {
103        self.working_dir = Some(dir.into());
104        self
105    }
106
107    pub fn with_timeout(mut self, timeout: Duration) -> Self {
108        self.timeout = Some(timeout);
109        self
110    }
111
112    pub fn with_stdin(mut self, stdin: impl Into<String>) -> Self {
113        self.stdin = Some(stdin.into());
114        self
115    }
116}
117
118/// Command execution result
119#[derive(Debug, Clone, Serialize, Deserialize)]
120pub struct CommandResult {
121    pub id: Uuid,
122    pub exit_code: Option<i32>,
123    pub stdout: String,
124    pub stderr: String,
125    pub duration_ms: u64,
126    pub cancelled: bool,
127}
128
129/// Command executor trait for different implementations
130#[async_trait]
131pub trait CommandExecutor: Send + Sync {
132    /// Execute a command and return a stream of events
133    async fn execute(
134        &self,
135        request: CommandRequest,
136    ) -> Result<Pin<Box<dyn Stream<Item = CommandEvent> + Send>>, UbiquityError>;
137
138    /// Cancel a running command
139    async fn cancel(&self, command_id: Uuid) -> Result<(), UbiquityError>;
140
141    /// Get the status of a command
142    async fn status(&self, command_id: Uuid) -> Result<Option<CommandResult>, UbiquityError>;
143}
144
145/// Command execution context for tracking active commands
146pub struct CommandContext {
147    active_commands: Arc<RwLock<HashMap<Uuid, CommandHandle>>>,
148}
149
150impl CommandContext {
151    pub fn new() -> Self {
152        Self {
153            active_commands: Arc::new(RwLock::new(HashMap::new())),
154        }
155    }
156
157    pub async fn register(&self, id: Uuid, handle: CommandHandle) {
158        self.active_commands.write().await.insert(id, handle);
159    }
160
161    pub async fn unregister(&self, id: &Uuid) {
162        self.active_commands.write().await.remove(id);
163    }
164
165    pub async fn get(&self, id: &Uuid) -> Option<CommandHandle> {
166        self.active_commands.read().await.get(id).cloned()
167    }
168
169    pub async fn cancel(&self, id: &Uuid) -> Result<(), UbiquityError> {
170        if let Some(handle) = self.get(id).await {
171            handle.cancel().await
172        } else {
173            Err(UbiquityError::NotFound(format!("Command {} not found", id)))
174        }
175    }
176
177    pub async fn cancel_all(&self) -> Vec<(Uuid, Result<(), UbiquityError>)> {
178        let commands: Vec<(Uuid, CommandHandle)> = {
179            let guard = self.active_commands.read().await;
180            guard.iter().map(|(id, h)| (*id, h.clone())).collect()
181        };
182
183        let mut results = Vec::new();
184        for (id, handle) in commands {
185            let result = handle.cancel().await;
186            results.push((id, result));
187        }
188        results
189    }
190}
191
192/// Handle to a running command
193#[derive(Clone)]
194pub struct CommandHandle {
195    pub id: Uuid,
196    cancel_tx: mpsc::Sender<()>,
197    status_tx: mpsc::Sender<oneshot::Sender<CommandResult>>,
198}
199
200impl CommandHandle {
201    pub fn new(
202        id: Uuid,
203        cancel_tx: mpsc::Sender<()>,
204        status_tx: mpsc::Sender<oneshot::Sender<CommandResult>>,
205    ) -> Self {
206        Self {
207            id,
208            cancel_tx,
209            status_tx,
210        }
211    }
212
213    pub async fn cancel(&self) -> Result<(), UbiquityError> {
214        self.cancel_tx
215            .send(())
216            .await
217            .map_err(|_| UbiquityError::Internal("Failed to send cancel signal".to_string()))
218    }
219
220    pub async fn status(&self) -> Result<CommandResult, UbiquityError> {
221        let (tx, rx) = oneshot::channel();
222        self.status_tx
223            .send(tx)
224            .await
225            .map_err(|_| UbiquityError::Internal("Failed to request status".to_string()))?;
226        
227        rx.await
228            .map_err(|_| UbiquityError::Internal("Failed to receive status".to_string()))
229    }
230}
231
232/// Progress tracker for commands that support progress reporting
233pub struct ProgressTracker {
234    command_id: Uuid,
235    tx: mpsc::Sender<CommandEvent>,
236}
237
238impl ProgressTracker {
239    pub fn new(command_id: Uuid, tx: mpsc::Sender<CommandEvent>) -> Self {
240        Self { command_id, tx }
241    }
242
243    pub async fn update(&self, percentage: f32, message: impl Into<String>) -> Result<(), UbiquityError> {
244        self.tx
245            .send(CommandEvent::Progress {
246                command_id: self.command_id,
247                percentage: percentage.clamp(0.0, 100.0),
248                message: message.into(),
249                timestamp: chrono::Utc::now(),
250            })
251            .await
252            .map_err(|_| UbiquityError::Internal("Failed to send progress update".to_string()))
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259
260    #[test]
261    fn test_command_request_builder() {
262        let request = CommandRequest::new("echo")
263            .with_args(vec!["hello".to_string(), "world".to_string()])
264            .with_env("PATH", "/usr/bin")
265            .with_working_dir("/tmp")
266            .with_timeout(Duration::from_secs(30))
267            .with_stdin("test input");
268
269        assert_eq!(request.command, "echo");
270        assert_eq!(request.args, vec!["hello", "world"]);
271        assert_eq!(request.env.get("PATH"), Some(&"/usr/bin".to_string()));
272        assert_eq!(request.working_dir, Some("/tmp".to_string()));
273        assert_eq!(request.timeout, Some(Duration::from_secs(30)));
274        assert_eq!(request.stdin, Some("test input".to_string()));
275    }
276
277    #[tokio::test]
278    async fn test_command_context() {
279        let context = CommandContext::new();
280        let id = Uuid::new_v4();
281        
282        let (cancel_tx, _cancel_rx) = mpsc::channel(1);
283        let (status_tx, _status_rx) = mpsc::channel(1);
284        let handle = CommandHandle::new(id, cancel_tx, status_tx);
285        
286        context.register(id, handle.clone()).await;
287        
288        let retrieved = context.get(&id).await;
289        assert!(retrieved.is_some());
290        assert_eq!(retrieved.unwrap().id, id);
291        
292        context.unregister(&id).await;
293        assert!(context.get(&id).await.is_none());
294    }
295
296    #[tokio::test]
297    async fn test_progress_tracker() {
298        let id = Uuid::new_v4();
299        let (tx, mut rx) = mpsc::channel(10);
300        let tracker = ProgressTracker::new(id, tx);
301        
302        tracker.update(50.0, "Half way there").await.unwrap();
303        
304        let event = rx.recv().await.unwrap();
305        match event {
306            CommandEvent::Progress { command_id, percentage, message, .. } => {
307                assert_eq!(command_id, id);
308                assert_eq!(percentage, 50.0);
309                assert_eq!(message, "Half way there");
310            }
311            _ => panic!("Expected Progress event"),
312        }
313    }
314}
ubiquity_core/command.rs

ubiquity_core/
command.rs
