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    project_scope: String,
35    child: SharedChild,
36    stdin: Arc<Mutex<Option<tokio::process::ChildStdin>>>,
37    ring: Arc<Mutex<Ring>>,
38    exit_code: Option<i32>,
39    running: bool,
40}
41
42/// One read off a pipe, with its UTF-8 byte length measured once.
43struct Chunk {
44    text: String,
45    bytes: usize,
46}
47
48/**
49 * Output kept for one process, oldest chunk dropped first.
50 *
51 * Lengths and cursors count UTF-8 bytes, matching the shared protocol limits.
52 * Chunks keep their own length so neither trimming nor reading has to measure
53 * the whole buffer: a reader asking for 100,000 bytes out of a full 256,000
54 * should pay for what it asked for, not for what is being held.
55 */
56#[derive(Default)]
57struct Ring {
58    chunks: VecDeque<Chunk>,
59    bytes: usize,
60    dropped: usize,
61}
62
63impl Ring {
64    fn append(&mut self, text: String) {
65        let bytes = text.len();
66        self.bytes += bytes;
67        self.chunks.push_back(Chunk { text, bytes });
68        while self.bytes > MAX_PROCESS_BUFFER_BYTES {
69            let overflow = self.bytes - MAX_PROCESS_BUFFER_BYTES;
70            let Some(oldest) = self.chunks.front_mut() else {
71                break;
72            };
73            if oldest.bytes <= overflow {
74                let oldest = self.chunks.pop_front().expect("front exists");
75                self.bytes -= oldest.bytes;
76                self.dropped += oldest.bytes;
77                continue;
78            }
79
80            let mut cut = overflow;
81            while !oldest.text.is_char_boundary(cut) {
82                cut += 1;
83            }
84            oldest.text = oldest.text.split_off(cut);
85            oldest.bytes -= cut;
86            self.bytes -= cut;
87            self.dropped += cut;
88        }
89    }
90
91    /// Copies at most `max` bytes starting `offset` bytes into what is still held.
92    fn slice(&self, offset: usize, max: usize) -> (String, usize) {
93        let mut skipped = offset;
94        let mut output = String::with_capacity(max);
95        let mut consumed = 0;
96
97        for chunk in &self.chunks {
98            if skipped >= chunk.bytes {
99                skipped -= chunk.bytes;
100                continue;
101            }
102            let mut start = skipped;
103            while !chunk.text.is_char_boundary(start) {
104                start += 1;
105            }
106            consumed += start - skipped;
107            let budget = max.saturating_sub(output.len());
108            let mut end = (start + budget).min(chunk.bytes);
109            while end > start && !chunk.text.is_char_boundary(end) {
110                end -= 1;
111            }
112            output.push_str(&chunk.text[start..end]);
113            consumed += end - start;
114            skipped = 0;
115            if end < chunk.bytes || output.len() >= max {
116                break;
117            }
118        }
119        (output, consumed)
120    }
121}
122
123pub struct ProcessRegistry {
124    entries: Mutex<HashMap<String, Running>>,
125}
126
127impl Default for ProcessRegistry {
128    fn default() -> Self {
129        Self::new()
130    }
131}
132
133impl ProcessRegistry {
134    pub fn new() -> Self {
135        Self {
136            entries: Mutex::new(HashMap::new()),
137        }
138    }
139
140    pub async fn run_command(
141        &self,
142        root: &Path,
143        value: Value,
144        cancel: CancellationToken,
145    ) -> Result<Value, ExeoraError> {
146        let args: RunArgs = parse(value)?;
147        let (real_root, cwd) = resolve_in_project(root, args.cwd.as_deref().unwrap_or("."))?;
148        let timeout_ms = args.timeout_ms.unwrap_or(DEFAULT_COMMAND_TIMEOUT_MS);
149        let mut child = spawn_wrapped(&args.command, &real_root.join(cwd), false)?;
150        let stdout = child.stdout().take();
151        let stderr = child.stderr().take();
152        let captured = Arc::new(Mutex::new(CapturedOutput::default()));
153        let stdout_task = tokio::spawn(capture(stdout, OutputStream::Stdout, captured.clone()));
154        let stderr_task = tokio::spawn(capture(stderr, OutputStream::Stderr, captured.clone()));
155
156        let mut timed_out = false;
157        let mut cancelled = false;
158        let status = {
159            let wait = child.wait();
160            tokio::pin!(wait);
161            tokio::select! {
162                status = &mut wait => Some(status.map_err(|error| ExeoraError::tool(error.to_string()))?),
163                _ = tokio::time::sleep(Duration::from_millis(timeout_ms)) => { timed_out = true; None },
164                _ = cancel.cancelled() => { cancelled = true; None },
165            }
166        };
167        if status.is_none() {
168            let _ = kill_child(child.as_mut()).await;
169        }
170        stdout_task.await.map_err(join_error)??;
171        stderr_task.await.map_err(join_error)??;
172        let captured = std::mem::take(&mut *captured.lock().await);
173        let truncated = captured.truncated;
174        let (stdout, stderr) = captured.into_strings();
175        if cancelled {
176            return Err(ExeoraError::new(
177                ErrorCode::Cancelled,
178                "The call was cancelled while the command was running.",
179            ));
180        }
181        Ok(json!({
182            "command": args.command,
183            "exitCode": status.and_then(|status| status.code()),
184            "stdout": stdout,
185            "stderr": stderr,
186            "truncated": truncated,
187            "timedOut": timed_out,
188        }))
189    }
190
191    pub async fn start_command(
192        &self,
193        root: &Path,
194        project_scope: &str,
195        value: Value,
196    ) -> Result<Value, ExeoraError> {
197        let args: StartArgs = parse(value)?;
198        let (real_root, cwd) = resolve_in_project(root, args.cwd.as_deref().unwrap_or("."))?;
199        let mut entries = self.entries.lock().await;
200        for entry in entries.values_mut() {
201            refresh(entry).await;
202        }
203        if entries
204            .values()
205            .filter(|entry| entry.project_scope == project_scope && entry.running)
206            .count()
207            >= MAX_PROCESSES_PER_PROJECT
208        {
209            return Err(ExeoraError::tool(format!(
210                "This project already has {MAX_PROCESSES_PER_PROJECT} processes running. Stop one with kill_command before starting another."
211            )));
212        }
213        let mut child = spawn_wrapped(&args.command, &real_root.join(cwd), true)?;
214        let pid = child.id();
215        let stdin = Arc::new(Mutex::new(child.stdin().take()));
216        let stdout = child.stdout().take();
217        let stderr = child.stderr().take();
218        let ring = Arc::new(Mutex::new(Ring::default()));
219        spawn_reader(stdout, ring.clone());
220        spawn_reader(stderr, ring.clone());
221        let id = format!("proc_{}", Uuid::new_v4().simple());
222        entries.insert(
223            id.clone(),
224            Running {
225                root: real_root,
226                project_scope: project_scope.to_owned(),
227                child: Arc::new(Mutex::new(child)),
228                stdin,
229                ring,
230                exit_code: None,
231                running: true,
232            },
233        );
234        Ok(json!({ "processId": id, "command": args.command, "pid": pid }))
235    }
236
237    pub async fn get_output(&self, root: &Path, value: Value) -> Result<Value, ExeoraError> {
238        let args: OutputArgs = parse(value)?;
239        let mut entries = self.entries.lock().await;
240        let entry = find_entry(&mut entries, root, &args.process_id)?;
241        refresh(entry).await;
242        let ring = entry.ring.lock().await;
243        let total = ring.dropped + ring.bytes;
244        let from = args.cursor.unwrap_or(0);
245        let start = from.max(ring.dropped).min(total);
246        let (chunk, read) = ring.slice(start - ring.dropped, MAX_PROCESS_CHUNK_BYTES);
247        Ok(json!({
248            "processId": args.process_id,
249            "chunk": chunk,
250            "nextCursor": start + read,
251            "skipped": from < ring.dropped,
252            "running": entry.running,
253            "exitCode": entry.exit_code,
254        }))
255    }
256
257    pub async fn send_input(&self, root: &Path, value: Value) -> Result<Value, ExeoraError> {
258        let args: InputArgs = parse(value)?;
259        let mut entries = self.entries.lock().await;
260        let entry = find_entry(&mut entries, root, &args.process_id)?;
261        if !entry.running {
262            return Err(ExeoraError::tool("That process is not accepting input."));
263        }
264        let payload = if args.newline.unwrap_or(true) {
265            format!("{}\n", args.data)
266        } else {
267            args.data
268        };
269
270        // Deliberately not refreshed first. Asking the kernel whether the child
271        // is still alive is a syscall on every keystroke to learn what a failed
272        // write reports anyway, and the answer would be stale by the time it is
273        // used. The exit is confirmed only once writing has actually failed.
274        let mut stdin = entry.stdin.lock().await;
275        let written = match stdin.as_mut() {
276            None => Err(std::io::ErrorKind::BrokenPipe.into()),
277            Some(stdin) => match stdin.write_all(payload.as_bytes()).await {
278                Ok(()) => stdin.flush().await,
279                Err(error) => Err(error),
280            },
281        };
282        drop(stdin);
283
284        if let Err(error) = written {
285            refresh(entry).await;
286            return Err(if entry.running {
287                ExeoraError::tool(error.to_string())
288            } else {
289                ExeoraError::tool("That process is not accepting input.")
290            });
291        }
292        Ok(json!({ "processId": args.process_id, "bytesWritten": payload.len() }))
293    }
294
295    pub async fn kill_command(&self, root: &Path, value: Value) -> Result<Value, ExeoraError> {
296        let args: ProcessArgs = parse(value)?;
297        let mut entries = self.entries.lock().await;
298        let entry = find_entry(&mut entries, root, &args.process_id)?;
299        refresh(entry).await;
300        if !entry.running {
301            return Ok(
302                json!({ "processId": args.process_id, "killed": false, "exitCode": entry.exit_code }),
303            );
304        }
305        // Signal the group and answer, rather than waiting for the reap. The
306        // status is not in the reply either way: `exit_code` is whatever the
307        // refresh above saw, and a process still alive a moment ago has none.
308        // Waiting costs the caller a full wait-and-retry loop to learn nothing.
309        let mut child = entry.child.lock().await;
310        let _ = child.start_kill();
311        drop(child);
312        entry.running = false;
313
314        // The reap still has to happen somewhere. Nothing else will do it: the
315        // entry stays in the map, so its child is never dropped, and `refresh`
316        // walks away from an entry already marked stopped. Left alone the
317        // killed group is a zombie for the rest of the session.
318        let child = entry.child.clone();
319        tokio::spawn(async move {
320            let mut child = child.lock().await;
321            let _ = child.wait().await;
322        });
323        Ok(json!({ "processId": args.process_id, "killed": true, "exitCode": entry.exit_code }))
324    }
325
326    pub async fn kill_all(&self) {
327        let mut entries = self.entries.lock().await;
328        for entry in entries.values_mut() {
329            if entry.running {
330                let mut child = entry.child.lock().await;
331                let _ = kill_child(child.as_mut()).await;
332            }
333        }
334        entries.clear();
335    }
336
337    pub async fn kill_root(&self, root: &Path) {
338        let root = std::fs::canonicalize(root).unwrap_or_else(|_| root.to_path_buf());
339        let mut entries = self.entries.lock().await;
340        let ids: Vec<_> = entries
341            .iter()
342            .filter(|(_, entry)| entry.root == root)
343            .map(|(id, _)| id.clone())
344            .collect();
345        for id in ids {
346            if let Some(mut entry) = entries.remove(&id)
347                && entry.running
348            {
349                let mut child = entry.child.lock().await;
350                let _ = kill_child(child.as_mut()).await;
351                entry.running = false;
352            }
353        }
354    }
355}
356
357#[derive(Deserialize)]
358#[serde(rename_all = "camelCase")]
359struct RunArgs {
360    command: String,
361    cwd: Option<String>,
362    timeout_ms: Option<u64>,
363}
364#[derive(Deserialize)]
365struct StartArgs {
366    command: String,
367    cwd: Option<String>,
368}
369#[derive(Deserialize)]
370#[serde(rename_all = "camelCase")]
371struct OutputArgs {
372    process_id: String,
373    cursor: Option<usize>,
374}
375#[derive(Deserialize)]
376#[serde(rename_all = "camelCase")]
377struct InputArgs {
378    process_id: String,
379    data: String,
380    newline: Option<bool>,
381}
382#[derive(Deserialize)]
383#[serde(rename_all = "camelCase")]
384struct ProcessArgs {
385    process_id: String,
386}
387
388fn spawn_wrapped(
389    command: &str,
390    cwd: &Path,
391    input: bool,
392) -> Result<Box<dyn ChildWrapper>, ExeoraError> {
393    let (program, shell_args) = shell(command);
394    let mut wrapped = CommandWrap::with_new(program, |cmd| {
395        cmd.args(shell_args)
396            .current_dir(cwd)
397            .stdout(Stdio::piped())
398            .stderr(Stdio::piped())
399            .stdin(if input { Stdio::piped() } else { Stdio::null() });
400    });
401    #[cfg(unix)]
402    wrapped.wrap(ProcessGroup::leader());
403    #[cfg(windows)]
404    wrapped.wrap(JobObject);
405    wrapped.wrap(KillOnDrop);
406    wrapped
407        .spawn()
408        .map_err(|error| ExeoraError::tool(error.to_string()))
409}
410
411#[cfg(unix)]
412fn shell(command: &str) -> (&'static str, Vec<&str>) {
413    ("/bin/sh", vec!["-c", command])
414}
415#[cfg(windows)]
416fn shell(command: &str) -> (&'static str, Vec<&str>) {
417    ("cmd.exe", vec!["/d", "/s", "/c", command])
418}
419
420fn spawn_reader<R: AsyncRead + Unpin + Send + 'static>(reader: Option<R>, ring: Arc<Mutex<Ring>>) {
421    let Some(mut reader) = reader else {
422        return;
423    };
424    tokio::spawn(async move {
425        let mut buffer = vec![0; 8192];
426        while let Ok(count) = reader.read(&mut buffer).await {
427            if count == 0 {
428                break;
429            }
430            let mut guard = ring.lock().await;
431            guard.append(String::from_utf8_lossy(&buffer[..count]).into_owned());
432        }
433    });
434}
435
436#[derive(Clone, Copy)]
437enum OutputStream {
438    Stdout,
439    Stderr,
440}
441
442struct OutputChunk {
443    stream: OutputStream,
444    bytes: Vec<u8>,
445}
446
447#[derive(Default)]
448struct CapturedOutput {
449    chunks: VecDeque<OutputChunk>,
450    bytes: usize,
451    truncated: bool,
452}
453
454impl CapturedOutput {
455    fn append(&mut self, stream: OutputStream, mut bytes: Vec<u8>) {
456        if bytes.len() > MAX_COMMAND_OUTPUT_BYTES {
457            self.truncated = true;
458            bytes.drain(..bytes.len() - MAX_COMMAND_OUTPUT_BYTES);
459        }
460        self.bytes += bytes.len();
461        self.chunks.push_back(OutputChunk { stream, bytes });
462        while self.bytes > MAX_COMMAND_OUTPUT_BYTES {
463            self.truncated = true;
464            let overflow = self.bytes - MAX_COMMAND_OUTPUT_BYTES;
465            let Some(oldest) = self.chunks.front_mut() else {
466                break;
467            };
468            if oldest.bytes.len() <= overflow {
469                let oldest = self.chunks.pop_front().expect("front exists");
470                self.bytes -= oldest.bytes.len();
471            } else {
472                oldest.bytes.drain(..overflow);
473                self.bytes -= overflow;
474            }
475        }
476    }
477
478    fn into_strings(self) -> (String, String) {
479        let mut stdout = Vec::new();
480        let mut stderr = Vec::new();
481        for chunk in self.chunks {
482            match chunk.stream {
483                OutputStream::Stdout => stdout.extend(chunk.bytes),
484                OutputStream::Stderr => stderr.extend(chunk.bytes),
485            }
486        }
487        (
488            String::from_utf8_lossy(&stdout).into_owned(),
489            String::from_utf8_lossy(&stderr).into_owned(),
490        )
491    }
492}
493
494async fn capture<R: AsyncRead + Unpin>(
495    reader: Option<R>,
496    stream: OutputStream,
497    captured: Arc<Mutex<CapturedOutput>>,
498) -> Result<(), ExeoraError> {
499    let Some(mut reader) = reader else {
500        return Ok(());
501    };
502    let mut buffer = vec![0; 8192];
503    loop {
504        let count = reader
505            .read(&mut buffer)
506            .await
507            .map_err(|error| ExeoraError::tool(error.to_string()))?;
508        if count == 0 {
509            break;
510        }
511        captured
512            .lock()
513            .await
514            .append(stream, buffer[..count].to_vec());
515    }
516    Ok(())
517}
518
519async fn refresh(entry: &mut Running) {
520    if !entry.running {
521        return;
522    }
523    if let Ok(Some(status)) = entry.child.lock().await.try_wait() {
524        entry.running = false;
525        entry.exit_code = status.code();
526    }
527}
528
529fn find_entry<'a>(
530    entries: &'a mut HashMap<String, Running>,
531    root: &Path,
532    id: &str,
533) -> Result<&'a mut Running, ExeoraError> {
534    let real_root = std::fs::canonicalize(root).unwrap_or_else(|_| root.to_owned());
535    entries
536        .get_mut(id)
537        .filter(|entry| entry.root == real_root)
538        .ok_or_else(|| {
539            ExeoraError::tool(
540                "No such process. It may have been stopped, or it belongs to another project.",
541            )
542        })
543}
544
545fn parse<T: for<'de> Deserialize<'de>>(value: Value) -> Result<T, ExeoraError> {
546    serde_json::from_value(value)
547        .map_err(|error| ExeoraError::new(ErrorCode::InvalidArguments, error.to_string()))
548}
549fn join_error(error: tokio::task::JoinError) -> ExeoraError {
550    ExeoraError::tool(error.to_string())
551}
552
553async fn kill_child(child: &mut dyn ChildWrapper) -> std::io::Result<()> {
554    Box::into_pin(child.kill()).await
555}
556
557#[cfg(test)]
558mod tests {
559    use super::{CapturedOutput, OutputStream, Ring};
560    use crate::protocol::{MAX_COMMAND_OUTPUT_BYTES, MAX_PROCESS_BUFFER_BYTES};
561
562    #[test]
563    fn a_multibyte_character_at_the_byte_limit_waits_for_the_next_read() {
564        let mut ring = Ring::default();
565        ring.append("a".repeat(9));
566        ring.append("\u{1f600}tail".to_owned());
567        ring.append("later".to_owned());
568
569        let (head, read) = ring.slice(0, 10);
570        assert_eq!(head, "a".repeat(9));
571        assert_eq!(read, 9, "the character is left for the next read");
572
573        let (tail, read) = ring.slice(read, 8);
574        assert_eq!(tail, "\u{1f600}tail");
575        assert_eq!(read, 8);
576    }
577
578    #[test]
579    fn a_cursor_inside_a_character_advances_past_it() {
580        let mut ring = Ring::default();
581        ring.append("\u{1f600}tail".to_owned());
582
583        let (chunk, read) = ring.slice(1, 10);
584        assert_eq!(chunk, "tail");
585        assert_eq!(read, 7, "three skipped bytes and four bytes of tail");
586    }
587
588    #[test]
589    fn one_large_chunk_is_trimmed_to_the_process_byte_limit() {
590        let mut ring = Ring::default();
591        let input = "\u{00e9}".repeat(MAX_PROCESS_BUFFER_BYTES);
592        let input_bytes = input.len();
593        ring.append(input);
594
595        assert!(ring.bytes <= MAX_PROCESS_BUFFER_BYTES);
596        assert_eq!(ring.dropped + ring.bytes, input_bytes);
597    }
598
599    #[test]
600    fn stdout_and_stderr_share_one_command_output_budget() {
601        let mut output = CapturedOutput::default();
602        output.append(OutputStream::Stdout, vec![b'o'; 150_000]);
603        output.append(OutputStream::Stderr, vec![b'e'; 100_000]);
604        assert!(output.truncated);
605        assert_eq!(output.bytes, MAX_COMMAND_OUTPUT_BYTES);
606
607        let (stdout, stderr) = output.into_strings();
608        assert_eq!(stdout.len() + stderr.len(), MAX_COMMAND_OUTPUT_BYTES);
609    }
610}
exeora_cli/tools/processes.rs

exeora_cli/tools/
processes.rs
