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::{ByteStream, Subscription};
const PROCESS_STREAM_CAPACITY: usize = 1024;
const PROCESS_REGISTRY_LIMIT: usize = 1024;
const OBSERVED_PROCESS_TIME_LIMIT: usize = 4096;
const EXEC_OUTPUT_CAPTURE_LIMIT_BYTES: usize = 16 * 1024 * 1024;
pub(crate) const DEFAULT_EXEC_CWD: &str = "/workspace";
pub(crate) const SYNTHETIC_PID_BASE: u64 = 1_000_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum TimingMitigation {
#[default]
Off,
Freeze,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StdinInput {
Text(String),
Bytes(Vec<u8>),
}
pub type OutputCallback = Box<dyn FnMut(&[u8]) + Send>;
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,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExecResult {
pub exit_code: i32,
pub stdout: String,
pub stderr: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum SpawnStdio {
#[default]
Pipe,
Inherit,
}
#[derive(Default)]
pub struct SpawnOptions {
pub base: ExecOptions,
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)]
pub struct SpawnedProcessInfo {
pub pid: u32,
pub command: String,
pub args: Vec<String>,
pub running: bool,
#[serde(rename = "exitCode")]
pub exit_code: Option<i32>,
#[serde(rename = "startedAt")]
pub started_at: i64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct SpawnHandle {
pub pid: u32,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ProcessStatus {
Running,
Exited,
}
#[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>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ProcessTreeNode {
#[serde(flatten)]
pub info: ProcessInfo,
pub children: Vec<ProcessTreeNode>,
}
impl AgentOs {
pub async fn exec(&self, command: &str, options: ExecOptions) -> Result<ExecResult> {
let (resolved_command, resolved_args) = resolve_exec_command(command)?;
self.exec_argv(&resolved_command, &resolved_args, options)
.await
}
pub async fn exec_argv(
&self,
command: &str,
args: &[String],
mut options: ExecOptions,
) -> Result<ExecResult> {
let process_id = self.next_process_id();
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);
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();
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(),
})
}
pub fn spawn(
&self,
command: &str,
args: Vec<String>,
mut 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());
}
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 (exit_tx, _) = watch::channel::<Option<i32>>(None);
let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
let mut output_tasks = Vec::new();
if let Some(cb) = options.base.on_stdout.take() {
output_tasks.push(install_output_callback(stdout_tx.clone(), cb));
}
if let Some(cb) = options.base.on_stderr.take() {
output_tasks.push(install_output_callback(stderr_tx.clone(), cb));
}
let entry = ProcessEntry {
command: command.to_owned(),
args: args.clone(),
stdout_tx: stdout_tx.clone(),
stderr_tx: stderr_tx.clone(),
exit_tx: exit_tx.clone(),
process_id: process_id.clone(),
kernel_pid: kernel_pid_tx.clone(),
output_tasks,
started_at: epoch_ms_now() as i64,
};
let _ = self.inner().processes.insert(pid, entry);
drop(registry_guard);
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,
exit_tx,
kernel_pid_tx,
)
.await;
});
Ok(SpawnHandle { pid })
}
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();
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(())
}
pub fn close_process_stdin(&self, pid: u32) -> std::result::Result<(), ClientError> {
let process_id = self.lookup_process_id(pid)?;
let this = self.clone();
tokio::spawn(async move {
let ownership = this.vm_scope();
let _ = this
.transport()
.request_wire(
ownership,
wire::RequestPayload::CloseStdinRequest(wire::CloseStdinRequest { process_id }),
)
.await;
});
Ok(())
}
pub fn on_process_stdout(&self, pid: u32) -> std::result::Result<ByteStream, ClientError> {
let rx = self
.inner()
.processes
.read(&pid, |_, entry| entry.stdout_tx.subscribe())
.ok_or(ClientError::ProcessNotFound(pid))?;
Ok(ByteStream::new(rx))
}
pub fn on_process_stderr(&self, pid: u32) -> std::result::Result<ByteStream, ClientError> {
let rx = self
.inner()
.processes
.read(&pid, |_, entry| entry.stderr_tx.subscribe())
.ok_or(ClientError::ProcessNotFound(pid))?;
Ok(ByteStream::new(rx))
}
pub fn on_process_exit(
&self,
pid: u32,
handler: impl FnOnce(i32) + Send + 'static,
) -> std::result::Result<Subscription, ClientError> {
let mut rx = self
.inner()
.processes
.read(&pid, |_, entry| entry.exit_tx.subscribe())
.ok_or(ClientError::ProcessNotFound(pid))?;
if let Some(code) = *rx.borrow() {
handler(code);
return Ok(Subscription::noop());
}
let task = tokio::spawn(async move {
while rx.changed().await.is_ok() {
if let Some(code) = *rx.borrow() {
handler(code);
return;
}
}
});
Ok(Subscription::new(move || task.abort()))
}
pub async fn wait_process(&self, pid: u32) -> std::result::Result<i32, ClientError> {
let mut rx = self
.inner()
.processes
.read(&pid, |_, entry| entry.exit_tx.subscribe())
.ok_or(ClientError::ProcessNotFound(pid))?;
if let Some(code) = *rx.borrow() {
return Ok(code);
}
while rx.changed().await.is_ok() {
if let Some(code) = *rx.borrow() {
return Ok(code);
}
}
Err(ClientError::Sidecar(format!(
"wait_process: exit channel closed before process {pid} reported an exit code"
)))
}
pub fn list_processes(&self) -> Vec<SpawnedProcessInfo> {
let mut out = Vec::new();
self.inner().processes.scan(|pid, entry| {
let exit_code = *entry.exit_tx.borrow();
out.push(SpawnedProcessInfo {
pid: *pid,
command: entry.command.clone(),
args: entry.args.clone(),
running: exit_code.is_none(),
exit_code,
started_at: entry.started_at,
});
});
out
}
pub async fn all_processes(&self) -> Result<Vec<ProcessInfo>> {
let ownership = self.vm_scope();
let response = self
.transport()
.request_wire(ownership, wire::RequestPayload::GetProcessSnapshotRequest)
.await
.context("all_processes: GetProcessSnapshot request failed")?;
let snapshot = match response {
wire::ResponsePayload::ProcessSnapshotResponse(snapshot) => snapshot,
wire::ResponsePayload::RejectedResponse(wire::RejectedResponse { code, message }) => {
return Err(ClientError::Kernel { code, message }.into());
}
other => {
return Err(ClientError::Sidecar(format!(
"all_processes: unexpected response {other:?}"
))
.into());
}
};
struct Tracked {
exit_code: Option<i32>,
command: String,
args: Vec<String>,
}
let mut tracked_by_process_id: BTreeMap<String, Tracked> = BTreeMap::new();
let mut display_pid_by_kernel_pid: BTreeMap<u32, u32> = BTreeMap::new();
self.inner().processes.scan(|display_pid, entry| {
let exit_code = *entry.exit_tx.borrow();
if let Some(kernel_pid) = *entry.kernel_pid.borrow() {
display_pid_by_kernel_pid.insert(kernel_pid, *display_pid);
}
tracked_by_process_id.insert(
entry.process_id.clone(),
Tracked {
exit_code,
command: entry.command.clone(),
args: entry.args.clone(),
},
);
});
let now_ms = epoch_ms_now();
let mut seen_display_pids: std::collections::BTreeSet<u32> =
std::collections::BTreeSet::new();
let mut out: Vec<ProcessInfo> = Vec::new();
for entry in snapshot.processes {
let tracked = tracked_by_process_id.get(&entry.process_id);
let display_pid = display_pid_by_kernel_pid
.get(&entry.pid)
.copied()
.unwrap_or(entry.pid);
let display_ppid = display_pid_by_kernel_pid
.get(&entry.ppid)
.copied()
.unwrap_or(entry.ppid);
let display_pgid = display_pid_by_kernel_pid
.get(&entry.pgid)
.copied()
.unwrap_or(entry.pgid);
let display_sid = display_pid_by_kernel_pid
.get(&entry.sid)
.copied()
.unwrap_or(entry.sid);
let process_key = format!("{}:{}", entry.process_id, entry.pid);
let start_time = self.observed_start_time(&process_key, now_ms);
let (status, exit_code) = match tracked {
Some(t) => match t.exit_code {
Some(code) => (ProcessStatus::Exited, Some(code)),
None => (ProcessStatus::Running, entry.exit_code),
},
None => {
let status = match entry.status {
ProcessSnapshotStatus::Running | ProcessSnapshotStatus::Stopped => {
ProcessStatus::Running
}
ProcessSnapshotStatus::Exited => ProcessStatus::Exited,
};
(status, entry.exit_code)
}
};
let exit_time = match (tracked, status) {
(Some(_), ProcessStatus::Exited) => {
Some(self.observed_exit_time(&entry.process_id, now_ms))
}
_ => None,
};
let (command, args) = match tracked {
Some(t) => (t.command.clone(), t.args.clone()),
None => (entry.command, entry.args),
};
seen_display_pids.insert(display_pid);
out.push(ProcessInfo {
pid: display_pid,
ppid: display_ppid,
pgid: display_pgid,
sid: display_sid,
driver: entry.driver,
command,
args,
cwd: entry.cwd,
status,
exit_code,
start_time,
exit_time,
});
}
self.inner().processes.scan(|display_pid, entry| {
if seen_display_pids.contains(display_pid) {
return;
}
let exit_code = *entry.exit_tx.borrow();
let process_key = format!("{}:{}", entry.process_id, display_pid);
let start_time = self.observed_start_time(&process_key, now_ms);
let (status, exit_time) = match exit_code {
Some(_) => (
ProcessStatus::Exited,
Some(self.observed_exit_time(&entry.process_id, now_ms)),
),
None => (ProcessStatus::Running, None),
};
out.push(ProcessInfo {
pid: *display_pid,
ppid: 0,
pgid: *display_pid,
sid: *display_pid,
driver: String::new(),
command: entry.command.clone(),
args: entry.args.clone(),
cwd: String::new(),
status,
exit_code,
start_time,
exit_time,
});
});
out.sort_by_key(|info| info.pid);
Ok(out)
}
fn observed_start_time(&self, process_key: &str, now_ms: f64) -> f64 {
let _guard = self.inner().observed_process_time_lock.lock();
if let Some(existing) = self
.inner()
.observed_process_start_times
.read(process_key, |_, value| *value)
{
return existing;
}
let _ = self
.inner()
.observed_process_start_times
.insert(process_key.to_owned(), now_ms);
prune_string_f64_map(
&self.inner().observed_process_start_times,
OBSERVED_PROCESS_TIME_LIMIT,
);
self.inner()
.observed_process_start_times
.read(process_key, |_, value| *value)
.unwrap_or(now_ms)
}
fn observed_exit_time(&self, process_id: &str, now_ms: f64) -> f64 {
let _guard = self.inner().observed_process_time_lock.lock();
if let Some(existing) = self
.inner()
.observed_process_exit_times
.read(process_id, |_, value| *value)
{
return existing;
}
let _ = self
.inner()
.observed_process_exit_times
.insert(process_id.to_owned(), now_ms);
prune_string_f64_map(
&self.inner().observed_process_exit_times,
OBSERVED_PROCESS_TIME_LIMIT,
);
self.inner()
.observed_process_exit_times
.read(process_id, |_, value| *value)
.unwrap_or(now_ms)
}
pub async fn process_tree(&self) -> Result<Vec<ProcessTreeNode>> {
let processes = self.all_processes().await?;
Ok(build_process_forest(processes))
}
pub fn get_process(&self, pid: u32) -> std::result::Result<SpawnedProcessInfo, ClientError> {
self.inner()
.processes
.read(&pid, |pid, entry| {
let exit_code = *entry.exit_tx.borrow();
SpawnedProcessInfo {
pid: *pid,
command: entry.command.clone(),
args: entry.args.clone(),
running: exit_code.is_none(),
exit_code,
started_at: entry.started_at,
}
})
.ok_or(ClientError::ProcessNotFound(pid))
}
pub fn stop_process(&self, pid: u32) -> std::result::Result<(), ClientError> {
self.signal_process(pid, "SIGTERM")
}
pub fn kill_process(&self, pid: u32) -> std::result::Result<(), ClientError> {
self.signal_process(pid, "SIGKILL")
}
fn vm_scope(&self) -> wire::OwnershipScope {
wire::OwnershipScope::VmOwnership(wire::VmOwnership {
connection_id: self.connection_id().to_string(),
session_id: self.wire_session_id().to_string(),
vm_id: self.vm_id().to_string(),
})
}
fn next_process_id(&self) -> String {
let n = self.inner().process_counter.fetch_add(1, Ordering::SeqCst);
format!("proc-{n}-{}", uuid::Uuid::new_v4())
}
fn lookup_process_id(&self, pid: u32) -> std::result::Result<String, ClientError> {
self.inner()
.processes
.read(&pid, |_, entry| entry.process_id.clone())
.ok_or(ClientError::ProcessNotFound(pid))
}
async fn send_execute(
&self,
process_id: &str,
command: Option<String>,
args: Vec<String>,
env: BTreeMap<String, String>,
cwd: Option<String>,
) -> std::result::Result<wire::ProcessStartedResponse, ClientError> {
let ownership = self.vm_scope();
let response = self
.transport()
.request_wire(
ownership,
wire::RequestPayload::ExecuteRequest(wire::ExecuteRequest {
process_id: process_id.to_owned(),
command,
runtime: None,
entrypoint: None,
args,
env: env.into_iter().collect(),
cwd,
wasm_permission_tier: None,
}),
)
.await?;
match response {
wire::ResponsePayload::ProcessStartedResponse(started) => Ok(started),
wire::ResponsePayload::RejectedResponse(wire::RejectedResponse { code, message }) => {
Err(ClientError::Kernel { code, message })
}
other => Err(ClientError::Sidecar(format!(
"Execute: unexpected response {other:?}"
))),
}
}
fn kill_wire_process(&self, process_id: &str, signal: &str) {
let process_id = process_id.to_owned();
let signal = signal.to_owned();
let this = self.clone();
tokio::spawn(async move {
let ownership = this.vm_scope();
let _ = this
.transport()
.request_wire(
ownership,
wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
process_id,
signal,
}),
)
.await;
});
}
fn signal_process(&self, pid: u32, signal: &str) -> std::result::Result<(), ClientError> {
let (process_id, already_exited) = self
.inner()
.processes
.read(&pid, |_, entry| {
(entry.process_id.clone(), entry.exit_tx.borrow().is_some())
})
.ok_or(ClientError::ProcessNotFound(pid))?;
if already_exited {
return Ok(());
}
let signal = signal.to_owned();
let this = self.clone();
tokio::spawn(async move {
let ownership = this.vm_scope();
let _ = this
.transport()
.request_wire(
ownership,
wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
process_id,
signal,
}),
)
.await;
});
Ok(())
}
fn process_registry_len_locked(&self) -> usize {
let mut count = 0usize;
self.inner().processes.scan(|_, _| {
count += 1;
});
count
}
fn prune_exited_processes_locked(&self, reserve_slots: usize) {
let mut entries = Vec::new();
self.inner().processes.scan(|pid, entry| {
entries.push((*pid, entry.exit_tx.borrow().is_some()));
});
let target_len = PROCESS_REGISTRY_LIMIT.saturating_sub(reserve_slots);
if entries.len() <= target_len {
return;
}
for pid in exited_pids_to_prune(entries, target_len) {
self.remove_process_tracking_locked(pid);
}
}
fn remove_process_tracking_locked(&self, pid: u32) {
if let Some((_, entry)) = self.inner().processes.remove(&pid) {
let _time_guard = self.inner().observed_process_time_lock.lock();
let _ = self
.inner()
.observed_process_exit_times
.remove(&entry.process_id);
let fallback_start_key = format!("{}:{pid}", entry.process_id);
let _ = self
.inner()
.observed_process_start_times
.remove(&fallback_start_key);
if let Some(kernel_pid) = *entry.kernel_pid.borrow() {
let start_key = format!("{}:{kernel_pid}", entry.process_id);
let _ = self.inner().observed_process_start_times.remove(&start_key);
}
}
}
#[allow(clippy::too_many_arguments)]
async fn run_spawn(
self,
pid: u32,
process_id: String,
command: String,
args: Vec<String>,
options: SpawnOptions,
mut events: broadcast::Receiver<(wire::OwnershipScope, EventPayload)>,
stdout_tx: broadcast::Sender<Vec<u8>>,
stderr_tx: broadcast::Sender<Vec<u8>>,
exit_tx: watch::Sender<Option<i32>>,
kernel_pid_tx: watch::Sender<Option<u32>>,
) {
match self
.send_execute(
&process_id,
Some(command),
args,
options.base.env.clone(),
options.base.cwd.clone(),
)
.await
{
Ok(started) => {
if let Some(kernel_pid) = started.pid {
let _ = kernel_pid_tx.send(Some(kernel_pid));
}
}
Err(error) => {
let message = format!("{error}\n");
let _ = stderr_tx.send(message.into_bytes());
tracing::error!(?error, pid, %process_id, "spawn: Execute request failed");
let _ = exit_tx.send(Some(1));
let _guard = self.inner().process_registry_lock.lock();
self.prune_exited_processes_locked(0);
return;
}
}
loop {
let (_, payload) = match events.recv().await {
Ok(frame) => frame,
Err(broadcast::error::RecvError::Lagged(_)) => continue,
Err(broadcast::error::RecvError::Closed) => {
let _ = exit_tx.send(Some(0));
break;
}
};
match payload {
EventPayload::ProcessOutputEvent(output) if output.process_id == process_id => {
let bytes = output.chunk;
match output.channel {
StreamChannel::Stdout => {
let _ = stdout_tx.send(bytes);
}
StreamChannel::Stderr => {
let _ = stderr_tx.send(bytes);
}
}
}
EventPayload::ProcessExitedEvent(exited) if exited.process_id == process_id => {
let _ = exit_tx.send(Some(exited.exit_code));
break;
}
EventPayload::ProcessOutputEvent(_)
| EventPayload::ProcessExitedEvent(_)
| EventPayload::VmLifecycleEvent(_)
| EventPayload::StructuredEvent(_)
| EventPayload::ExtEnvelope(_) => {}
}
}
let _guard = self.inner().process_registry_lock.lock();
self.prune_exited_processes_locked(0);
}
}
fn build_process_forest(processes: Vec<ProcessInfo>) -> Vec<ProcessTreeNode> {
use std::collections::BTreeMap as Map;
let pids: std::collections::BTreeSet<u32> = processes.iter().map(|p| p.pid).collect();
let mut children_of: Map<u32, Vec<usize>> = Map::new();
let mut roots: Vec<usize> = Vec::new();
for (index, proc) in processes.iter().enumerate() {
if pids.contains(&proc.ppid) {
children_of.entry(proc.ppid).or_default().push(index);
} else {
roots.push(index);
}
}
fn build_node(
index: usize,
processes: &[ProcessInfo],
children_of: &Map<u32, Vec<usize>>,
seen: &mut std::collections::BTreeSet<usize>,
) -> ProcessTreeNode {
let info = processes[index].clone();
seen.insert(index);
let child_indices: Vec<usize> = children_of
.get(&info.pid)
.map(|indices| {
indices
.iter()
.copied()
.filter(|child_index| !seen.contains(child_index))
.collect()
})
.unwrap_or_default();
let children = child_indices
.into_iter()
.map(|child_index| build_node(child_index, processes, children_of, seen))
.collect();
ProcessTreeNode { info, children }
}
let mut seen = std::collections::BTreeSet::new();
roots
.into_iter()
.map(|index| build_node(index, &processes, &children_of, &mut seen))
.collect()
}
fn stdin_to_bytes(input: StdinInput) -> Vec<u8> {
match input {
StdinInput::Text(text) => text.into_bytes(),
StdinInput::Bytes(bytes) => bytes,
}
}
fn append_exec_output(
buffer: &mut Vec<u8>,
chunk: &[u8],
captured_output_bytes: &mut usize,
channel: &str,
) -> std::result::Result<(), ClientError> {
let next_total = captured_output_bytes
.checked_add(chunk.len())
.ok_or_else(|| exec_output_limit_error(channel, usize::MAX))?;
if next_total > EXEC_OUTPUT_CAPTURE_LIMIT_BYTES {
return Err(exec_output_limit_error(channel, next_total));
}
buffer.extend_from_slice(chunk);
*captured_output_bytes = next_total;
Ok(())
}
fn exec_output_limit_error(channel: &str, size: usize) -> ClientError {
ClientError::Sidecar(format!(
"exec {channel} capture is {size} bytes, limit is {EXEC_OUTPUT_CAPTURE_LIMIT_BYTES}"
))
}
fn exited_pids_to_prune(mut entries: Vec<(u32, bool)>, target_len: usize) -> Vec<u32> {
if entries.len() <= target_len {
return Vec::new();
}
let mut remove_count = entries.len() - target_len;
entries.sort_by_key(|(pid, _)| *pid);
let mut out = Vec::new();
for (pid, exited) in entries {
if remove_count == 0 {
break;
}
if !exited {
continue;
}
out.push(pid);
remove_count -= 1;
}
out
}
fn prune_string_f64_map(map: &SccHashMap<String, f64>, limit: usize) {
let mut keys = Vec::new();
map.scan(|key, _| {
keys.push(key.clone());
});
if keys.len() <= limit {
return;
}
let remove_count = keys.len() - limit;
keys.sort();
for key in keys.into_iter().take(remove_count) {
let _ = map.remove(&key);
}
}
pub(crate) fn install_output_callback(
tx: broadcast::Sender<Vec<u8>>,
mut callback: OutputCallback,
) -> JoinHandle<()> {
let mut rx = tx.subscribe();
tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(chunk) => callback(&chunk),
Err(broadcast::error::RecvError::Lagged(_)) => continue,
Err(broadcast::error::RecvError::Closed) => break,
}
}
})
}
pub(crate) fn drain_process_output_tasks(processes: &SccHashMap<u32, ProcessEntry>) {
let mut tasks = Vec::new();
processes.retain(|_, entry| {
tasks.append(&mut entry.output_tasks);
false
});
for task in tasks {
task.abort();
}
}
fn epoch_ms_now() -> f64 {
use std::time::{SystemTime, UNIX_EPOCH};
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs_f64() * 1000.0)
.unwrap_or(0.0)
}
#[cfg(test)]
mod tests {
use super::{
append_exec_output, drain_process_output_tasks, exited_pids_to_prune,
install_output_callback, prune_string_f64_map, ExecOptions, OutputCallback,
DEFAULT_EXEC_CWD, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES,
};
use crate::agent_os::ProcessEntry;
use scc::HashMap as SccHashMap;
use tokio::sync::{broadcast, watch};
#[tokio::test]
async fn drain_process_output_tasks_clears_registry_and_aborts_tasks() {
let processes: SccHashMap<u32, ProcessEntry> = SccHashMap::new();
let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(8);
let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(8);
let (exit_tx, _) = watch::channel::<Option<i32>>(None);
let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
let task = tokio::spawn(async {
loop {
tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
}
});
let abort_handle = task.abort_handle();
let entry = ProcessEntry {
command: "sleep".to_string(),
args: vec!["3600".to_string()],
stdout_tx,
stderr_tx,
exit_tx,
process_id: "proc-test".to_string(),
kernel_pid: kernel_pid_tx,
output_tasks: vec![task],
started_at: 0,
};
let _ = processes.insert(1, entry);
assert!(!abort_handle.is_finished(), "task should start alive");
drain_process_output_tasks(&processes);
assert!(processes.is_empty(), "registry must be cleared on drain");
for _ in 0..100 {
if abort_handle.is_finished() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
}
assert!(
abort_handle.is_finished(),
"output task must be aborted after drain"
);
}
#[tokio::test]
async fn install_output_callback_handle_is_captured_into_process_entry() {
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(8);
let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(8);
let (exit_tx, _) = watch::channel::<Option<i32>>(None);
let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
let calls = Arc::new(AtomicUsize::new(0));
let calls_cb = Arc::clone(&calls);
let cb: OutputCallback = Box::new(move |_chunk: &[u8]| {
calls_cb.fetch_add(1, Ordering::SeqCst);
});
let output_tasks = vec![install_output_callback(stdout_tx.clone(), cb)];
let entry = ProcessEntry {
command: "sleep".to_string(),
args: vec!["3600".to_string()],
stdout_tx: stdout_tx.clone(),
stderr_tx,
exit_tx,
process_id: "proc-test".to_string(),
kernel_pid: kernel_pid_tx,
output_tasks,
started_at: 0,
};
assert_eq!(
entry.output_tasks.len(),
1,
"the install_output_callback handle must be captured on the entry"
);
stdout_tx
.send(b"hello".to_vec())
.expect("broadcast send to subscribed callback task");
for _ in 0..100 {
if calls.load(Ordering::SeqCst) > 0 {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
}
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"the stored handle must drive the registered callback"
);
let processes: SccHashMap<u32, ProcessEntry> = SccHashMap::new();
let _ = processes.insert(1, entry);
drain_process_output_tasks(&processes);
assert!(processes.is_empty(), "registry must be cleared on drain");
}
#[test]
fn exec_options_default_uses_workspace_cwd() {
assert_eq!(
ExecOptions::default().cwd.as_deref(),
Some(DEFAULT_EXEC_CWD)
);
}
#[test]
fn append_exec_output_rejects_capture_over_limit() {
let mut buffer = vec![0u8; EXEC_OUTPUT_CAPTURE_LIMIT_BYTES - 1];
let mut captured = buffer.len();
append_exec_output(&mut buffer, &[1], &mut captured, "stdout")
.expect("chunk at limit should fit");
assert_eq!(captured, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
let error = append_exec_output(&mut buffer, &[2], &mut captured, "stdout")
.expect_err("chunk over limit should fail");
assert!(
error.to_string().contains("exec stdout capture is"),
"unexpected error: {error}"
);
assert_eq!(captured, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
assert_eq!(buffer.len(), EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
}
#[test]
fn exited_pid_pruning_keeps_live_entries_and_removes_oldest_exited() {
let pids = exited_pids_to_prune(vec![(3, true), (1, false), (2, true), (4, true)], 2);
assert_eq!(pids, vec![2, 3]);
}
#[test]
fn observed_time_pruning_enforces_limit() {
let map = SccHashMap::new();
let _ = map.insert("b".to_string(), 2.0);
let _ = map.insert("a".to_string(), 1.0);
let _ = map.insert("c".to_string(), 3.0);
prune_string_f64_map(&map, 2);
assert!(map.read("a", |_, _| ()).is_none());
assert!(map.read("b", |_, _| ()).is_some());
assert!(map.read("c", |_, _| ()).is_some());
}
}