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
41struct Chunk {
43 text: String,
44 units: usize,
45}
46
47#[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 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 if taken < wanted || written >= max {
97 break;
98 }
99 }
100 (output, written)
101 }
102}
103
104fn 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 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 let mut child = entry.child.lock().await;
328 let _ = child.start_kill();
329 drop(child);
330 entry.running = false;
331
332 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 #[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 #[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}