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}