Skip to main content

exeora_cli/tools/
processes.rs

1use super::path::resolve_in_project;
2use crate::{
3    error::{ErrorCode, ExeoraError},
4    protocol::{
5        DEFAULT_COMMAND_TIMEOUT_MS, MAX_COMMAND_OUTPUT_BYTES, MAX_PROCESS_BUFFER_BYTES,
6        MAX_PROCESS_CHUNK_BYTES, MAX_PROCESSES_PER_PROJECT,
7    },
8};
9#[cfg(windows)]
10use process_wrap::tokio::JobObject;
11#[cfg(unix)]
12use process_wrap::tokio::ProcessGroup;
13use process_wrap::tokio::{ChildWrapper, CommandWrap, KillOnDrop};
14use serde::Deserialize;
15use serde_json::{Value, json};
16use std::{
17    collections::{HashMap, VecDeque},
18    path::{Path, PathBuf},
19    process::Stdio,
20    sync::Arc,
21    time::Duration,
22};
23use tokio::{
24    io::{AsyncRead, AsyncReadExt, AsyncWriteExt},
25    sync::Mutex,
26};
27use tokio_util::sync::CancellationToken;
28use uuid::Uuid;
29
30type SharedChild = Arc<Mutex<Box<dyn ChildWrapper>>>;
31
32struct Running {
33    root: PathBuf,
34    child: SharedChild,
35    stdin: Arc<Mutex<Option<tokio::process::ChildStdin>>>,
36    ring: Arc<Mutex<Ring>>,
37    exit_code: Option<i32>,
38    running: bool,
39}
40
41#[derive(Default)]
42struct Ring {
43    chunks: VecDeque<String>,
44    units: usize,
45    dropped: usize,
46}
47
48pub struct ProcessRegistry {
49    entries: Mutex<HashMap<String, Running>>,
50}
51
52impl Default for ProcessRegistry {
53    fn default() -> Self {
54        Self::new()
55    }
56}
57
58impl ProcessRegistry {
59    pub fn new() -> Self {
60        Self {
61            entries: Mutex::new(HashMap::new()),
62        }
63    }
64
65    pub async fn run_command(
66        &self,
67        root: &Path,
68        value: Value,
69        cancel: CancellationToken,
70    ) -> Result<Value, ExeoraError> {
71        let args: RunArgs = parse(value)?;
72        let (real_root, cwd) = resolve_in_project(root, args.cwd.as_deref().unwrap_or("."))?;
73        let timeout_ms = args.timeout_ms.unwrap_or(DEFAULT_COMMAND_TIMEOUT_MS);
74        let mut child = spawn_wrapped(&args.command, &real_root.join(cwd), false)?;
75        let stdout = child.stdout().take();
76        let stderr = child.stderr().take();
77        let stdout_task = tokio::spawn(capture(stdout, MAX_COMMAND_OUTPUT_BYTES));
78        let stderr_task = tokio::spawn(capture(stderr, MAX_COMMAND_OUTPUT_BYTES));
79
80        let mut timed_out = false;
81        let mut cancelled = false;
82        let status = {
83            let wait = child.wait();
84            tokio::pin!(wait);
85            tokio::select! {
86                status = &mut wait => Some(status.map_err(|error| ExeoraError::tool(error.to_string()))?),
87                _ = tokio::time::sleep(Duration::from_millis(timeout_ms)) => { timed_out = true; None },
88                _ = cancel.cancelled() => { cancelled = true; None },
89            }
90        };
91        if status.is_none() {
92            let _ = kill_child(child.as_mut()).await;
93        }
94        let (stdout, stdout_cut) = stdout_task.await.map_err(join_error)??;
95        let (stderr, stderr_cut) = stderr_task.await.map_err(join_error)??;
96        if cancelled {
97            return Err(ExeoraError::new(
98                ErrorCode::Cancelled,
99                "The call was cancelled while the command was running.",
100            ));
101        }
102        Ok(json!({
103            "command": args.command,
104            "exitCode": status.and_then(|status| status.code()),
105            "stdout": stdout,
106            "stderr": stderr,
107            "truncated": stdout_cut || stderr_cut,
108            "timedOut": timed_out,
109        }))
110    }
111
112    pub async fn start_command(&self, root: &Path, value: Value) -> Result<Value, ExeoraError> {
113        let args: StartArgs = parse(value)?;
114        let (real_root, cwd) = resolve_in_project(root, args.cwd.as_deref().unwrap_or("."))?;
115        let root_key = real_root.clone();
116        let mut entries = self.entries.lock().await;
117        if entries
118            .values()
119            .filter(|entry| entry.root == root_key && entry.running)
120            .count()
121            >= MAX_PROCESSES_PER_PROJECT
122        {
123            return Err(ExeoraError::tool(format!(
124                "This project already has {MAX_PROCESSES_PER_PROJECT} processes running. Stop one with kill_command before starting another."
125            )));
126        }
127        let mut child = spawn_wrapped(&args.command, &real_root.join(cwd), true)?;
128        let pid = child.id();
129        let stdin = Arc::new(Mutex::new(child.stdin().take()));
130        let stdout = child.stdout().take();
131        let stderr = child.stderr().take();
132        let ring = Arc::new(Mutex::new(Ring::default()));
133        spawn_reader(stdout, ring.clone());
134        spawn_reader(stderr, ring.clone());
135        let id = format!("proc_{}", Uuid::new_v4().simple());
136        entries.insert(
137            id.clone(),
138            Running {
139                root: real_root,
140                child: Arc::new(Mutex::new(child)),
141                stdin,
142                ring,
143                exit_code: None,
144                running: true,
145            },
146        );
147        Ok(json!({ "processId": id, "command": args.command, "pid": pid }))
148    }
149
150    pub async fn get_output(&self, root: &Path, value: Value) -> Result<Value, ExeoraError> {
151        let args: OutputArgs = parse(value)?;
152        let mut entries = self.entries.lock().await;
153        let entry = find_entry(&mut entries, root, &args.process_id)?;
154        refresh(entry).await;
155        let ring = entry.ring.lock().await;
156        let total = ring.dropped + ring.units;
157        let from = args.cursor.unwrap_or(0);
158        let start = from.max(ring.dropped).min(total);
159        let available = ring.chunks.iter().cloned().collect::<String>();
160        let chunk = slice_utf16(&available, start - ring.dropped, MAX_PROCESS_CHUNK_BYTES);
161        let read = utf16_len(&chunk);
162        Ok(json!({
163            "processId": args.process_id,
164            "chunk": chunk,
165            "nextCursor": start + read,
166            "skipped": from < ring.dropped,
167            "running": entry.running,
168            "exitCode": entry.exit_code,
169        }))
170    }
171
172    pub async fn send_input(&self, root: &Path, value: Value) -> Result<Value, ExeoraError> {
173        let args: InputArgs = parse(value)?;
174        let mut entries = self.entries.lock().await;
175        let entry = find_entry(&mut entries, root, &args.process_id)?;
176        refresh(entry).await;
177        if !entry.running {
178            return Err(ExeoraError::tool("That process is not accepting input."));
179        }
180        let payload = if args.newline.unwrap_or(true) {
181            format!("{}\n", args.data)
182        } else {
183            args.data
184        };
185        let mut stdin = entry.stdin.lock().await;
186        let Some(stdin) = stdin.as_mut() else {
187            return Err(ExeoraError::tool("That process is not accepting input."));
188        };
189        stdin
190            .write_all(payload.as_bytes())
191            .await
192            .map_err(|error| ExeoraError::tool(error.to_string()))?;
193        stdin
194            .flush()
195            .await
196            .map_err(|error| ExeoraError::tool(error.to_string()))?;
197        Ok(json!({ "processId": args.process_id, "bytesWritten": payload.len() }))
198    }
199
200    pub async fn kill_command(&self, root: &Path, value: Value) -> Result<Value, ExeoraError> {
201        let args: ProcessArgs = parse(value)?;
202        let mut entries = self.entries.lock().await;
203        let entry = find_entry(&mut entries, root, &args.process_id)?;
204        refresh(entry).await;
205        if !entry.running {
206            return Ok(
207                json!({ "processId": args.process_id, "killed": false, "exitCode": entry.exit_code }),
208            );
209        }
210        let mut child = entry.child.lock().await;
211        let _ = kill_child(child.as_mut()).await;
212        entry.running = false;
213        Ok(json!({ "processId": args.process_id, "killed": true, "exitCode": entry.exit_code }))
214    }
215
216    pub async fn kill_all(&self) {
217        let mut entries = self.entries.lock().await;
218        for entry in entries.values_mut() {
219            if entry.running {
220                let mut child = entry.child.lock().await;
221                let _ = kill_child(child.as_mut()).await;
222            }
223        }
224        entries.clear();
225    }
226}
227
228#[derive(Deserialize)]
229#[serde(rename_all = "camelCase")]
230struct RunArgs {
231    command: String,
232    cwd: Option<String>,
233    timeout_ms: Option<u64>,
234}
235#[derive(Deserialize)]
236struct StartArgs {
237    command: String,
238    cwd: Option<String>,
239}
240#[derive(Deserialize)]
241#[serde(rename_all = "camelCase")]
242struct OutputArgs {
243    process_id: String,
244    cursor: Option<usize>,
245}
246#[derive(Deserialize)]
247#[serde(rename_all = "camelCase")]
248struct InputArgs {
249    process_id: String,
250    data: String,
251    newline: Option<bool>,
252}
253#[derive(Deserialize)]
254#[serde(rename_all = "camelCase")]
255struct ProcessArgs {
256    process_id: String,
257}
258
259fn spawn_wrapped(
260    command: &str,
261    cwd: &Path,
262    input: bool,
263) -> Result<Box<dyn ChildWrapper>, ExeoraError> {
264    let (program, shell_args) = shell(command);
265    let mut wrapped = CommandWrap::with_new(program, |cmd| {
266        cmd.args(shell_args)
267            .current_dir(cwd)
268            .stdout(Stdio::piped())
269            .stderr(Stdio::piped())
270            .stdin(if input { Stdio::piped() } else { Stdio::null() });
271    });
272    #[cfg(unix)]
273    wrapped.wrap(ProcessGroup::leader());
274    #[cfg(windows)]
275    wrapped.wrap(JobObject);
276    wrapped.wrap(KillOnDrop);
277    wrapped
278        .spawn()
279        .map_err(|error| ExeoraError::tool(error.to_string()))
280}
281
282#[cfg(unix)]
283fn shell(command: &str) -> (&'static str, Vec<&str>) {
284    ("/bin/sh", vec!["-c", command])
285}
286#[cfg(windows)]
287fn shell(command: &str) -> (&'static str, Vec<&str>) {
288    ("cmd.exe", vec!["/d", "/s", "/c", command])
289}
290
291fn spawn_reader<R: AsyncRead + Unpin + Send + 'static>(reader: Option<R>, ring: Arc<Mutex<Ring>>) {
292    let Some(mut reader) = reader else {
293        return;
294    };
295    tokio::spawn(async move {
296        let mut buffer = vec![0; 8192];
297        while let Ok(count) = reader.read(&mut buffer).await {
298            if count == 0 {
299                break;
300            }
301            let mut guard = ring.lock().await;
302            append(
303                &mut guard,
304                String::from_utf8_lossy(&buffer[..count]).into_owned(),
305            );
306        }
307    });
308}
309
310async fn capture<R: AsyncRead + Unpin>(
311    reader: Option<R>,
312    max: usize,
313) -> Result<(String, bool), ExeoraError> {
314    let Some(mut reader) = reader else {
315        return Ok((String::new(), false));
316    };
317    let mut all = Vec::new();
318    reader
319        .read_to_end(&mut all)
320        .await
321        .map_err(|error| ExeoraError::tool(error.to_string()))?;
322    let cut = all.len() > max;
323    let kept = if cut { &all[all.len() - max..] } else { &all };
324    Ok((String::from_utf8_lossy(kept).into_owned(), cut))
325}
326
327fn append(ring: &mut Ring, text: String) {
328    let units = utf16_len(&text);
329    ring.units += units;
330    ring.chunks.push_back(text);
331    while ring.units > MAX_PROCESS_BUFFER_BYTES && ring.chunks.len() > 1 {
332        if let Some(oldest) = ring.chunks.pop_front() {
333            let old_units = utf16_len(&oldest);
334            ring.units -= old_units;
335            ring.dropped += old_units;
336        }
337    }
338}
339
340async fn refresh(entry: &mut Running) {
341    if !entry.running {
342        return;
343    }
344    if let Ok(Some(status)) = entry.child.lock().await.try_wait() {
345        entry.running = false;
346        entry.exit_code = status.code();
347    }
348}
349
350fn find_entry<'a>(
351    entries: &'a mut HashMap<String, Running>,
352    root: &Path,
353    id: &str,
354) -> Result<&'a mut Running, ExeoraError> {
355    let real_root = std::fs::canonicalize(root).unwrap_or_else(|_| root.to_owned());
356    entries
357        .get_mut(id)
358        .filter(|entry| entry.root == real_root)
359        .ok_or_else(|| {
360            ExeoraError::tool(
361                "No such process. It may have been stopped, or it belongs to another project.",
362            )
363        })
364}
365
366fn utf16_len(text: &str) -> usize {
367    if text.is_ascii() {
368        text.len()
369    } else {
370        text.encode_utf16().count()
371    }
372}
373fn slice_utf16(text: &str, start: usize, max: usize) -> String {
374    if text.is_ascii() {
375        let start = start.min(text.len());
376        return text[start..(start + max).min(text.len())].to_owned();
377    }
378    let mut position = 0;
379    let mut written = 0;
380    let mut output = String::new();
381    for ch in text.chars() {
382        let units = ch.len_utf16();
383        if position + units <= start {
384            position += units;
385            continue;
386        }
387        if written + units > max {
388            break;
389        }
390        output.push(ch);
391        position += units;
392        written += units;
393    }
394    output
395}
396
397fn parse<T: for<'de> Deserialize<'de>>(value: Value) -> Result<T, ExeoraError> {
398    serde_json::from_value(value)
399        .map_err(|error| ExeoraError::new(ErrorCode::InvalidArguments, error.to_string()))
400}
401fn join_error(error: tokio::task::JoinError) -> ExeoraError {
402    ExeoraError::tool(error.to_string())
403}
404
405async fn kill_child(child: &mut dyn ChildWrapper) -> std::io::Result<()> {
406    Box::into_pin(child.kill()).await
407}