1mod lifecycle;
7#[cfg(test)]
8mod pipe_tests;
9
10pub use lifecycle::{ExecSessionEvent, ExecSessionSnapshot, ExecSessionStatus};
11
12use std::collections::{BTreeMap, VecDeque};
13use std::path::PathBuf;
14use std::process::ExitStatus;
15use std::sync::atomic::{AtomicU64, Ordering};
16use std::sync::{Arc, Mutex};
17use std::time::{Duration, Instant};
18
19use orchestral_core::agent_protocol::wire::RunId;
20use orchestral_core::tool_protocol::ToolOperationPlan;
21use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt};
22use tokio::process::{Child, ChildStdin, Command};
23use tokio::sync::{broadcast, Mutex as AsyncMutex, Notify};
24use tokio_util::sync::CancellationToken;
25
26use crate::pty_process::{PtyProcessId, PtyProcessManager, PtyReadOptions, PtySpawnSpec};
27
28mod runtime_temp;
29use runtime_temp::{RuntimeTempDirectory, RuntimeTempRoot};
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
32pub struct ExecSessionId(u64);
33
34impl ExecSessionId {
35 pub fn new(value: u64) -> Result<Self, ExecProcessError> {
36 if value == 0 {
37 return Err(ExecProcessError::Invalid(
38 "exec session ID must be positive".to_owned(),
39 ));
40 }
41 Ok(Self(value))
42 }
43
44 pub fn get(self) -> u64 {
45 self.0
46 }
47}
48
49#[derive(Debug, Clone)]
50pub struct ExecSpawnSpec {
51 pub run_id: RunId,
52 pub program: String,
53 pub args: Vec<String>,
54 pub cwd: PathBuf,
55 pub environment: BTreeMap<String, String>,
56 pub tty: bool,
57 pub backend_starts_new_session: bool,
58 pub operation: ToolOperationPlan,
60}
61
62#[derive(Debug, Clone, PartialEq)]
63pub struct ExecPollResult {
64 pub stdout: String,
65 pub stderr: String,
66 pub dropped_bytes: u64,
67 pub alive: bool,
70 pub exit_code: Option<i32>,
71 pub wall_time_seconds: f64,
72}
73
74#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76pub enum ExecWaitMode {
77 Output,
79 Completion,
81}
82
83impl ExecWaitMode {
84 pub fn for_interaction(tty: bool, has_input: bool) -> Self {
86 if tty || has_input {
87 Self::Output
88 } else {
89 Self::Completion
90 }
91 }
92}
93
94#[derive(Debug, Clone)]
96pub struct ExecWaitOptions {
97 pub duration: Duration,
98 pub mode: ExecWaitMode,
99 pub yield_requested: CancellationToken,
101}
102
103#[derive(Debug, thiserror::Error)]
104#[non_exhaustive]
105pub enum ExecProcessError {
106 #[error("invalid exec operation: {0}")]
107 Invalid(String),
108 #[error("exec session was not found in this Run: {0}")]
109 NotFound(u64),
110 #[error("exec process operation was cancelled")]
111 Cancelled,
112 #[error("exec process manager state is unavailable")]
113 Unavailable,
114 #[error("exec process I/O failed: {0}")]
115 Io(String),
116}
117
118type SessionKey = (RunId, ExecSessionId);
119
120#[derive(Clone)]
121enum ManagedProcess {
122 Pipe(Arc<PipeSession>),
123 Pty { process_id: PtyProcessId },
124}
125
126#[derive(Clone)]
127struct ManagedSession {
128 process: ManagedProcess,
129 lifecycle: Arc<SessionLifecycle>,
130 started: Instant,
131 tty: bool,
132 operation: ToolOperationPlan,
133 _runtime_temp: Option<Arc<RuntimeTempDirectory>>,
135}
136
137use lifecycle::SessionLifecycle;
138
139#[derive(Default)]
140struct OutputState {
141 bytes: VecDeque<u8>,
142 dropped_bytes: u64,
143 generation: u64,
144 closed: bool,
145}
146
147struct SharedOutput {
148 state: Mutex<OutputState>,
149 changed: Notify,
150 max_bytes: usize,
151}
152
153impl SharedOutput {
154 fn new(max_bytes: usize) -> Arc<Self> {
155 Arc::new(Self {
156 state: Mutex::new(OutputState::default()),
157 changed: Notify::new(),
158 max_bytes,
159 })
160 }
161
162 fn push(&self, bytes: &[u8]) {
163 if let Ok(mut state) = self.state.lock() {
164 for byte in bytes {
165 if state.bytes.len() == self.max_bytes {
166 state.bytes.pop_front();
167 state.dropped_bytes = state.dropped_bytes.saturating_add(1);
168 }
169 state.bytes.push_back(*byte);
170 }
171 state.generation = state.generation.saturating_add(1);
172 }
173 self.changed.notify_waiters();
174 }
175
176 fn close(&self) {
177 if let Ok(mut state) = self.state.lock() {
178 state.closed = true;
179 state.generation = state.generation.saturating_add(1);
180 }
181 self.changed.notify_waiters();
182 }
183
184 fn snapshot(&self) -> Result<(u64, bool, bool), ExecProcessError> {
185 let state = self
186 .state
187 .lock()
188 .map_err(|_| ExecProcessError::Unavailable)?;
189 Ok((state.generation, state.closed, !state.bytes.is_empty()))
190 }
191
192 fn drain(&self) -> Result<(String, u64), ExecProcessError> {
193 let mut state = self
194 .state
195 .lock()
196 .map_err(|_| ExecProcessError::Unavailable)?;
197 let raw = state.bytes.drain(..).collect::<Vec<_>>();
198 let dropped = std::mem::take(&mut state.dropped_bytes);
199 Ok((String::from_utf8_lossy(&raw).into_owned(), dropped))
200 }
201
202 async fn wait_closed(&self) -> Result<(), ExecProcessError> {
203 loop {
204 let changed = self.changed.notified();
205 tokio::pin!(changed);
206 changed.as_mut().enable();
209 if self.snapshot()?.1 {
210 return Ok(());
211 }
212 changed.await;
213 }
214 }
215}
216
217struct PipeSession {
218 child: AsyncMutex<Child>,
219 stdin: AsyncMutex<Option<ChildStdin>>,
220 stdout: Arc<SharedOutput>,
221 stderr: Arc<SharedOutput>,
222 stop_readers: CancellationToken,
223 process_group_id: Option<u32>,
224 #[cfg(windows)]
225 job: crate::windows_process_job::ProcessJob,
226}
227
228impl PipeSession {
229 fn spawn(spec: &ExecSpawnSpec, max_output_bytes: usize) -> Result<Arc<Self>, ExecProcessError> {
230 if spec.run_id.is_empty() || spec.program.trim().is_empty() || !spec.cwd.is_absolute() {
231 return Err(ExecProcessError::Invalid(
232 "exec spawn requires run/program/absolute cwd".to_owned(),
233 ));
234 }
235 let mut command = Command::new(&spec.program);
236 command
237 .args(&spec.args)
238 .env_clear()
239 .envs(&spec.environment)
240 .current_dir(&spec.cwd)
241 .stdin(std::process::Stdio::piped())
242 .stdout(std::process::Stdio::piped())
243 .stderr(std::process::Stdio::piped())
244 .kill_on_drop(true);
245 isolate_process_group(&mut command, spec.backend_starts_new_session);
246 let mut child = command
247 .spawn()
248 .map_err(|error| ExecProcessError::Io(error.to_string()))?;
249 #[cfg(windows)]
250 let job = match child
251 .raw_handle()
252 .ok_or_else(|| std::io::Error::other("child process handle unavailable"))
253 .and_then(crate::windows_process_job::ProcessJob::attach)
254 {
255 Ok(job) => job,
256 Err(error) => {
257 let _ = child.start_kill();
258 return Err(ExecProcessError::Io(format!(
259 "could not supervise Windows process tree: {error}"
260 )));
261 }
262 };
263 let process_group_id = child.id();
264 let stdin = child.stdin.take();
265 let stdout = child
266 .stdout
267 .take()
268 .ok_or_else(|| ExecProcessError::Io("exec stdout pipe was not created".to_owned()))?;
269 let stderr = child
270 .stderr
271 .take()
272 .ok_or_else(|| ExecProcessError::Io("exec stderr pipe was not created".to_owned()))?;
273 let stdout_buffer = SharedOutput::new(max_output_bytes);
274 let stderr_buffer = SharedOutput::new(max_output_bytes);
275 let stop_readers = CancellationToken::new();
276 spawn_reader(stdout, stdout_buffer.clone(), stop_readers.clone());
277 spawn_reader(stderr, stderr_buffer.clone(), stop_readers.clone());
278 Ok(Arc::new(Self {
279 child: AsyncMutex::new(child),
280 stdin: AsyncMutex::new(stdin),
281 stdout: stdout_buffer,
282 stderr: stderr_buffer,
283 stop_readers,
284 process_group_id,
285 #[cfg(windows)]
286 job,
287 }))
288 }
289
290 async fn send(&self, input: &str) -> Result<(), ExecProcessError> {
291 if input.is_empty() {
292 return Ok(());
293 }
294 let mut stdin = self.stdin.lock().await;
295 let stdin = stdin
296 .as_mut()
297 .ok_or_else(|| ExecProcessError::Io("exec stdin is closed".to_owned()))?;
298 stdin
299 .write_all(input.as_bytes())
300 .await
301 .map_err(|error| ExecProcessError::Io(error.to_string()))?;
302 stdin
303 .flush()
304 .await
305 .map_err(|error| ExecProcessError::Io(error.to_string()))
306 }
307
308 async fn poll(
309 &self,
310 lifecycle: &SessionLifecycle,
311 started_at: Instant,
312 options: &ExecWaitOptions,
313 cancellation: &CancellationToken,
314 ) -> Result<ExecPollResult, ExecProcessError> {
315 poll_pipe_output(
316 &self.stdout,
317 &self.stderr,
318 lifecycle,
319 started_at,
320 options,
321 cancellation,
322 )
323 .await
324 }
325
326 async fn terminate(&self) {
327 self.stdin.lock().await.take();
328 let mut child = self.child.lock().await;
329 terminate_process_group(self.process_group_id);
330 #[cfg(windows)]
331 self.job.terminate();
332 let _ = child.start_kill();
333 let _ = child.wait().await;
334 self.stop_readers.cancel();
335 let _ = tokio::join!(self.stdout.wait_closed(), self.stderr.wait_closed());
336 }
337}
338
339impl Drop for PipeSession {
340 fn drop(&mut self) {
341 self.stop_readers.cancel();
344 }
345}
346
347async fn poll_pipe_output(
348 stdout: &SharedOutput,
349 stderr: &SharedOutput,
350 lifecycle: &SessionLifecycle,
351 started_at: Instant,
352 options: &ExecWaitOptions,
353 cancellation: &CancellationToken,
354) -> Result<ExecPollResult, ExecProcessError> {
355 let wait = options.duration;
356 let started = Instant::now();
357 let settle = Duration::from_millis(50).min(wait);
358 let mut observed = (u64::MAX, u64::MAX);
359 let mut last_change = Instant::now();
360 let exit_code = loop {
361 if cancellation.is_cancelled() {
362 return Err(ExecProcessError::Cancelled);
363 }
364 let status = lifecycle.status()?;
365 let exit_code = match &status {
366 ExecSessionStatus::Running => None,
367 ExecSessionStatus::Exited { exit_code } => Some(*exit_code),
368 ExecSessionStatus::Terminated => return Err(ExecProcessError::Cancelled),
369 ExecSessionStatus::Failed { message } => {
370 return Err(ExecProcessError::Io(message.clone()))
371 }
372 };
373 let stdout_state = stdout.snapshot()?;
374 let stderr_state = stderr.snapshot()?;
375 let generation = (stdout_state.0, stderr_state.0);
376 if generation != observed {
377 observed = generation;
378 last_change = Instant::now();
379 }
380 if exit_code.is_some() && stdout_state.1 && stderr_state.1 {
381 break exit_code;
382 }
383 if options.yield_requested.is_cancelled()
384 || started.elapsed() >= wait
385 || (options.mode == ExecWaitMode::Output
386 && (stdout_state.2 || stderr_state.2)
387 && last_change.elapsed() >= settle)
388 {
389 break None;
393 }
394 let remaining = wait.saturating_sub(started.elapsed());
395 let pause = remaining.min(Duration::from_millis(20));
396 tokio::select! {
397 _ = cancellation.cancelled() => return Err(ExecProcessError::Cancelled),
398 _ = options.yield_requested.cancelled() => {},
399 _ = stdout.changed.notified() => {},
400 _ = stderr.changed.notified() => {},
401 _ = lifecycle.changed.notified() => {},
402 _ = tokio::time::sleep(pause) => {},
403 }
404 };
405 let (stdout, stdout_dropped) = stdout.drain()?;
406 let (stderr, stderr_dropped) = stderr.drain()?;
407 Ok(ExecPollResult {
408 stdout,
409 stderr,
410 dropped_bytes: stdout_dropped.saturating_add(stderr_dropped),
411 alive: exit_code.is_none(),
412 exit_code,
413 wall_time_seconds: started_at.elapsed().as_secs_f64(),
414 })
415}
416
417fn spawn_reader<R>(mut reader: R, output: Arc<SharedOutput>, stop: CancellationToken)
418where
419 R: AsyncRead + Unpin + Send + 'static,
420{
421 tokio::spawn(async move {
422 let mut chunk = [0_u8; 8192];
423 loop {
424 let read = tokio::select! {
425 biased;
426 _ = stop.cancelled() => break,
427 read = reader.read(&mut chunk) => read,
428 };
429 match read {
430 Ok(0) => break,
431 Ok(count) => output.push(&chunk[..count]),
432 Err(_) => break,
433 }
434 }
435 output.close();
436 });
437}
438
439const PROCESS_EVENT_BUFFER: usize = 256;
440const PROCESS_WATCH_INTERVAL: Duration = Duration::from_millis(20);
441
442pub struct ProcessSupervisor {
444 sessions: Mutex<BTreeMap<SessionKey, ManagedSession>>,
445 next_session_id: AtomicU64,
446 pty: Arc<PtyProcessManager>,
447 max_output_bytes: usize,
448 events: broadcast::Sender<ExecSessionEvent>,
449 runtime_temp_root: Arc<RuntimeTempRoot>,
450 runtime_temps: Mutex<BTreeMap<RunId, Arc<RuntimeTempDirectory>>>,
451}
452
453impl ProcessSupervisor {
454 pub fn new(max_output_bytes: usize) -> Result<Self, ExecProcessError> {
455 Self::with_temp_root(max_output_bytes, RuntimeTempRoot::temporary()?)
456 }
457
458 pub fn new_with_runtime_temp_root(
462 max_output_bytes: usize,
463 root: impl AsRef<std::path::Path>,
464 ) -> Result<Self, ExecProcessError> {
465 Self::with_temp_root(max_output_bytes, RuntimeTempRoot::open(root.as_ref())?)
466 }
467
468 fn with_temp_root(
469 max_output_bytes: usize,
470 runtime_temp_root: RuntimeTempRoot,
471 ) -> Result<Self, ExecProcessError> {
472 if max_output_bytes == 0 {
473 return Err(ExecProcessError::Invalid(
474 "exec output limit must be positive".to_owned(),
475 ));
476 }
477 let pty = PtyProcessManager::new(max_output_bytes, Duration::from_secs(10 * 60))
478 .map_err(|error| ExecProcessError::Io(error.to_string()))?;
479 let (events, _) = broadcast::channel(PROCESS_EVENT_BUFFER);
480 Ok(Self {
481 sessions: Mutex::new(BTreeMap::new()),
482 next_session_id: AtomicU64::new(1),
483 pty: Arc::new(pty),
484 max_output_bytes,
485 events,
486 runtime_temp_root: Arc::new(runtime_temp_root),
487 runtime_temps: Mutex::new(BTreeMap::new()),
488 })
489 }
490
491 pub fn runtime_temp_root(&self) -> &std::path::Path {
494 self.runtime_temp_root.path()
495 }
496
497 pub(crate) fn runtime_temp_path(&self, run_id: &RunId) -> PathBuf {
498 self.runtime_temp_root.run_path(run_id)
499 }
500
501 pub(crate) fn prepare_runtime_temp(
502 self: &Arc<Self>,
503 run_id: &RunId,
504 run_cancellation: CancellationToken,
505 ) -> Result<PathBuf, ExecProcessError> {
506 if run_cancellation.is_cancelled() {
507 return Err(ExecProcessError::Cancelled);
508 }
509 let mut directories = self
510 .runtime_temps
511 .lock()
512 .map_err(|_| ExecProcessError::Unavailable)?;
513 if let Some(directory) = directories.get(run_id) {
514 return Ok(directory.path().to_owned());
515 }
516 let directory = self.runtime_temp_root.create_run(run_id)?;
517 let path = directory.path().to_owned();
518 directories.insert(run_id.clone(), directory);
519 let manager = Arc::downgrade(self);
520 let run_id = run_id.clone();
521 tokio::spawn(async move {
522 run_cancellation.cancelled().await;
523 if let Some(manager) = manager.upgrade() {
524 let _ = manager.close_run(&run_id).await;
525 }
526 });
527 Ok(path)
528 }
529
530 pub fn subscribe(&self) -> broadcast::Receiver<ExecSessionEvent> {
531 self.events.subscribe()
532 }
533
534 pub async fn spawn(&self, spec: ExecSpawnSpec) -> Result<ExecSessionId, ExecProcessError> {
535 spec.operation
536 .validate_shape()
537 .map_err(|error| ExecProcessError::Invalid(error.message))?;
538 let session_id = ExecSessionId::new(self.next_session_id.fetch_add(1, Ordering::Relaxed))?;
539 let lifecycle = SessionLifecycle::running();
540 let started = Instant::now();
541 let process = if spec.tty {
542 let process_id = PtyProcessId::new(format!("exec-{}", session_id.get()))
543 .map_err(|error| ExecProcessError::Invalid(error.to_string()))?;
544 let pty_spec = PtySpawnSpec {
545 run_id: spec.run_id.clone(),
546 process_id: process_id.clone(),
547 program: spec.program.clone(),
548 args: spec.args.clone(),
549 cwd: spec.cwd.clone(),
550 environment: spec.environment.clone(),
551 rows: 24,
552 cols: 120,
553 };
554 let pty = self.pty.clone();
555 tokio::task::spawn_blocking(move || pty.create(pty_spec))
556 .await
557 .map_err(|error| ExecProcessError::Io(error.to_string()))?
558 .map_err(|error| ExecProcessError::Io(error.to_string()))?;
559 ManagedProcess::Pty { process_id }
560 } else {
561 ManagedProcess::Pipe(PipeSession::spawn(&spec, self.max_output_bytes)?)
562 };
563 let session = ManagedSession {
564 process,
565 lifecycle,
566 started,
567 tty: spec.tty,
568 operation: spec.operation,
569 _runtime_temp: self
570 .runtime_temps
571 .lock()
572 .map_err(|_| ExecProcessError::Unavailable)?
573 .get(&spec.run_id)
574 .cloned(),
575 };
576 let key = (spec.run_id, session_id);
577 self.sessions
578 .lock()
579 .map_err(|_| ExecProcessError::Unavailable)?
580 .insert(key.clone(), session.clone());
581 publish_session_event(&self.events, &key, &session);
582 spawn_exit_watcher(key, session, self.pty.clone(), self.events.clone());
583 Ok(session_id)
584 }
585
586 pub async fn write_and_poll(
589 &self,
590 run_id: &RunId,
591 session_id: ExecSessionId,
592 input: Option<&str>,
593 wait: Duration,
594 cancellation: &CancellationToken,
595 ) -> Result<ExecPollResult, ExecProcessError> {
596 let tty = self.snapshot(run_id, session_id)?.tty;
597 self.write_and_poll_with_options(
598 run_id,
599 session_id,
600 input,
601 ExecWaitOptions {
602 duration: wait,
603 mode: ExecWaitMode::for_interaction(tty, input.is_some_and(|v| !v.is_empty())),
604 yield_requested: CancellationToken::new(),
605 },
606 cancellation,
607 )
608 .await
609 }
610
611 pub async fn write_and_poll_with_options(
614 &self,
615 run_id: &RunId,
616 session_id: ExecSessionId,
617 input: Option<&str>,
618 options: ExecWaitOptions,
619 cancellation: &CancellationToken,
620 ) -> Result<ExecPollResult, ExecProcessError> {
621 if options.duration.is_zero() {
622 return Err(ExecProcessError::Invalid(
623 "exec poll duration must be positive".to_owned(),
624 ));
625 }
626 let session = self.session(run_id, session_id)?;
627 let requested_input = input.filter(|input| !input.is_empty());
628 let input_to_deliver = if requested_input.is_some() {
629 match session.lifecycle.status()? {
630 ExecSessionStatus::Running => requested_input,
631 ExecSessionStatus::Exited { .. } => None,
636 ExecSessionStatus::Terminated => return Err(ExecProcessError::Cancelled),
637 ExecSessionStatus::Failed { message } => return Err(ExecProcessError::Io(message)),
638 }
639 } else {
640 None
641 };
642 let result = match session.process.clone() {
643 ManagedProcess::Pipe(process) => {
644 if let Some(input) = input_to_deliver {
645 if let Err(error) = process.send(input).await {
646 if !matches!(
647 session.lifecycle.status()?,
648 ExecSessionStatus::Exited { .. }
649 ) {
650 return Err(error);
651 }
652 }
653 }
654 process
655 .poll(&session.lifecycle, session.started, &options, cancellation)
656 .await?
657 }
658 ManagedProcess::Pty { process_id } => {
659 if let Some(input) = input_to_deliver {
660 let pty = self.pty.clone();
661 let run_id = run_id.clone();
662 let process_id = process_id.clone();
663 let input = input.to_owned();
664 tokio::task::spawn_blocking(move || pty.send(&run_id, &process_id, &input))
665 .await
666 .map_err(|error| ExecProcessError::Io(error.to_string()))?
667 .map_err(|error| ExecProcessError::Io(error.to_string()))?;
668 }
669 let pty = self.pty.clone();
670 let run_id = run_id.clone();
671 let process_id = process_id.clone();
672 let cancellation = cancellation.clone();
673 let read = tokio::task::spawn_blocking(move || {
674 pty.read_with_options(
675 &run_id,
676 &process_id,
677 PtyReadOptions {
678 timeout: options.duration,
679 settle: match options.mode {
680 ExecWaitMode::Output => {
681 Duration::from_millis(50).min(options.duration)
682 }
683 ExecWaitMode::Completion => options.duration,
684 },
685 yield_requested: options.yield_requested,
686 },
687 &cancellation,
688 )
689 })
690 .await
691 .map_err(|error| ExecProcessError::Io(error.to_string()))?
692 .map_err(|error| match error {
693 crate::pty_process::PtyProcessError::Cancelled => ExecProcessError::Cancelled,
694 error => ExecProcessError::Io(error.to_string()),
695 })?;
696 ExecPollResult {
697 stdout: read.output,
698 stderr: String::new(),
699 dropped_bytes: read.dropped_bytes,
700 alive: read.alive,
701 exit_code: read.exit_code,
702 wall_time_seconds: session.started.elapsed().as_secs_f64(),
703 }
704 }
705 };
706 if !result.alive {
707 if let Some(exit_code) = result.exit_code {
708 transition_session(
709 &self.events,
710 &(run_id.clone(), session_id),
711 &session,
712 ExecSessionStatus::Exited { exit_code },
713 )?;
714 }
715 self.remove_finished(run_id, session_id, session).await;
716 }
717 Ok(result)
718 }
719
720 pub async fn close(
721 &self,
722 run_id: &RunId,
723 session_id: ExecSessionId,
724 ) -> Result<(), ExecProcessError> {
725 let session = self
726 .sessions
727 .lock()
728 .map_err(|_| ExecProcessError::Unavailable)?
729 .remove(&(run_id.clone(), session_id))
730 .ok_or(ExecProcessError::NotFound(session_id.get()))?;
731 transition_session(
732 &self.events,
733 &(run_id.clone(), session_id),
734 &session,
735 ExecSessionStatus::Terminated,
736 )?;
737 self.terminate_session(run_id, session).await;
738 Ok(())
739 }
740
741 pub async fn close_run(&self, run_id: &RunId) -> Result<usize, ExecProcessError> {
742 let owned = {
743 let mut sessions = self
744 .sessions
745 .lock()
746 .map_err(|_| ExecProcessError::Unavailable)?;
747 let keys = sessions
748 .keys()
749 .filter(|(owner, _)| owner == run_id)
750 .cloned()
751 .collect::<Vec<_>>();
752 keys.into_iter()
753 .filter_map(|key| sessions.remove(&key).map(|session| (key, session)))
754 .collect::<Vec<_>>()
755 };
756 let count = owned.len();
757 for (key, session) in owned {
758 transition_session(&self.events, &key, &session, ExecSessionStatus::Terminated)?;
759 self.terminate_session(run_id, session).await;
760 }
761 self.runtime_temps
762 .lock()
763 .map_err(|_| ExecProcessError::Unavailable)?
764 .remove(run_id);
765 Ok(count)
766 }
767
768 pub fn list(&self, run_id: &RunId) -> Result<Vec<ExecSessionId>, ExecProcessError> {
769 let sessions = self
770 .sessions
771 .lock()
772 .map_err(|_| ExecProcessError::Unavailable)?;
773 let mut active = Vec::new();
774 for ((owner, session_id), session) in sessions.iter() {
775 if owner == run_id && session.lifecycle.status()? == ExecSessionStatus::Running {
776 active.push(*session_id);
777 }
778 }
779 Ok(active)
780 }
781
782 pub fn snapshot(
783 &self,
784 run_id: &RunId,
785 session_id: ExecSessionId,
786 ) -> Result<ExecSessionSnapshot, ExecProcessError> {
787 let session = self.session(run_id, session_id)?;
788 session_snapshot(&(run_id.clone(), session_id), &session)
789 }
790
791 pub fn operation_plan(
792 &self,
793 run_id: &RunId,
794 session_id: ExecSessionId,
795 ) -> Result<ToolOperationPlan, ExecProcessError> {
796 Ok(self.session(run_id, session_id)?.operation)
797 }
798
799 fn session(
800 &self,
801 run_id: &RunId,
802 session_id: ExecSessionId,
803 ) -> Result<ManagedSession, ExecProcessError> {
804 self.sessions
805 .lock()
806 .map_err(|_| ExecProcessError::Unavailable)?
807 .get(&(run_id.clone(), session_id))
808 .cloned()
809 .ok_or(ExecProcessError::NotFound(session_id.get()))
810 }
811
812 async fn remove_finished(
813 &self,
814 run_id: &RunId,
815 session_id: ExecSessionId,
816 session: ManagedSession,
817 ) {
818 if let Ok(mut sessions) = self.sessions.lock() {
819 sessions.remove(&(run_id.clone(), session_id));
820 }
821 if let ManagedProcess::Pty { process_id } = session.process {
822 let pty = self.pty.clone();
823 let run_id = run_id.clone();
824 let _ = tokio::task::spawn_blocking(move || pty.close(&run_id, &process_id)).await;
825 }
826 }
827
828 async fn terminate_session(&self, run_id: &RunId, session: ManagedSession) {
829 match session {
830 ManagedSession {
831 process: ManagedProcess::Pipe(process),
832 ..
833 } => process.terminate().await,
834 ManagedSession {
835 process: ManagedProcess::Pty { process_id },
836 ..
837 } => {
838 let pty = self.pty.clone();
839 let run_id = run_id.clone();
840 let _ = tokio::task::spawn_blocking(move || pty.close(&run_id, &process_id)).await;
841 }
842 }
843 }
844}
845
846fn session_snapshot(
847 (run_id, session_id): &SessionKey,
848 session: &ManagedSession,
849) -> Result<ExecSessionSnapshot, ExecProcessError> {
850 Ok(ExecSessionSnapshot {
851 run_id: run_id.clone(),
852 session_id: *session_id,
853 tty: session.tty,
854 status: session.lifecycle.status()?,
855 operation: session.operation.clone(),
856 wall_time_seconds: session.started.elapsed().as_secs_f64(),
857 })
858}
859
860fn publish_session_event(
861 events: &broadcast::Sender<ExecSessionEvent>,
862 key: &SessionKey,
863 session: &ManagedSession,
864) {
865 if let Ok(snapshot) = session_snapshot(key, session) {
866 let _ = events.send(ExecSessionEvent { snapshot });
867 }
868}
869
870fn transition_session(
871 events: &broadcast::Sender<ExecSessionEvent>,
872 key: &SessionKey,
873 session: &ManagedSession,
874 status: ExecSessionStatus,
875) -> Result<bool, ExecProcessError> {
876 let changed = session.lifecycle.transition(status)?;
877 if changed {
878 publish_session_event(events, key, session);
879 }
880 Ok(changed)
881}
882
883fn spawn_exit_watcher(
884 key: SessionKey,
885 session: ManagedSession,
886 pty: Arc<PtyProcessManager>,
887 events: broadcast::Sender<ExecSessionEvent>,
888) {
889 tokio::spawn(async move {
890 loop {
891 if session
892 .lifecycle
893 .status()
894 .is_ok_and(|status| status.is_terminal())
895 {
896 return;
897 }
898 let observed = match &session.process {
899 ManagedProcess::Pipe(process) => process
900 .child
901 .lock()
902 .await
903 .try_wait()
904 .map(|status| status.map(|status| exit_status_code(&status)))
905 .map_err(|error| error.to_string()),
906 ManagedProcess::Pty { process_id } => {
907 let pty = pty.clone();
908 let run_id = key.0.clone();
909 let process_id = process_id.clone();
910 match tokio::task::spawn_blocking(move || pty.status(&run_id, &process_id))
911 .await
912 {
913 Ok(result) => result.map_err(|error| error.to_string()),
914 Err(error) => Err(error.to_string()),
915 }
916 }
917 };
918 match observed {
919 Ok(Some(exit_code)) => {
920 let _ = transition_session(
921 &events,
922 &key,
923 &session,
924 ExecSessionStatus::Exited { exit_code },
925 );
926 if let ManagedProcess::Pipe(process) = &session.process {
927 process.stdin.lock().await.take();
928 }
929 return;
930 }
931 Ok(None) => {}
932 Err(message) => {
933 let _ = transition_session(
934 &events,
935 &key,
936 &session,
937 ExecSessionStatus::Failed { message },
938 );
939 match &session.process {
940 ManagedProcess::Pipe(process) => process.terminate().await,
941 ManagedProcess::Pty { process_id } => {
942 let pty = pty.clone();
943 let run_id = key.0.clone();
944 let process_id = process_id.clone();
945 let _ = tokio::task::spawn_blocking(move || {
946 pty.close(&run_id, &process_id)
947 })
948 .await;
949 }
950 }
951 return;
952 }
953 }
954 tokio::time::sleep(PROCESS_WATCH_INTERVAL).await;
955 }
956 });
957}
958
959fn exit_status_code(status: &ExitStatus) -> i32 {
960 status.code().unwrap_or(-1)
961}
962
963#[cfg(unix)]
964fn isolate_process_group(command: &mut Command, backend_starts_new_session: bool) {
965 if !backend_starts_new_session {
966 command.process_group(0);
967 }
968}
969
970#[cfg(not(unix))]
971fn isolate_process_group(_command: &mut Command, _backend_starts_new_session: bool) {}
972
973#[cfg(unix)]
974fn terminate_process_group(process_group_id: Option<u32>) {
975 if let Some(process_group_id) = process_group_id.filter(|id| *id <= i32::MAX as u32) {
976 unsafe {
978 libc::kill(-(process_group_id as i32), libc::SIGKILL);
979 }
980 }
981}
982
983#[cfg(not(unix))]
984fn terminate_process_group(_process_group_id: Option<u32>) {}