1479 lines · 56.4 KB · Rust
RawDownload
//! Process execution & management methods + supporting types.//!//! Ported from `packages/core/src/agent-os.ts` (process methods) and `runtime-compat.ts`//! (`ExecOptions`, `ExecResult`, `ProcessInfo`, etc.).//!//! Two distinct process views: SDK-spawned processes (`processes` map, keyed by user-facing pid)//! back `spawn` + the stdin/stdout/stderr/exit subscriptions + `wait/list/get/stop/kill`; the kernel//! process table backs `exec`, `all_processes`, `process_tree`. use std::collections::BTreeMap;use std::sync::atomic::Ordering; use anyhow::{Context, Result};use scc::HashMap as SccHashMap;use serde::{Deserialize, Serialize};use tokio::sync::{broadcast, watch};use tokio::task::JoinHandle; use agentos_sidecar_client::wire::{self, EventPayload, ProcessSnapshotStatus, StreamChannel}; use crate::agent_os::{AgentOs, ProcessEntry};use crate::command_line::resolve_exec_command;use crate::error::ClientError;use crate::stream::Subscription; /// Broadcast channel capacity for a spawned process's stdout/stderr fan-out.const PROCESS_STREAM_CAPACITY: usize = 1024; /// Maximum SDK-spawned process entries retained per VM.const PROCESS_REGISTRY_LIMIT: usize = 1024; /// Maximum first-observed process timestamp entries retained per VM.const OBSERVED_PROCESS_TIME_LIMIT: usize = 4096; /// Maximum bytes captured by `exec` across stdout and stderr.const EXEC_OUTPUT_CAPTURE_LIMIT_BYTES: usize = 16 * 1024 * 1024; /// Default guest working directory for `exec`/`spawn`, matching the TS sidecar client.pub(crate) const DEFAULT_EXEC_CWD: &str = "/workspace"; /// Base value for the synthetic display-pid sequence used by `spawn` (TS `SYNTHETIC_PID_BASE`). The/// first spawned process is assigned exactly this value.pub(crate) const SYNTHETIC_PID_BASE: u64 = 1_000_000; // ---------------------------------------------------------------------------// Supporting types// --------------------------------------------------------------------------- /// Timing-mitigation mode for an execution.#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]#[serde(rename_all = "lowercase")]pub enum TimingMitigation {    #[default]    Off,    Freeze,} /// `stdin` value: a string or raw bytes.#[derive(Debug, Clone, PartialEq, Eq)]pub enum StdinInput {    Text(String),    Bytes(Vec<u8>),} /// A raw-byte streaming callback for stdout/stderr (TS `(data: Uint8Array) => void`). Invoked once/// per output chunk as it arrives. Never assume UTF-8: chunks are delivered as raw bytes.pub type OutputCallback = Box<dyn FnMut(&[u8]) + Send>; /// Base options shared by `exec` and `spawn`.////// `on_stdout`/`on_stderr` mirror the TS `ExecOptions.onStdout`/`onStderr` raw-byte streaming/// callbacks. For `exec` they fire for the duration of the call; for `spawn` they are seeded into the/// stdout/stderr fan-out at spawn time (matching the TS initial-handler-set behavior).pub struct ExecOptions {    pub env: BTreeMap<String, String>,    pub cwd: Option<String>,    pub stdin: Option<StdinInput>,    pub timeout: Option<f64>,    pub on_stdout: Option<OutputCallback>,    pub on_stderr: Option<OutputCallback>,    pub capture_stdio: Option<bool>,    pub file_path: Option<String>,    pub cpu_time_limit_ms: Option<f64>,    pub timing_mitigation: Option<TimingMitigation>,} impl Default for ExecOptions {    fn default() -> Self {        Self {            env: BTreeMap::new(),            cwd: Some(DEFAULT_EXEC_CWD.to_string()),            stdin: None,            timeout: None,            on_stdout: None,            on_stderr: None,            capture_stdio: None,            file_path: None,            cpu_time_limit_ms: None,            timing_mitigation: None,        }    }} /// Result of `exec`.#[derive(Debug, Clone, PartialEq, Eq)]pub struct ExecResult {    pub exit_code: i32,    pub stdout: String,    pub stderr: String,} /// `stdio` mode for a spawn.#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]#[serde(rename_all = "lowercase")]pub enum SpawnStdio {    #[default]    Pipe,    Inherit,} /// Callback-free options for portable `spawn`.#[derive(Default)]pub struct SpawnOptions {    pub env: BTreeMap<String, String>,    pub cwd: Option<String>,    pub stdio: Option<SpawnStdio>,    pub stdin_fd: Option<i32>,    pub stdout_fd: Option<i32>,    pub stderr_fd: Option<i32>,    pub stream_stdin: Option<bool>,} #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]#[serde(rename_all = "lowercase")]pub enum ProcessStream {    Stdout,    Stderr,} #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]pub struct ProcessOutput {    pub pid: u32,    pub stream: ProcessStream,    pub data: Vec<u8>,} #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]pub struct ProcessExit {    pub pid: u32,    #[serde(rename = "exitCode")]    pub exit_code: i32,} /// Public JSON info for SDK-spawned processes.#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]pub struct SpawnedProcessInfo {    pub pid: u32,    pub command: String,    pub args: Vec<String>,    pub running: bool,    #[serde(rename = "exitCode")]    pub exit_code: Option<i32>,    /// Epoch milliseconds when `spawn` registered the process.    #[serde(rename = "startedAt")]    pub started_at: i64,} /// The pid returned by `spawn`.#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]pub struct SpawnHandle {    pub pid: u32,} /// Process status from the kernel process table.#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]#[serde(rename_all = "lowercase")]pub enum ProcessStatus {    Running,    Exited,} /// Full kernel process info (TS `KernelProcessInfo`).#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]pub struct ProcessInfo {    pub pid: u32,    pub ppid: u32,    pub pgid: u32,    pub sid: u32,    pub driver: String,    pub command: String,    pub args: Vec<String>,    pub cwd: String,    pub status: ProcessStatus,    #[serde(rename = "exitCode")]    pub exit_code: Option<i32>,    #[serde(rename = "startTime")]    pub start_time: f64,    #[serde(rename = "exitTime")]    pub exit_time: Option<f64>,} /// A node in the process forest (`ProcessInfo` + children).#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]pub struct ProcessTreeNode {    #[serde(flatten)]    pub info: ProcessInfo,    pub children: Vec<ProcessTreeNode>,} // ---------------------------------------------------------------------------// Methods// --------------------------------------------------------------------------- impl AgentOs {    /// Run a command to completion. The wire `Execute` request starts the process and returns a    /// process id immediately; stdout/stderr are accumulated and the call resolves once the matching    /// `ProcessExited` event arrives. This mirrors the TS pass-through to `kernel.exec` semantically:    /// the result is the full captured stdout/stderr plus exit code.    pub async fn exec(&self, command: &str, options: ExecOptions) -> Result<ExecResult> {        // Parse the command line into a `(command, args)` pair the same way the sidecar's        // child_process path does: shell-free argv lists spawn directly (preserving the command's        // real exit code), while shell syntax or a builtin head runs under `sh -c <line>`.        let (resolved_command, resolved_args) = resolve_exec_command(command)?;        self.exec_argv(&resolved_command, &resolved_args, options)            .await    }     /// Run a command to completion from an already-structured `(command, args)` argv, bypassing the    /// `exec` command-line parser. Each `args` element is sent verbatim as a distinct argv element —    /// no whitespace re-splitting, no shell metacharacter detection, and no routing through    /// `sh -c`. Callers that already hold a structured argv (for example the cron `Exec` action)    /// must use this so the structured-argv contract is preserved end to end.    pub async fn exec_argv(        &self,        command: &str,        args: &[String],        mut options: ExecOptions,    ) -> Result<ExecResult> {        let process_id = self.next_process_id();         // Subscribe to events BEFORE issuing the request so no output/exit is missed between the        // request landing and the subscription being installed.        let mut events = self.transport().subscribe_wire_events();         let resolved_command = command.to_owned();        let resolved_args = args.to_vec();        let started = self            .send_execute(                &process_id,                Some(resolved_command),                resolved_args,                options.env.clone(),                options.cwd.clone(),            )            .await            .context("exec: Execute request failed")?;        debug_assert_eq!(started.process_id, process_id);         // Deliver any provided stdin, then close stdin so a non-interactive run observes EOF. This        // mirrors the TS `runAndCapture` path (`proc.writeStdin(options.stdin); proc.closeStdin()`).        if let Some(stdin) = options.stdin.take() {            let chunk = stdin_to_bytes(stdin);            let ownership = self.vm_scope();            let _ = self                .transport()                .request_wire(                    ownership,                    wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {                        process_id: process_id.clone(),                        chunk,                    }),                )                .await;        }        {            let ownership = self.vm_scope();            let _ = self                .transport()                .request_wire(                    ownership,                    wire::RequestPayload::CloseStdinRequest(wire::CloseStdinRequest {                        process_id: process_id.clone(),                    }),                )                .await;        }         let mut on_stdout = options.on_stdout.take();        let mut on_stderr = options.on_stderr.take();         // A `timeout` (ms) bounds the run: when it elapses, SIGKILL the process and keep draining        // until the exit event lands. This mirrors the TS `runAndCapture` timeout race that kills the        // process and then awaits its exit code.        let timeout_deadline = options            .timeout            .filter(|ms| ms.is_finite() && *ms >= 0.0)            .map(|ms| {                tokio::time::Instant::now() + std::time::Duration::from_secs_f64(ms / 1000.0)            });        let mut killed_for_timeout = false;         let capture_stdio = options.capture_stdio.unwrap_or(true);        let mut stdout = Vec::<u8>::new();        let mut stderr = Vec::<u8>::new();        let mut captured_output_bytes = 0usize;        let mut capture_error: Option<ClientError> = None;        let exit_code = loop {            let recv = events.recv();            let frame = match timeout_deadline {                Some(deadline) => {                    tokio::select! {                        result = recv => result,                        _ = tokio::time::sleep_until(deadline), if !killed_for_timeout => {                            killed_for_timeout = true;                            self.kill_wire_process(&process_id, "SIGKILL");                            continue;                        }                    }                }                None => recv.await,            };            let (_, payload) = match frame {                Ok(frame) => frame,                Err(broadcast::error::RecvError::Lagged(_)) => continue,                Err(broadcast::error::RecvError::Closed) => {                    return Err(ClientError::Sidecar(                        "exec: event stream closed before process exit".to_owned(),                    )                    .into());                }            };            match payload {                EventPayload::ProcessOutputEvent(output) if output.process_id == process_id => {                    match output.channel {                        StreamChannel::Stdout => {                            if let Some(cb) = on_stdout.as_mut() {                                cb(&output.chunk);                            }                            if capture_stdio && capture_error.is_none() {                                match append_exec_output(                                    &mut stdout,                                    &output.chunk,                                    &mut captured_output_bytes,                                    "stdout",                                ) {                                    Ok(()) => {}                                    Err(error) => {                                        self.kill_wire_process(&process_id, "SIGKILL");                                        capture_error = Some(error);                                    }                                }                            }                        }                        StreamChannel::Stderr => {                            if let Some(cb) = on_stderr.as_mut() {                                cb(&output.chunk);                            }                            if capture_stdio && capture_error.is_none() {                                match append_exec_output(                                    &mut stderr,                                    &output.chunk,                                    &mut captured_output_bytes,                                    "stderr",                                ) {                                    Ok(()) => {}                                    Err(error) => {                                        self.kill_wire_process(&process_id, "SIGKILL");                                        capture_error = Some(error);                                    }                                }                            }                        }                    }                }                EventPayload::ProcessExitedEvent(exited) if exited.process_id == process_id => {                    break exited.exit_code;                }                EventPayload::ProcessOutputEvent(_)                | EventPayload::ProcessExitedEvent(_)                | EventPayload::VmLifecycleEvent(_)                | EventPayload::StructuredEvent(_)                | EventPayload::ExtEnvelope(_) => {}            }        };         if let Some(error) = capture_error {            return Err(error.into());        }         Ok(ExecResult {            exit_code,            stdout: String::from_utf8_lossy(&stdout).into_owned(),            stderr: String::from_utf8_lossy(&stderr).into_owned(),        })    }     /// Spawn a process. SYNC; returns `{ pid }` only. Installs stdout/stderr fan-out over broadcast    /// channels and wires exit via a background event-pump task. The user-facing `pid` is the    /// SDK-allocated map key (the wire `process_id` is held inside the [`ProcessEntry`]).    pub fn spawn(        &self,        command: &str,        args: Vec<String>,        options: SpawnOptions,    ) -> Result<SpawnHandle> {        let registry_guard = self.inner().process_registry_lock.lock();        self.prune_exited_processes_locked(1);        if self.process_registry_len_locked() >= PROCESS_REGISTRY_LIMIT {            return Err(ClientError::Sidecar(format!(                "process registry limit exceeded: at most {PROCESS_REGISTRY_LIMIT} processes can be tracked per VM"            ))            .into());        }         // Draw the public pid from the dedicated synthetic-pid space (TS `nextSyntheticPid`), seeded        // at `SYNTHETIC_PID_BASE`. `exec` uses a separate counter so it never perturbs this sequence.        let pid = self            .inner()            .synthetic_pid_counter            .fetch_add(1, Ordering::SeqCst) as u32;        let process_id = format!("proc-{pid}-{}", uuid::Uuid::new_v4());         let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(PROCESS_STREAM_CAPACITY);        let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(PROCESS_STREAM_CAPACITY);        let (output_tx, _) = broadcast::channel::<ProcessOutput>(PROCESS_STREAM_CAPACITY);        // Seeded `None`; the already-exited branch of `on_process_exit` fires immediately once this        // watch holds `Some(code)`.        let (exit_tx, _) = watch::channel::<Option<i32>>(None);        // Seeded `None`; filled with the kernel pid once the `Execute` response lands so        // `all_processes`/`process_tree` can remap the kernel snapshot back to this display pid.        let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);         let entry = ProcessEntry {            command: command.to_owned(),            args: args.clone(),            stdout_tx: stdout_tx.clone(),            stderr_tx: stderr_tx.clone(),            output_tx: output_tx.clone(),            exit_tx: exit_tx.clone(),            process_id: process_id.clone(),            kernel_pid: kernel_pid_tx.clone(),            output_tasks: Vec::new(),            started_at: epoch_ms_now() as i64,        };        // `spawn` is documented as overwriting any prior entry for a freshly allocated pid; the pid        // is monotonic so a collision is not expected.        let _ = self.inner().processes.insert(pid, entry);        drop(registry_guard);         // Subscribe to events before issuing the request so the pump sees everything.        let events = self.transport().subscribe_wire_events();         let this = self.clone();        let command = command.to_owned();        tokio::spawn(async move {            this.run_spawn(                pid,                process_id,                command,                args,                options,                events,                stdout_tx,                stderr_tx,                output_tx,                exit_tx,                kernel_pid_tx,            )            .await;        });         Ok(SpawnHandle { pid })    }     /// Write to a spawned process's stdin. SYNC. Errors with `ProcessNotFound`.    pub fn write_process_stdin(        &self,        pid: u32,        data: StdinInput,    ) -> std::result::Result<(), ClientError> {        let process_id = self.lookup_process_id(pid)?;        let chunk: Vec<u8> = stdin_to_bytes(data);        let this = self.clone();        // Fire-and-forget: the TS API is synchronous and does not surface a write error.        tokio::spawn(async move {            let ownership = this.vm_scope();            let _ = this                .transport()                .request_wire(                    ownership,                    wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {                        process_id,                        chunk,                    }),                )                .await;        });        Ok(())    } 

Lines 1–500 of 1479.

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t