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

exeora_cli/tools/
processes.rs
