1use std::collections::BTreeMap;
11use std::sync::atomic::Ordering;
12
13use anyhow::{Context, Result};
14use scc::HashMap as SccHashMap;
15use serde::{Deserialize, Serialize};
16use tokio::sync::{broadcast, watch};
17use tokio::task::JoinHandle;
18
19use agentos_sidecar_client::wire::{self, EventPayload, ProcessSnapshotStatus, StreamChannel};
20
21use crate::agent_os::{AgentOs, ProcessEntry};
22use crate::command_line::resolve_exec_command;
23use crate::error::ClientError;
24use crate::stream::{ByteStream, Subscription};
25
26const PROCESS_STREAM_CAPACITY: usize = 1024;
28
29const PROCESS_REGISTRY_LIMIT: usize = 1024;
31
32const OBSERVED_PROCESS_TIME_LIMIT: usize = 4096;
34
35const EXEC_OUTPUT_CAPTURE_LIMIT_BYTES: usize = 16 * 1024 * 1024;
37
38pub(crate) const DEFAULT_EXEC_CWD: &str = "/workspace";
40
41pub(crate) const SYNTHETIC_PID_BASE: u64 = 1_000_000;
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
51#[serde(rename_all = "lowercase")]
52pub enum TimingMitigation {
53 #[default]
54 Off,
55 Freeze,
56}
57
58#[derive(Debug, Clone, PartialEq, Eq)]
60pub enum StdinInput {
61 Text(String),
62 Bytes(Vec<u8>),
63}
64
65pub type OutputCallback = Box<dyn FnMut(&[u8]) + Send>;
68
69pub struct ExecOptions {
75 pub env: BTreeMap<String, String>,
76 pub cwd: Option<String>,
77 pub stdin: Option<StdinInput>,
78 pub timeout: Option<f64>,
79 pub on_stdout: Option<OutputCallback>,
80 pub on_stderr: Option<OutputCallback>,
81 pub capture_stdio: Option<bool>,
82 pub file_path: Option<String>,
83 pub cpu_time_limit_ms: Option<f64>,
84 pub timing_mitigation: Option<TimingMitigation>,
85}
86
87impl Default for ExecOptions {
88 fn default() -> Self {
89 Self {
90 env: BTreeMap::new(),
91 cwd: Some(DEFAULT_EXEC_CWD.to_string()),
92 stdin: None,
93 timeout: None,
94 on_stdout: None,
95 on_stderr: None,
96 capture_stdio: None,
97 file_path: None,
98 cpu_time_limit_ms: None,
99 timing_mitigation: None,
100 }
101 }
102}
103
104#[derive(Debug, Clone, PartialEq, Eq)]
106pub struct ExecResult {
107 pub exit_code: i32,
108 pub stdout: String,
109 pub stderr: String,
110}
111
112#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
114#[serde(rename_all = "lowercase")]
115pub enum SpawnStdio {
116 #[default]
117 Pipe,
118 Inherit,
119}
120
121#[derive(Default)]
123pub struct SpawnOptions {
124 pub base: ExecOptions,
125 pub stdio: Option<SpawnStdio>,
126 pub stdin_fd: Option<i32>,
127 pub stdout_fd: Option<i32>,
128 pub stderr_fd: Option<i32>,
129 pub stream_stdin: Option<bool>,
130}
131
132#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
134pub struct SpawnedProcessInfo {
135 pub pid: u32,
136 pub command: String,
137 pub args: Vec<String>,
138 pub running: bool,
139 #[serde(rename = "exitCode")]
140 pub exit_code: Option<i32>,
141 #[serde(rename = "startedAt")]
143 pub started_at: i64,
144}
145
146#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
148pub struct SpawnHandle {
149 pub pid: u32,
150}
151
152#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
154#[serde(rename_all = "lowercase")]
155pub enum ProcessStatus {
156 Running,
157 Exited,
158}
159
160#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
162pub struct ProcessInfo {
163 pub pid: u32,
164 pub ppid: u32,
165 pub pgid: u32,
166 pub sid: u32,
167 pub driver: String,
168 pub command: String,
169 pub args: Vec<String>,
170 pub cwd: String,
171 pub status: ProcessStatus,
172 #[serde(rename = "exitCode")]
173 pub exit_code: Option<i32>,
174 #[serde(rename = "startTime")]
175 pub start_time: f64,
176 #[serde(rename = "exitTime")]
177 pub exit_time: Option<f64>,
178}
179
180#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
182pub struct ProcessTreeNode {
183 #[serde(flatten)]
184 pub info: ProcessInfo,
185 pub children: Vec<ProcessTreeNode>,
186}
187
188impl AgentOs {
193 pub async fn exec(&self, command: &str, options: ExecOptions) -> Result<ExecResult> {
198 let (resolved_command, resolved_args) = resolve_exec_command(command)?;
202 self.exec_argv(&resolved_command, &resolved_args, options)
203 .await
204 }
205
206 pub async fn exec_argv(
212 &self,
213 command: &str,
214 args: &[String],
215 mut options: ExecOptions,
216 ) -> Result<ExecResult> {
217 let process_id = self.next_process_id();
218
219 let mut events = self.transport().subscribe_wire_events();
222
223 let resolved_command = command.to_owned();
224 let resolved_args = args.to_vec();
225 let started = self
226 .send_execute(
227 &process_id,
228 Some(resolved_command),
229 resolved_args,
230 options.env.clone(),
231 options.cwd.clone(),
232 )
233 .await
234 .context("exec: Execute request failed")?;
235 debug_assert_eq!(started.process_id, process_id);
236
237 if let Some(stdin) = options.stdin.take() {
240 let chunk = stdin_to_bytes(stdin);
241 let ownership = self.vm_scope();
242 let _ = self
243 .transport()
244 .request_wire(
245 ownership,
246 wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {
247 process_id: process_id.clone(),
248 chunk,
249 }),
250 )
251 .await;
252 }
253 {
254 let ownership = self.vm_scope();
255 let _ = self
256 .transport()
257 .request_wire(
258 ownership,
259 wire::RequestPayload::CloseStdinRequest(wire::CloseStdinRequest {
260 process_id: process_id.clone(),
261 }),
262 )
263 .await;
264 }
265
266 let mut on_stdout = options.on_stdout.take();
267 let mut on_stderr = options.on_stderr.take();
268
269 let timeout_deadline = options
273 .timeout
274 .filter(|ms| ms.is_finite() && *ms >= 0.0)
275 .map(|ms| {
276 tokio::time::Instant::now() + std::time::Duration::from_secs_f64(ms / 1000.0)
277 });
278 let mut killed_for_timeout = false;
279
280 let capture_stdio = options.capture_stdio.unwrap_or(true);
281 let mut stdout = Vec::<u8>::new();
282 let mut stderr = Vec::<u8>::new();
283 let mut captured_output_bytes = 0usize;
284 let mut capture_error: Option<ClientError> = None;
285 let exit_code = loop {
286 let recv = events.recv();
287 let frame = match timeout_deadline {
288 Some(deadline) => {
289 tokio::select! {
290 result = recv => result,
291 _ = tokio::time::sleep_until(deadline), if !killed_for_timeout => {
292 killed_for_timeout = true;
293 self.kill_wire_process(&process_id, "SIGKILL");
294 continue;
295 }
296 }
297 }
298 None => recv.await,
299 };
300 let (_, payload) = match frame {
301 Ok(frame) => frame,
302 Err(broadcast::error::RecvError::Lagged(_)) => continue,
303 Err(broadcast::error::RecvError::Closed) => {
304 return Err(ClientError::Sidecar(
305 "exec: event stream closed before process exit".to_owned(),
306 )
307 .into());
308 }
309 };
310 match payload {
311 EventPayload::ProcessOutputEvent(output) if output.process_id == process_id => {
312 match output.channel {
313 StreamChannel::Stdout => {
314 if let Some(cb) = on_stdout.as_mut() {
315 cb(&output.chunk);
316 }
317 if capture_stdio && capture_error.is_none() {
318 match append_exec_output(
319 &mut stdout,
320 &output.chunk,
321 &mut captured_output_bytes,
322 "stdout",
323 ) {
324 Ok(()) => {}
325 Err(error) => {
326 self.kill_wire_process(&process_id, "SIGKILL");
327 capture_error = Some(error);
328 }
329 }
330 }
331 }
332 StreamChannel::Stderr => {
333 if let Some(cb) = on_stderr.as_mut() {
334 cb(&output.chunk);
335 }
336 if capture_stdio && capture_error.is_none() {
337 match append_exec_output(
338 &mut stderr,
339 &output.chunk,
340 &mut captured_output_bytes,
341 "stderr",
342 ) {
343 Ok(()) => {}
344 Err(error) => {
345 self.kill_wire_process(&process_id, "SIGKILL");
346 capture_error = Some(error);
347 }
348 }
349 }
350 }
351 }
352 }
353 EventPayload::ProcessExitedEvent(exited) if exited.process_id == process_id => {
354 break exited.exit_code;
355 }
356 EventPayload::ProcessOutputEvent(_)
357 | EventPayload::ProcessExitedEvent(_)
358 | EventPayload::VmLifecycleEvent(_)
359 | EventPayload::StructuredEvent(_)
360 | EventPayload::ExtEnvelope(_) => {}
361 }
362 };
363
364 if let Some(error) = capture_error {
365 return Err(error.into());
366 }
367
368 Ok(ExecResult {
369 exit_code,
370 stdout: String::from_utf8_lossy(&stdout).into_owned(),
371 stderr: String::from_utf8_lossy(&stderr).into_owned(),
372 })
373 }
374
375 pub fn spawn(
379 &self,
380 command: &str,
381 args: Vec<String>,
382 mut options: SpawnOptions,
383 ) -> Result<SpawnHandle> {
384 let registry_guard = self.inner().process_registry_lock.lock();
385 self.prune_exited_processes_locked(1);
386 if self.process_registry_len_locked() >= PROCESS_REGISTRY_LIMIT {
387 return Err(ClientError::Sidecar(format!(
388 "process registry limit exceeded: at most {PROCESS_REGISTRY_LIMIT} processes can be tracked per VM"
389 ))
390 .into());
391 }
392
393 let pid = self
396 .inner()
397 .synthetic_pid_counter
398 .fetch_add(1, Ordering::SeqCst) as u32;
399 let process_id = format!("proc-{pid}-{}", uuid::Uuid::new_v4());
400
401 let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(PROCESS_STREAM_CAPACITY);
402 let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(PROCESS_STREAM_CAPACITY);
403 let (exit_tx, _) = watch::channel::<Option<i32>>(None);
406 let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
409
410 let mut output_tasks = Vec::new();
415 if let Some(cb) = options.base.on_stdout.take() {
416 output_tasks.push(install_output_callback(stdout_tx.clone(), cb));
417 }
418 if let Some(cb) = options.base.on_stderr.take() {
419 output_tasks.push(install_output_callback(stderr_tx.clone(), cb));
420 }
421
422 let entry = ProcessEntry {
423 command: command.to_owned(),
424 args: args.clone(),
425 stdout_tx: stdout_tx.clone(),
426 stderr_tx: stderr_tx.clone(),
427 exit_tx: exit_tx.clone(),
428 process_id: process_id.clone(),
429 kernel_pid: kernel_pid_tx.clone(),
430 output_tasks,
431 started_at: epoch_ms_now() as i64,
432 };
433 let _ = self.inner().processes.insert(pid, entry);
436 drop(registry_guard);
437
438 let events = self.transport().subscribe_wire_events();
440
441 let this = self.clone();
442 let command = command.to_owned();
443 tokio::spawn(async move {
444 this.run_spawn(
445 pid,
446 process_id,
447 command,
448 args,
449 options,
450 events,
451 stdout_tx,
452 stderr_tx,
453 exit_tx,
454 kernel_pid_tx,
455 )
456 .await;
457 });
458
459 Ok(SpawnHandle { pid })
460 }
461
462 pub fn write_process_stdin(
464 &self,
465 pid: u32,
466 data: StdinInput,
467 ) -> std::result::Result<(), ClientError> {
468 let process_id = self.lookup_process_id(pid)?;
469 let chunk: Vec<u8> = stdin_to_bytes(data);
470 let this = self.clone();
471 tokio::spawn(async move {
473 let ownership = this.vm_scope();
474 let _ = this
475 .transport()
476 .request_wire(
477 ownership,
478 wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {
479 process_id,
480 chunk,
481 }),
482 )
483 .await;
484 });
485 Ok(())
486 }
487
488 pub fn close_process_stdin(&self, pid: u32) -> std::result::Result<(), ClientError> {
490 let process_id = self.lookup_process_id(pid)?;
491 let this = self.clone();
492 tokio::spawn(async move {
493 let ownership = this.vm_scope();
494 let _ = this
495 .transport()
496 .request_wire(
497 ownership,
498 wire::RequestPayload::CloseStdinRequest(wire::CloseStdinRequest { process_id }),
499 )
500 .await;
501 });
502 Ok(())
503 }
504
505 pub fn on_process_stdout(&self, pid: u32) -> std::result::Result<ByteStream, ClientError> {
507 let rx = self
508 .inner()
509 .processes
510 .read(&pid, |_, entry| entry.stdout_tx.subscribe())
511 .ok_or(ClientError::ProcessNotFound(pid))?;
512 Ok(ByteStream::new(rx))
513 }
514
515 pub fn on_process_stderr(&self, pid: u32) -> std::result::Result<ByteStream, ClientError> {
517 let rx = self
518 .inner()
519 .processes
520 .read(&pid, |_, entry| entry.stderr_tx.subscribe())
521 .ok_or(ClientError::ProcessNotFound(pid))?;
522 Ok(ByteStream::new(rx))
523 }
524
525 pub fn on_process_exit(
530 &self,
531 pid: u32,
532 handler: impl FnOnce(i32) + Send + 'static,
533 ) -> std::result::Result<Subscription, ClientError> {
534 let mut rx = self
535 .inner()
536 .processes
537 .read(&pid, |_, entry| entry.exit_tx.subscribe())
538 .ok_or(ClientError::ProcessNotFound(pid))?;
539
540 if let Some(code) = *rx.borrow() {
542 handler(code);
543 return Ok(Subscription::noop());
544 }
545
546 let task = tokio::spawn(async move {
549 while rx.changed().await.is_ok() {
550 if let Some(code) = *rx.borrow() {
551 handler(code);
552 return;
553 }
554 }
555 });
556 Ok(Subscription::new(move || task.abort()))
557 }
558
559 pub async fn wait_process(&self, pid: u32) -> std::result::Result<i32, ClientError> {
562 let mut rx = self
563 .inner()
564 .processes
565 .read(&pid, |_, entry| entry.exit_tx.subscribe())
566 .ok_or(ClientError::ProcessNotFound(pid))?;
567
568 if let Some(code) = *rx.borrow() {
569 return Ok(code);
570 }
571 while rx.changed().await.is_ok() {
572 if let Some(code) = *rx.borrow() {
573 return Ok(code);
574 }
575 }
576 Err(ClientError::Sidecar(format!(
577 "wait_process: exit channel closed before process {pid} reported an exit code"
578 )))
579 }
580
581 pub fn list_processes(&self) -> Vec<SpawnedProcessInfo> {
583 let mut out = Vec::new();
584 self.inner().processes.scan(|pid, entry| {
585 let exit_code = *entry.exit_tx.borrow();
586 out.push(SpawnedProcessInfo {
587 pid: *pid,
588 command: entry.command.clone(),
589 args: entry.args.clone(),
590 running: exit_code.is_none(),
591 exit_code,
592 started_at: entry.started_at,
593 });
594 });
595 out
596 }
597
598 pub async fn all_processes(&self) -> Result<Vec<ProcessInfo>> {
606 let ownership = self.vm_scope();
607 let response = self
608 .transport()
609 .request_wire(ownership, wire::RequestPayload::GetProcessSnapshotRequest)
610 .await
611 .context("all_processes: GetProcessSnapshot request failed")?;
612 let snapshot = match response {
613 wire::ResponsePayload::ProcessSnapshotResponse(snapshot) => snapshot,
614 wire::ResponsePayload::RejectedResponse(wire::RejectedResponse { code, message }) => {
615 return Err(ClientError::Kernel { code, message }.into());
616 }
617 other => {
618 return Err(ClientError::Sidecar(format!(
619 "all_processes: unexpected response {other:?}"
620 ))
621 .into());
622 }
623 };
624
625 struct Tracked {
629 exit_code: Option<i32>,
630 command: String,
631 args: Vec<String>,
632 }
633 let mut tracked_by_process_id: BTreeMap<String, Tracked> = BTreeMap::new();
634 let mut display_pid_by_kernel_pid: BTreeMap<u32, u32> = BTreeMap::new();
635 self.inner().processes.scan(|display_pid, entry| {
636 let exit_code = *entry.exit_tx.borrow();
637 if let Some(kernel_pid) = *entry.kernel_pid.borrow() {
638 display_pid_by_kernel_pid.insert(kernel_pid, *display_pid);
639 }
640 tracked_by_process_id.insert(
641 entry.process_id.clone(),
642 Tracked {
643 exit_code,
644 command: entry.command.clone(),
645 args: entry.args.clone(),
646 },
647 );
648 });
649
650 let now_ms = epoch_ms_now();
651 let mut seen_display_pids: std::collections::BTreeSet<u32> =
652 std::collections::BTreeSet::new();
653 let mut out: Vec<ProcessInfo> = Vec::new();
654
655 for entry in snapshot.processes {
656 let tracked = tracked_by_process_id.get(&entry.process_id);
657 let display_pid = display_pid_by_kernel_pid
658 .get(&entry.pid)
659 .copied()
660 .unwrap_or(entry.pid);
661 let display_ppid = display_pid_by_kernel_pid
662 .get(&entry.ppid)
663 .copied()
664 .unwrap_or(entry.ppid);
665 let display_pgid = display_pid_by_kernel_pid
666 .get(&entry.pgid)
667 .copied()
668 .unwrap_or(entry.pgid);
669 let display_sid = display_pid_by_kernel_pid
670 .get(&entry.sid)
671 .copied()
672 .unwrap_or(entry.sid);
673
674 let process_key = format!("{}:{}", entry.process_id, entry.pid);
676 let start_time = self.observed_start_time(&process_key, now_ms);
677
678 let (status, exit_code) = match tracked {
681 Some(t) => match t.exit_code {
682 Some(code) => (ProcessStatus::Exited, Some(code)),
683 None => (ProcessStatus::Running, entry.exit_code),
684 },
685 None => {
686 let status = match entry.status {
687 ProcessSnapshotStatus::Running | ProcessSnapshotStatus::Stopped => {
688 ProcessStatus::Running
689 }
690 ProcessSnapshotStatus::Exited => ProcessStatus::Exited,
691 };
692 (status, entry.exit_code)
693 }
694 };
695
696 let exit_time = match (tracked, status) {
698 (Some(_), ProcessStatus::Exited) => {
699 Some(self.observed_exit_time(&entry.process_id, now_ms))
700 }
701 _ => None,
702 };
703
704 let (command, args) = match tracked {
705 Some(t) => (t.command.clone(), t.args.clone()),
706 None => (entry.command, entry.args),
707 };
708
709 seen_display_pids.insert(display_pid);
710 out.push(ProcessInfo {
711 pid: display_pid,
712 ppid: display_ppid,
713 pgid: display_pgid,
714 sid: display_sid,
715 driver: entry.driver,
716 command,
717 args,
718 cwd: entry.cwd,
719 status,
720 exit_code,
721 start_time,
722 exit_time,
723 });
724 }
725
726 self.inner().processes.scan(|display_pid, entry| {
729 if seen_display_pids.contains(display_pid) {
730 return;
731 }
732 let exit_code = *entry.exit_tx.borrow();
733 let process_key = format!("{}:{}", entry.process_id, display_pid);
734 let start_time = self.observed_start_time(&process_key, now_ms);
735 let (status, exit_time) = match exit_code {
736 Some(_) => (
737 ProcessStatus::Exited,
738 Some(self.observed_exit_time(&entry.process_id, now_ms)),
739 ),
740 None => (ProcessStatus::Running, None),
741 };
742 out.push(ProcessInfo {
743 pid: *display_pid,
744 ppid: 0,
745 pgid: *display_pid,
746 sid: *display_pid,
747 driver: String::new(),
748 command: entry.command.clone(),
749 args: entry.args.clone(),
750 cwd: String::new(),
751 status,
752 exit_code,
753 start_time,
754 exit_time,
755 });
756 });
757
758 out.sort_by_key(|info| info.pid);
759 Ok(out)
760 }
761
762 fn observed_start_time(&self, process_key: &str, now_ms: f64) -> f64 {
765 let _guard = self.inner().observed_process_time_lock.lock();
766 if let Some(existing) = self
767 .inner()
768 .observed_process_start_times
769 .read(process_key, |_, value| *value)
770 {
771 return existing;
772 }
773 let _ = self
774 .inner()
775 .observed_process_start_times
776 .insert(process_key.to_owned(), now_ms);
777 prune_string_f64_map(
778 &self.inner().observed_process_start_times,
779 OBSERVED_PROCESS_TIME_LIMIT,
780 );
781 self.inner()
784 .observed_process_start_times
785 .read(process_key, |_, value| *value)
786 .unwrap_or(now_ms)
787 }
788
789 fn observed_exit_time(&self, process_id: &str, now_ms: f64) -> f64 {
791 let _guard = self.inner().observed_process_time_lock.lock();
792 if let Some(existing) = self
793 .inner()
794 .observed_process_exit_times
795 .read(process_id, |_, value| *value)
796 {
797 return existing;
798 }
799 let _ = self
800 .inner()
801 .observed_process_exit_times
802 .insert(process_id.to_owned(), now_ms);
803 prune_string_f64_map(
804 &self.inner().observed_process_exit_times,
805 OBSERVED_PROCESS_TIME_LIMIT,
806 );
807 self.inner()
808 .observed_process_exit_times
809 .read(process_id, |_, value| *value)
810 .unwrap_or(now_ms)
811 }
812
813 pub async fn process_tree(&self) -> Result<Vec<ProcessTreeNode>> {
815 let processes = self.all_processes().await?;
816 Ok(build_process_forest(processes))
817 }
818
819 pub fn get_process(&self, pid: u32) -> std::result::Result<SpawnedProcessInfo, ClientError> {
821 self.inner()
822 .processes
823 .read(&pid, |pid, entry| {
824 let exit_code = *entry.exit_tx.borrow();
825 SpawnedProcessInfo {
826 pid: *pid,
827 command: entry.command.clone(),
828 args: entry.args.clone(),
829 running: exit_code.is_none(),
830 exit_code,
831 started_at: entry.started_at,
832 }
833 })
834 .ok_or(ClientError::ProcessNotFound(pid))
835 }
836
837 pub fn stop_process(&self, pid: u32) -> std::result::Result<(), ClientError> {
839 self.signal_process(pid, "SIGTERM")
840 }
841
842 pub fn kill_process(&self, pid: u32) -> std::result::Result<(), ClientError> {
844 self.signal_process(pid, "SIGKILL")
845 }
846
847 fn vm_scope(&self) -> wire::OwnershipScope {
853 wire::OwnershipScope::VmOwnership(wire::VmOwnership {
854 connection_id: self.connection_id().to_string(),
855 session_id: self.wire_session_id().to_string(),
856 vm_id: self.vm_id().to_string(),
857 })
858 }
859
860 fn next_process_id(&self) -> String {
862 let n = self.inner().process_counter.fetch_add(1, Ordering::SeqCst);
863 format!("proc-{n}-{}", uuid::Uuid::new_v4())
864 }
865
866 fn lookup_process_id(&self, pid: u32) -> std::result::Result<String, ClientError> {
868 self.inner()
869 .processes
870 .read(&pid, |_, entry| entry.process_id.clone())
871 .ok_or(ClientError::ProcessNotFound(pid))
872 }
873
874 async fn send_execute(
876 &self,
877 process_id: &str,
878 command: Option<String>,
879 args: Vec<String>,
880 env: BTreeMap<String, String>,
881 cwd: Option<String>,
882 ) -> std::result::Result<wire::ProcessStartedResponse, ClientError> {
883 let ownership = self.vm_scope();
884 let response = self
885 .transport()
886 .request_wire(
887 ownership,
888 wire::RequestPayload::ExecuteRequest(wire::ExecuteRequest {
889 process_id: process_id.to_owned(),
890 command,
891 runtime: None,
892 entrypoint: None,
893 args,
894 env: env.into_iter().collect(),
895 cwd,
896 wasm_permission_tier: None,
897 }),
898 )
899 .await?;
900 match response {
901 wire::ResponsePayload::ProcessStartedResponse(started) => Ok(started),
902 wire::ResponsePayload::RejectedResponse(wire::RejectedResponse { code, message }) => {
903 Err(ClientError::Kernel { code, message })
904 }
905 other => Err(ClientError::Sidecar(format!(
906 "Execute: unexpected response {other:?}"
907 ))),
908 }
909 }
910
911 fn kill_wire_process(&self, process_id: &str, signal: &str) {
914 let process_id = process_id.to_owned();
915 let signal = signal.to_owned();
916 let this = self.clone();
917 tokio::spawn(async move {
918 let ownership = this.vm_scope();
919 let _ = this
920 .transport()
921 .request_wire(
922 ownership,
923 wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
924 process_id,
925 signal,
926 }),
927 )
928 .await;
929 });
930 }
931
932 fn signal_process(&self, pid: u32, signal: &str) -> std::result::Result<(), ClientError> {
935 let (process_id, already_exited) = self
936 .inner()
937 .processes
938 .read(&pid, |_, entry| {
939 (entry.process_id.clone(), entry.exit_tx.borrow().is_some())
940 })
941 .ok_or(ClientError::ProcessNotFound(pid))?;
942 if already_exited {
943 return Ok(());
944 }
945 let signal = signal.to_owned();
946 let this = self.clone();
947 tokio::spawn(async move {
948 let ownership = this.vm_scope();
949 let _ = this
950 .transport()
951 .request_wire(
952 ownership,
953 wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
954 process_id,
955 signal,
956 }),
957 )
958 .await;
959 });
960 Ok(())
961 }
962
963 fn process_registry_len_locked(&self) -> usize {
964 let mut count = 0usize;
965 self.inner().processes.scan(|_, _| {
966 count += 1;
967 });
968 count
969 }
970
971 fn prune_exited_processes_locked(&self, reserve_slots: usize) {
972 let mut entries = Vec::new();
973 self.inner().processes.scan(|pid, entry| {
974 entries.push((*pid, entry.exit_tx.borrow().is_some()));
975 });
976 let target_len = PROCESS_REGISTRY_LIMIT.saturating_sub(reserve_slots);
977 if entries.len() <= target_len {
978 return;
979 }
980
981 for pid in exited_pids_to_prune(entries, target_len) {
982 self.remove_process_tracking_locked(pid);
983 }
984 }
985
986 fn remove_process_tracking_locked(&self, pid: u32) {
987 if let Some((_, entry)) = self.inner().processes.remove(&pid) {
988 let _time_guard = self.inner().observed_process_time_lock.lock();
989 let _ = self
990 .inner()
991 .observed_process_exit_times
992 .remove(&entry.process_id);
993 let fallback_start_key = format!("{}:{pid}", entry.process_id);
994 let _ = self
995 .inner()
996 .observed_process_start_times
997 .remove(&fallback_start_key);
998 if let Some(kernel_pid) = *entry.kernel_pid.borrow() {
999 let start_key = format!("{}:{kernel_pid}", entry.process_id);
1000 let _ = self.inner().observed_process_start_times.remove(&start_key);
1001 }
1002 }
1003 }
1004
1005 #[allow(clippy::too_many_arguments)]
1010 async fn run_spawn(
1011 self,
1012 pid: u32,
1013 process_id: String,
1014 command: String,
1015 args: Vec<String>,
1016 options: SpawnOptions,
1017 mut events: broadcast::Receiver<(wire::OwnershipScope, EventPayload)>,
1018 stdout_tx: broadcast::Sender<Vec<u8>>,
1019 stderr_tx: broadcast::Sender<Vec<u8>>,
1020 exit_tx: watch::Sender<Option<i32>>,
1021 kernel_pid_tx: watch::Sender<Option<u32>>,
1022 ) {
1023 match self
1024 .send_execute(
1025 &process_id,
1026 Some(command),
1027 args,
1028 options.base.env.clone(),
1029 options.base.cwd.clone(),
1030 )
1031 .await
1032 {
1033 Ok(started) => {
1034 if let Some(kernel_pid) = started.pid {
1037 let _ = kernel_pid_tx.send(Some(kernel_pid));
1038 }
1039 }
1040 Err(error) => {
1041 let message = format!("{error}\n");
1045 let _ = stderr_tx.send(message.into_bytes());
1046 tracing::error!(?error, pid, %process_id, "spawn: Execute request failed");
1047 let _ = exit_tx.send(Some(1));
1048 let _guard = self.inner().process_registry_lock.lock();
1049 self.prune_exited_processes_locked(0);
1050 return;
1051 }
1052 }
1053
1054 loop {
1055 let (_, payload) = match events.recv().await {
1056 Ok(frame) => frame,
1057 Err(broadcast::error::RecvError::Lagged(_)) => continue,
1058 Err(broadcast::error::RecvError::Closed) => {
1059 let _ = exit_tx.send(Some(0));
1063 break;
1064 }
1065 };
1066 match payload {
1067 EventPayload::ProcessOutputEvent(output) if output.process_id == process_id => {
1068 let bytes = output.chunk;
1069 match output.channel {
1070 StreamChannel::Stdout => {
1071 let _ = stdout_tx.send(bytes);
1072 }
1073 StreamChannel::Stderr => {
1074 let _ = stderr_tx.send(bytes);
1075 }
1076 }
1077 }
1078 EventPayload::ProcessExitedEvent(exited) if exited.process_id == process_id => {
1079 let _ = exit_tx.send(Some(exited.exit_code));
1080 break;
1081 }
1082 EventPayload::ProcessOutputEvent(_)
1083 | EventPayload::ProcessExitedEvent(_)
1084 | EventPayload::VmLifecycleEvent(_)
1085 | EventPayload::StructuredEvent(_)
1086 | EventPayload::ExtEnvelope(_) => {}
1087 }
1088 }
1089 let _guard = self.inner().process_registry_lock.lock();
1090 self.prune_exited_processes_locked(0);
1091 }
1092}
1093
1094fn build_process_forest(processes: Vec<ProcessInfo>) -> Vec<ProcessTreeNode> {
1101 use std::collections::BTreeMap as Map;
1102
1103 let pids: std::collections::BTreeSet<u32> = processes.iter().map(|p| p.pid).collect();
1104 let mut children_of: Map<u32, Vec<usize>> = Map::new();
1106 let mut roots: Vec<usize> = Vec::new();
1107 for (index, proc) in processes.iter().enumerate() {
1108 if pids.contains(&proc.ppid) {
1109 children_of.entry(proc.ppid).or_default().push(index);
1110 } else {
1111 roots.push(index);
1112 }
1113 }
1114
1115 fn build_node(
1116 index: usize,
1117 processes: &[ProcessInfo],
1118 children_of: &Map<u32, Vec<usize>>,
1119 seen: &mut std::collections::BTreeSet<usize>,
1120 ) -> ProcessTreeNode {
1121 let info = processes[index].clone();
1122 seen.insert(index);
1123 let child_indices: Vec<usize> = children_of
1124 .get(&info.pid)
1125 .map(|indices| {
1126 indices
1127 .iter()
1128 .copied()
1129 .filter(|child_index| !seen.contains(child_index))
1130 .collect()
1131 })
1132 .unwrap_or_default();
1133 let children = child_indices
1134 .into_iter()
1135 .map(|child_index| build_node(child_index, processes, children_of, seen))
1136 .collect();
1137 ProcessTreeNode { info, children }
1138 }
1139
1140 let mut seen = std::collections::BTreeSet::new();
1141 roots
1142 .into_iter()
1143 .map(|index| build_node(index, &processes, &children_of, &mut seen))
1144 .collect()
1145}
1146
1147fn stdin_to_bytes(input: StdinInput) -> Vec<u8> {
1150 match input {
1151 StdinInput::Text(text) => text.into_bytes(),
1152 StdinInput::Bytes(bytes) => bytes,
1153 }
1154}
1155
1156fn append_exec_output(
1157 buffer: &mut Vec<u8>,
1158 chunk: &[u8],
1159 captured_output_bytes: &mut usize,
1160 channel: &str,
1161) -> std::result::Result<(), ClientError> {
1162 let next_total = captured_output_bytes
1163 .checked_add(chunk.len())
1164 .ok_or_else(|| exec_output_limit_error(channel, usize::MAX))?;
1165 if next_total > EXEC_OUTPUT_CAPTURE_LIMIT_BYTES {
1166 return Err(exec_output_limit_error(channel, next_total));
1167 }
1168 buffer.extend_from_slice(chunk);
1169 *captured_output_bytes = next_total;
1170 Ok(())
1171}
1172
1173fn exec_output_limit_error(channel: &str, size: usize) -> ClientError {
1174 ClientError::Sidecar(format!(
1175 "exec {channel} capture is {size} bytes, limit is {EXEC_OUTPUT_CAPTURE_LIMIT_BYTES}"
1176 ))
1177}
1178
1179fn exited_pids_to_prune(mut entries: Vec<(u32, bool)>, target_len: usize) -> Vec<u32> {
1180 if entries.len() <= target_len {
1181 return Vec::new();
1182 }
1183 let mut remove_count = entries.len() - target_len;
1184 entries.sort_by_key(|(pid, _)| *pid);
1185 let mut out = Vec::new();
1186 for (pid, exited) in entries {
1187 if remove_count == 0 {
1188 break;
1189 }
1190 if !exited {
1191 continue;
1192 }
1193 out.push(pid);
1194 remove_count -= 1;
1195 }
1196 out
1197}
1198
1199fn prune_string_f64_map(map: &SccHashMap<String, f64>, limit: usize) {
1200 let mut keys = Vec::new();
1201 map.scan(|key, _| {
1202 keys.push(key.clone());
1203 });
1204 if keys.len() <= limit {
1205 return;
1206 }
1207 let remove_count = keys.len() - limit;
1208 keys.sort();
1209 for key in keys.into_iter().take(remove_count) {
1210 let _ = map.remove(&key);
1211 }
1212}
1213
1214pub(crate) fn install_output_callback(
1223 tx: broadcast::Sender<Vec<u8>>,
1224 mut callback: OutputCallback,
1225) -> JoinHandle<()> {
1226 let mut rx = tx.subscribe();
1227 tokio::spawn(async move {
1228 loop {
1229 match rx.recv().await {
1230 Ok(chunk) => callback(&chunk),
1231 Err(broadcast::error::RecvError::Lagged(_)) => continue,
1232 Err(broadcast::error::RecvError::Closed) => break,
1233 }
1234 }
1235 })
1236}
1237
1238pub(crate) fn drain_process_output_tasks(processes: &SccHashMap<u32, ProcessEntry>) {
1243 let mut tasks = Vec::new();
1244 processes.retain(|_, entry| {
1245 tasks.append(&mut entry.output_tasks);
1246 false
1247 });
1248 for task in tasks {
1249 task.abort();
1250 }
1251}
1252
1253fn epoch_ms_now() -> f64 {
1255 use std::time::{SystemTime, UNIX_EPOCH};
1256 SystemTime::now()
1257 .duration_since(UNIX_EPOCH)
1258 .map(|d| d.as_secs_f64() * 1000.0)
1259 .unwrap_or(0.0)
1260}
1261
1262#[cfg(test)]
1263mod tests {
1264 use super::{
1265 append_exec_output, drain_process_output_tasks, exited_pids_to_prune,
1266 install_output_callback, prune_string_f64_map, ExecOptions, OutputCallback,
1267 DEFAULT_EXEC_CWD, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES,
1268 };
1269 use crate::agent_os::ProcessEntry;
1270 use scc::HashMap as SccHashMap;
1271 use tokio::sync::{broadcast, watch};
1272
1273 #[tokio::test]
1278 async fn drain_process_output_tasks_clears_registry_and_aborts_tasks() {
1279 let processes: SccHashMap<u32, ProcessEntry> = SccHashMap::new();
1280
1281 let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(8);
1282 let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(8);
1283 let (exit_tx, _) = watch::channel::<Option<i32>>(None);
1284 let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
1285
1286 let task = tokio::spawn(async {
1289 loop {
1290 tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
1291 }
1292 });
1293 let abort_handle = task.abort_handle();
1294
1295 let entry = ProcessEntry {
1296 command: "sleep".to_string(),
1297 args: vec!["3600".to_string()],
1298 stdout_tx,
1299 stderr_tx,
1300 exit_tx,
1301 process_id: "proc-test".to_string(),
1302 kernel_pid: kernel_pid_tx,
1303 output_tasks: vec![task],
1304 started_at: 0,
1305 };
1306 let _ = processes.insert(1, entry);
1307
1308 assert!(!abort_handle.is_finished(), "task should start alive");
1309
1310 drain_process_output_tasks(&processes);
1311
1312 assert!(processes.is_empty(), "registry must be cleared on drain");
1313
1314 for _ in 0..100 {
1316 if abort_handle.is_finished() {
1317 break;
1318 }
1319 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1320 }
1321 assert!(
1322 abort_handle.is_finished(),
1323 "output task must be aborted after drain"
1324 );
1325 }
1326
1327 #[tokio::test]
1333 async fn install_output_callback_handle_is_captured_into_process_entry() {
1334 use std::sync::atomic::{AtomicUsize, Ordering};
1335 use std::sync::Arc;
1336
1337 let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(8);
1338 let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(8);
1339 let (exit_tx, _) = watch::channel::<Option<i32>>(None);
1340 let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
1341
1342 let calls = Arc::new(AtomicUsize::new(0));
1343 let calls_cb = Arc::clone(&calls);
1344 let cb: OutputCallback = Box::new(move |_chunk: &[u8]| {
1345 calls_cb.fetch_add(1, Ordering::SeqCst);
1346 });
1347
1348 let output_tasks = vec![install_output_callback(stdout_tx.clone(), cb)];
1350
1351 let entry = ProcessEntry {
1352 command: "sleep".to_string(),
1353 args: vec!["3600".to_string()],
1354 stdout_tx: stdout_tx.clone(),
1355 stderr_tx,
1356 exit_tx,
1357 process_id: "proc-test".to_string(),
1358 kernel_pid: kernel_pid_tx,
1359 output_tasks,
1360 started_at: 0,
1361 };
1362
1363 assert_eq!(
1364 entry.output_tasks.len(),
1365 1,
1366 "the install_output_callback handle must be captured on the entry"
1367 );
1368
1369 stdout_tx
1371 .send(b"hello".to_vec())
1372 .expect("broadcast send to subscribed callback task");
1373 for _ in 0..100 {
1374 if calls.load(Ordering::SeqCst) > 0 {
1375 break;
1376 }
1377 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1378 }
1379 assert_eq!(
1380 calls.load(Ordering::SeqCst),
1381 1,
1382 "the stored handle must drive the registered callback"
1383 );
1384
1385 let processes: SccHashMap<u32, ProcessEntry> = SccHashMap::new();
1387 let _ = processes.insert(1, entry);
1388 drain_process_output_tasks(&processes);
1389 assert!(processes.is_empty(), "registry must be cleared on drain");
1390 }
1391
1392 #[test]
1393 fn exec_options_default_uses_workspace_cwd() {
1394 assert_eq!(
1395 ExecOptions::default().cwd.as_deref(),
1396 Some(DEFAULT_EXEC_CWD)
1397 );
1398 }
1399
1400 #[test]
1401 fn append_exec_output_rejects_capture_over_limit() {
1402 let mut buffer = vec![0u8; EXEC_OUTPUT_CAPTURE_LIMIT_BYTES - 1];
1403 let mut captured = buffer.len();
1404
1405 append_exec_output(&mut buffer, &[1], &mut captured, "stdout")
1406 .expect("chunk at limit should fit");
1407 assert_eq!(captured, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1408
1409 let error = append_exec_output(&mut buffer, &[2], &mut captured, "stdout")
1410 .expect_err("chunk over limit should fail");
1411 assert!(
1412 error.to_string().contains("exec stdout capture is"),
1413 "unexpected error: {error}"
1414 );
1415 assert_eq!(captured, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1416 assert_eq!(buffer.len(), EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1417 }
1418
1419 #[test]
1420 fn exited_pid_pruning_keeps_live_entries_and_removes_oldest_exited() {
1421 let pids = exited_pids_to_prune(vec![(3, true), (1, false), (2, true), (4, true)], 2);
1422 assert_eq!(pids, vec![2, 3]);
1423 }
1424
1425 #[test]
1426 fn observed_time_pruning_enforces_limit() {
1427 let map = SccHashMap::new();
1428 let _ = map.insert("b".to_string(), 2.0);
1429 let _ = map.insert("a".to_string(), 1.0);
1430 let _ = map.insert("c".to_string(), 3.0);
1431
1432 prune_string_f64_map(&map, 2);
1433
1434 assert!(map.read("a", |_, _| ()).is_none());
1435 assert!(map.read("b", |_, _| ()).is_some());
1436 assert!(map.read("c", |_, _| ()).is_some());
1437 }
1438}