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