Skip to main content

orchestral_runtime/
exec_process.rs

1//! Run-scoped process sessions behind the model-visible unified exec Tools.
2//!
3//! Pipe and PTY are execution details. Both are addressed by one integer
4//! session ID and remain strictly scoped to the owning Agent Run.
5
6mod 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    /// Exact Host-derived authority under which subsequent input executes.
59    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    /// The session still needs observation, including final pipe drainage
68    /// after its direct process has exited.
69    pub alive: bool,
70    pub exit_code: Option<i32>,
71    pub wall_time_seconds: f64,
72}
73
74/// Determines whether ordinary process output ends a wait early.
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76pub enum ExecWaitMode {
77    /// Return after a short pause in output, suitable for interactive prompts.
78    Output,
79    /// Aggregate output until exit, the wait deadline, or a Host yield request.
80    Completion,
81}
82
83impl ExecWaitMode {
84    /// Preserve prompt responsiveness for TTY sessions and delivered input.
85    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/// One observation window. Ending this window never terminates the process.
95#[derive(Debug, Clone)]
96pub struct ExecWaitOptions {
97    pub duration: Duration,
98    pub mode: ExecWaitMode,
99    /// Host request to return the current observation without cancelling work.
100    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    // Keep scratch files until this process and its exit watcher release them.
134    _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            // Register before inspection: EOF between inspecting and awaiting
207            // must not leave the completion waiter asleep forever.
208            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        // Reader tasks must not outlive their owning session and retain pipe
342        // handles when the supervisor itself is dropped.
343        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            // Exit alone is not a terminal observation: stdout/stderr
390            // readers may still own bytes that have not reached the buffers.
391            // Keep this session addressable across deadlines and yields.
392            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
442/// Run-scoped owner and observer for pipe and PTY execution resources.
443pub 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    /// Use a stable, private Host directory so execution policy identities can
459    /// survive Host restart. The root must be outside all workspace roots.
460    /// Its parent must already exist; Run children are exclusively created.
461    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    /// Host policy and exec Tool restrictions must explicitly grant read/write
492    /// access to this root. Each dispatched sandbox receives only its Run child.
493    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    /// Use output waits for TTY/input interactions and completion waits for
587    /// empty pipe observations. The duration bounds this observation only.
588    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    /// Observe one existing session with an explicit wait strategy. Yielding
612    /// preserves both the process and its Run-scoped authority.
613    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                // A process can exit after exec_command reports it as alive but
632                // before the model's follow-up write arrives. Preserve the
633                // terminal observation instead of turning that normal race into
634                // a Tool failure.
635                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        // SAFETY: the child is the leader of a fresh process group/session.
977        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>) {}