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::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 env: BTreeMap<String, String>,
125 pub cwd: Option<String>,
126 pub stdio: Option<SpawnStdio>,
127 pub stdin_fd: Option<i32>,
128 pub stdout_fd: Option<i32>,
129 pub stderr_fd: Option<i32>,
130 pub stream_stdin: Option<bool>,
131}
132
133#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
134#[serde(rename_all = "lowercase")]
135pub enum ProcessStream {
136 Stdout,
137 Stderr,
138}
139
140#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
141pub struct ProcessOutput {
142 pub pid: u32,
143 pub stream: ProcessStream,
144 pub data: Vec<u8>,
145}
146
147#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
148pub struct ProcessExit {
149 pub pid: u32,
150 #[serde(rename = "exitCode")]
151 pub exit_code: i32,
152}
153
154#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
156pub struct SpawnedProcessInfo {
157 pub pid: u32,
158 pub command: String,
159 pub args: Vec<String>,
160 pub running: bool,
161 #[serde(rename = "exitCode")]
162 pub exit_code: Option<i32>,
163 #[serde(rename = "startedAt")]
165 pub started_at: i64,
166}
167
168#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
170pub struct SpawnHandle {
171 pub pid: u32,
172}
173
174#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
176#[serde(rename_all = "lowercase")]
177pub enum ProcessStatus {
178 Running,
179 Exited,
180}
181
182#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
184pub struct ProcessInfo {
185 pub pid: u32,
186 pub ppid: u32,
187 pub pgid: u32,
188 pub sid: u32,
189 pub driver: String,
190 pub command: String,
191 pub args: Vec<String>,
192 pub cwd: String,
193 pub status: ProcessStatus,
194 #[serde(rename = "exitCode")]
195 pub exit_code: Option<i32>,
196 #[serde(rename = "startTime")]
197 pub start_time: f64,
198 #[serde(rename = "exitTime")]
199 pub exit_time: Option<f64>,
200}
201
202#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
204pub struct ProcessTreeNode {
205 #[serde(flatten)]
206 pub info: ProcessInfo,
207 pub children: Vec<ProcessTreeNode>,
208}
209
210impl AgentOs {
215 pub async fn exec(&self, command: &str, options: ExecOptions) -> Result<ExecResult> {
220 let (resolved_command, resolved_args) = resolve_exec_command(command)?;
224 self.exec_argv(&resolved_command, &resolved_args, options)
225 .await
226 }
227
228 pub async fn exec_argv(
234 &self,
235 command: &str,
236 args: &[String],
237 mut options: ExecOptions,
238 ) -> Result<ExecResult> {
239 let process_id = self.next_process_id();
240
241 let mut events = self.transport().subscribe_wire_events();
244
245 let resolved_command = command.to_owned();
246 let resolved_args = args.to_vec();
247 let started = self
248 .send_execute(
249 &process_id,
250 Some(resolved_command),
251 resolved_args,
252 options.env.clone(),
253 options.cwd.clone(),
254 )
255 .await
256 .context("exec: Execute request failed")?;
257 debug_assert_eq!(started.process_id, process_id);
258
259 if let Some(stdin) = options.stdin.take() {
262 let chunk = stdin_to_bytes(stdin);
263 let ownership = self.vm_scope();
264 let _ = self
265 .transport()
266 .request_wire(
267 ownership,
268 wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {
269 process_id: process_id.clone(),
270 chunk,
271 }),
272 )
273 .await;
274 }
275 {
276 let ownership = self.vm_scope();
277 let _ = self
278 .transport()
279 .request_wire(
280 ownership,
281 wire::RequestPayload::CloseStdinRequest(wire::CloseStdinRequest {
282 process_id: process_id.clone(),
283 }),
284 )
285 .await;
286 }
287
288 let mut on_stdout = options.on_stdout.take();
289 let mut on_stderr = options.on_stderr.take();
290
291 let timeout_deadline = options
295 .timeout
296 .filter(|ms| ms.is_finite() && *ms >= 0.0)
297 .map(|ms| {
298 tokio::time::Instant::now() + std::time::Duration::from_secs_f64(ms / 1000.0)
299 });
300 let mut killed_for_timeout = false;
301
302 let capture_stdio = options.capture_stdio.unwrap_or(true);
303 let mut stdout = Vec::<u8>::new();
304 let mut stderr = Vec::<u8>::new();
305 let mut captured_output_bytes = 0usize;
306 let mut capture_error: Option<ClientError> = None;
307 let exit_code = loop {
308 let recv = events.recv();
309 let frame = match timeout_deadline {
310 Some(deadline) => {
311 tokio::select! {
312 result = recv => result,
313 _ = tokio::time::sleep_until(deadline), if !killed_for_timeout => {
314 killed_for_timeout = true;
315 self.kill_wire_process(&process_id, "SIGKILL");
316 continue;
317 }
318 }
319 }
320 None => recv.await,
321 };
322 let (_, payload) = match frame {
323 Ok(frame) => frame,
324 Err(broadcast::error::RecvError::Lagged(_)) => continue,
325 Err(broadcast::error::RecvError::Closed) => {
326 return Err(ClientError::Sidecar(
327 "exec: event stream closed before process exit".to_owned(),
328 )
329 .into());
330 }
331 };
332 match payload {
333 EventPayload::ProcessOutputEvent(output) if output.process_id == process_id => {
334 match output.channel {
335 StreamChannel::Stdout => {
336 if let Some(cb) = on_stdout.as_mut() {
337 cb(&output.chunk);
338 }
339 if capture_stdio && capture_error.is_none() {
340 match append_exec_output(
341 &mut stdout,
342 &output.chunk,
343 &mut captured_output_bytes,
344 "stdout",
345 ) {
346 Ok(()) => {}
347 Err(error) => {
348 self.kill_wire_process(&process_id, "SIGKILL");
349 capture_error = Some(error);
350 }
351 }
352 }
353 }
354 StreamChannel::Stderr => {
355 if let Some(cb) = on_stderr.as_mut() {
356 cb(&output.chunk);
357 }
358 if capture_stdio && capture_error.is_none() {
359 match append_exec_output(
360 &mut stderr,
361 &output.chunk,
362 &mut captured_output_bytes,
363 "stderr",
364 ) {
365 Ok(()) => {}
366 Err(error) => {
367 self.kill_wire_process(&process_id, "SIGKILL");
368 capture_error = Some(error);
369 }
370 }
371 }
372 }
373 }
374 }
375 EventPayload::ProcessExitedEvent(exited) if exited.process_id == process_id => {
376 break exited.exit_code;
377 }
378 EventPayload::ProcessOutputEvent(_)
379 | EventPayload::ProcessExitedEvent(_)
380 | EventPayload::VmLifecycleEvent(_)
381 | EventPayload::StructuredEvent(_)
382 | EventPayload::ExtEnvelope(_) => {}
383 }
384 };
385
386 if let Some(error) = capture_error {
387 return Err(error.into());
388 }
389
390 Ok(ExecResult {
391 exit_code,
392 stdout: String::from_utf8_lossy(&stdout).into_owned(),
393 stderr: String::from_utf8_lossy(&stderr).into_owned(),
394 })
395 }
396
397 pub fn spawn(
401 &self,
402 command: &str,
403 args: Vec<String>,
404 options: SpawnOptions,
405 ) -> Result<SpawnHandle> {
406 let registry_guard = self.inner().process_registry_lock.lock();
407 self.prune_exited_processes_locked(1);
408 if self.process_registry_len_locked() >= PROCESS_REGISTRY_LIMIT {
409 return Err(ClientError::Sidecar(format!(
410 "process registry limit exceeded: at most {PROCESS_REGISTRY_LIMIT} processes can be tracked per VM"
411 ))
412 .into());
413 }
414
415 let pid = self
418 .inner()
419 .synthetic_pid_counter
420 .fetch_add(1, Ordering::SeqCst) as u32;
421 let process_id = format!("proc-{pid}-{}", uuid::Uuid::new_v4());
422
423 let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(PROCESS_STREAM_CAPACITY);
424 let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(PROCESS_STREAM_CAPACITY);
425 let (output_tx, _) = broadcast::channel::<ProcessOutput>(PROCESS_STREAM_CAPACITY);
426 let (exit_tx, _) = watch::channel::<Option<i32>>(None);
429 let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
432
433 let entry = ProcessEntry {
434 command: command.to_owned(),
435 args: args.clone(),
436 stdout_tx: stdout_tx.clone(),
437 stderr_tx: stderr_tx.clone(),
438 output_tx: output_tx.clone(),
439 exit_tx: exit_tx.clone(),
440 process_id: process_id.clone(),
441 kernel_pid: kernel_pid_tx.clone(),
442 output_tasks: Vec::new(),
443 started_at: epoch_ms_now() as i64,
444 };
445 let _ = self.inner().processes.insert(pid, entry);
448 drop(registry_guard);
449
450 let events = self.transport().subscribe_wire_events();
452
453 let this = self.clone();
454 let command = command.to_owned();
455 tokio::spawn(async move {
456 this.run_spawn(
457 pid,
458 process_id,
459 command,
460 args,
461 options,
462 events,
463 stdout_tx,
464 stderr_tx,
465 output_tx,
466 exit_tx,
467 kernel_pid_tx,
468 )
469 .await;
470 });
471
472 Ok(SpawnHandle { pid })
473 }
474
475 pub fn write_process_stdin(
477 &self,
478 pid: u32,
479 data: StdinInput,
480 ) -> std::result::Result<(), ClientError> {
481 let process_id = self.lookup_process_id(pid)?;
482 let chunk: Vec<u8> = stdin_to_bytes(data);
483 let this = self.clone();
484 tokio::spawn(async move {
486 let ownership = this.vm_scope();
487 let _ = this
488 .transport()
489 .request_wire(
490 ownership,
491 wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {
492 process_id,
493 chunk,
494 }),
495 )
496 .await;
497 });
498 Ok(())
499 }
500
501 pub fn close_process_stdin(&self, pid: u32) -> std::result::Result<(), ClientError> {
503 let process_id = self.lookup_process_id(pid)?;
504 let this = self.clone();
505 tokio::spawn(async move {
506 let ownership = this.vm_scope();
507 let _ = this
508 .transport()
509 .request_wire(
510 ownership,
511 wire::RequestPayload::CloseStdinRequest(wire::CloseStdinRequest { process_id }),
512 )
513 .await;
514 });
515 Ok(())
516 }
517
518 pub fn on_process_output(
520 &self,
521 pid: u32,
522 mut handler: impl FnMut(ProcessOutput) + Send + 'static,
523 ) -> std::result::Result<Subscription, ClientError> {
524 let mut rx = self
525 .inner()
526 .processes
527 .read(&pid, |_, entry| entry.output_tx.subscribe())
528 .ok_or(ClientError::ProcessNotFound(pid))?;
529 let task = tokio::spawn(async move {
530 loop {
531 match rx.recv().await {
532 Ok(event) => handler(event),
533 Err(broadcast::error::RecvError::Lagged(_)) => continue,
534 Err(broadcast::error::RecvError::Closed) => return,
535 }
536 }
537 });
538 Ok(Subscription::new(move || task.abort()))
539 }
540
541 pub fn on_process_exit(
546 &self,
547 pid: u32,
548 handler: impl FnOnce(ProcessExit) + Send + 'static,
549 ) -> std::result::Result<Subscription, ClientError> {
550 let mut rx = self
551 .inner()
552 .processes
553 .read(&pid, |_, entry| entry.exit_tx.subscribe())
554 .ok_or(ClientError::ProcessNotFound(pid))?;
555
556 if let Some(code) = *rx.borrow() {
558 handler(ProcessExit {
559 pid,
560 exit_code: code,
561 });
562 return Ok(Subscription::noop());
563 }
564
565 let task = tokio::spawn(async move {
568 while rx.changed().await.is_ok() {
569 if let Some(code) = *rx.borrow() {
570 handler(ProcessExit {
571 pid,
572 exit_code: code,
573 });
574 return;
575 }
576 }
577 });
578 Ok(Subscription::new(move || task.abort()))
579 }
580
581 pub async fn wait_process(&self, pid: u32) -> std::result::Result<i32, ClientError> {
584 let mut rx = self
585 .inner()
586 .processes
587 .read(&pid, |_, entry| entry.exit_tx.subscribe())
588 .ok_or(ClientError::ProcessNotFound(pid))?;
589
590 if let Some(code) = *rx.borrow() {
591 return Ok(code);
592 }
593 while rx.changed().await.is_ok() {
594 if let Some(code) = *rx.borrow() {
595 return Ok(code);
596 }
597 }
598 Err(ClientError::Sidecar(format!(
599 "wait_process: exit channel closed before process {pid} reported an exit code"
600 )))
601 }
602
603 pub fn list_processes(&self) -> Vec<SpawnedProcessInfo> {
605 let mut out = Vec::new();
606 self.inner().processes.scan(|pid, entry| {
607 let exit_code = *entry.exit_tx.borrow();
608 out.push(SpawnedProcessInfo {
609 pid: *pid,
610 command: entry.command.clone(),
611 args: entry.args.clone(),
612 running: exit_code.is_none(),
613 exit_code,
614 started_at: entry.started_at,
615 });
616 });
617 out
618 }
619
620 pub async fn all_processes(&self) -> Result<Vec<ProcessInfo>> {
628 let ownership = self.vm_scope();
629 let response = self
630 .transport()
631 .request_wire(ownership, wire::RequestPayload::GetProcessSnapshotRequest)
632 .await
633 .context("all_processes: GetProcessSnapshot request failed")?;
634 let snapshot = match response {
635 wire::ResponsePayload::ProcessSnapshotResponse(snapshot) => snapshot,
636 wire::ResponsePayload::RejectedResponse(rejected) => {
637 return Err(ClientError::from_rejection(rejected).into());
638 }
639 other => {
640 return Err(ClientError::Sidecar(format!(
641 "all_processes: unexpected response {other:?}"
642 ))
643 .into());
644 }
645 };
646
647 struct Tracked {
651 exit_code: Option<i32>,
652 command: String,
653 args: Vec<String>,
654 }
655 let mut tracked_by_process_id: BTreeMap<String, Tracked> = BTreeMap::new();
656 let mut display_pid_by_kernel_pid: BTreeMap<u32, u32> = BTreeMap::new();
657 self.inner().processes.scan(|display_pid, entry| {
658 let exit_code = *entry.exit_tx.borrow();
659 if let Some(kernel_pid) = *entry.kernel_pid.borrow() {
660 display_pid_by_kernel_pid.insert(kernel_pid, *display_pid);
661 }
662 tracked_by_process_id.insert(
663 entry.process_id.clone(),
664 Tracked {
665 exit_code,
666 command: entry.command.clone(),
667 args: entry.args.clone(),
668 },
669 );
670 });
671
672 let now_ms = epoch_ms_now();
673 let mut seen_display_pids: std::collections::BTreeSet<u32> =
674 std::collections::BTreeSet::new();
675 let mut out: Vec<ProcessInfo> = Vec::new();
676
677 for entry in snapshot.processes {
678 let tracked = tracked_by_process_id.get(&entry.process_id);
679 let display_pid = display_pid_by_kernel_pid
680 .get(&entry.pid)
681 .copied()
682 .unwrap_or(entry.pid);
683 let display_ppid = display_pid_by_kernel_pid
684 .get(&entry.ppid)
685 .copied()
686 .unwrap_or(entry.ppid);
687 let display_pgid = display_pid_by_kernel_pid
688 .get(&entry.pgid)
689 .copied()
690 .unwrap_or(entry.pgid);
691 let display_sid = display_pid_by_kernel_pid
692 .get(&entry.sid)
693 .copied()
694 .unwrap_or(entry.sid);
695
696 let process_key = format!("{}:{}", entry.process_id, entry.pid);
698 let start_time = self.observed_start_time(&process_key, now_ms);
699
700 let (status, exit_code) = match tracked {
703 Some(t) => match t.exit_code {
704 Some(code) => (ProcessStatus::Exited, Some(code)),
705 None => (ProcessStatus::Running, entry.exit_code),
706 },
707 None => {
708 let status = match entry.status {
709 ProcessSnapshotStatus::Running | ProcessSnapshotStatus::Stopped => {
710 ProcessStatus::Running
711 }
712 ProcessSnapshotStatus::Exited => ProcessStatus::Exited,
713 };
714 (status, entry.exit_code)
715 }
716 };
717
718 let exit_time = match (tracked, status) {
720 (Some(_), ProcessStatus::Exited) => {
721 Some(self.observed_exit_time(&entry.process_id, now_ms))
722 }
723 _ => None,
724 };
725
726 let (command, args) = match tracked {
727 Some(t) => (t.command.clone(), t.args.clone()),
728 None => (entry.command, entry.args),
729 };
730
731 seen_display_pids.insert(display_pid);
732 out.push(ProcessInfo {
733 pid: display_pid,
734 ppid: display_ppid,
735 pgid: display_pgid,
736 sid: display_sid,
737 driver: entry.driver,
738 command,
739 args,
740 cwd: entry.cwd,
741 status,
742 exit_code,
743 start_time,
744 exit_time,
745 });
746 }
747
748 self.inner().processes.scan(|display_pid, entry| {
751 if seen_display_pids.contains(display_pid) {
752 return;
753 }
754 let exit_code = *entry.exit_tx.borrow();
755 let process_key = format!("{}:{}", entry.process_id, display_pid);
756 let start_time = self.observed_start_time(&process_key, now_ms);
757 let (status, exit_time) = match exit_code {
758 Some(_) => (
759 ProcessStatus::Exited,
760 Some(self.observed_exit_time(&entry.process_id, now_ms)),
761 ),
762 None => (ProcessStatus::Running, None),
763 };
764 out.push(ProcessInfo {
765 pid: *display_pid,
766 ppid: 0,
767 pgid: *display_pid,
768 sid: *display_pid,
769 driver: String::new(),
770 command: entry.command.clone(),
771 args: entry.args.clone(),
772 cwd: String::new(),
773 status,
774 exit_code,
775 start_time,
776 exit_time,
777 });
778 });
779
780 out.sort_by_key(|info| info.pid);
781 Ok(out)
782 }
783
784 fn observed_start_time(&self, process_key: &str, now_ms: f64) -> f64 {
787 let _guard = self.inner().observed_process_time_lock.lock();
788 if let Some(existing) = self
789 .inner()
790 .observed_process_start_times
791 .read(process_key, |_, value| *value)
792 {
793 return existing;
794 }
795 let _ = self
796 .inner()
797 .observed_process_start_times
798 .insert(process_key.to_owned(), now_ms);
799 prune_string_f64_map(
800 &self.inner().observed_process_start_times,
801 OBSERVED_PROCESS_TIME_LIMIT,
802 );
803 self.inner()
806 .observed_process_start_times
807 .read(process_key, |_, value| *value)
808 .unwrap_or(now_ms)
809 }
810
811 fn observed_exit_time(&self, process_id: &str, now_ms: f64) -> f64 {
813 let _guard = self.inner().observed_process_time_lock.lock();
814 if let Some(existing) = self
815 .inner()
816 .observed_process_exit_times
817 .read(process_id, |_, value| *value)
818 {
819 return existing;
820 }
821 let _ = self
822 .inner()
823 .observed_process_exit_times
824 .insert(process_id.to_owned(), now_ms);
825 prune_string_f64_map(
826 &self.inner().observed_process_exit_times,
827 OBSERVED_PROCESS_TIME_LIMIT,
828 );
829 self.inner()
830 .observed_process_exit_times
831 .read(process_id, |_, value| *value)
832 .unwrap_or(now_ms)
833 }
834
835 pub async fn process_tree(&self) -> Result<Vec<ProcessTreeNode>> {
837 let processes = self.all_processes().await?;
838 Ok(build_process_forest(processes))
839 }
840
841 pub fn get_process(&self, pid: u32) -> std::result::Result<SpawnedProcessInfo, ClientError> {
843 self.inner()
844 .processes
845 .read(&pid, |pid, entry| {
846 let exit_code = *entry.exit_tx.borrow();
847 SpawnedProcessInfo {
848 pid: *pid,
849 command: entry.command.clone(),
850 args: entry.args.clone(),
851 running: exit_code.is_none(),
852 exit_code,
853 started_at: entry.started_at,
854 }
855 })
856 .ok_or(ClientError::ProcessNotFound(pid))
857 }
858
859 pub fn stop_process(&self, pid: u32) -> std::result::Result<(), ClientError> {
861 self.signal_process(pid, "SIGTERM")
862 }
863
864 pub fn kill_process(&self, pid: u32) -> std::result::Result<(), ClientError> {
866 self.signal_process(pid, "SIGKILL")
867 }
868
869 fn vm_scope(&self) -> wire::OwnershipScope {
875 wire::OwnershipScope::VmOwnership(wire::VmOwnership {
876 connection_id: self.connection_id().to_string(),
877 session_id: self.wire_session_id().to_string(),
878 vm_id: self.vm_id().to_string(),
879 })
880 }
881
882 fn next_process_id(&self) -> String {
884 let n = self.inner().process_counter.fetch_add(1, Ordering::SeqCst);
885 format!("proc-{n}-{}", uuid::Uuid::new_v4())
886 }
887
888 fn lookup_process_id(&self, pid: u32) -> std::result::Result<String, ClientError> {
890 self.inner()
891 .processes
892 .read(&pid, |_, entry| entry.process_id.clone())
893 .ok_or(ClientError::ProcessNotFound(pid))
894 }
895
896 async fn send_execute(
898 &self,
899 process_id: &str,
900 command: Option<String>,
901 args: Vec<String>,
902 env: BTreeMap<String, String>,
903 cwd: Option<String>,
904 ) -> std::result::Result<wire::ProcessStartedResponse, ClientError> {
905 let ownership = self.vm_scope();
906 let response = self
907 .transport()
908 .request_wire(
909 ownership,
910 wire::RequestPayload::ExecuteRequest(wire::ExecuteRequest {
911 process_id: process_id.to_owned(),
912 command,
913 runtime: None,
914 entrypoint: None,
915 args,
916 env: env.into_iter().collect(),
917 cwd,
918 wasm_permission_tier: None,
919 }),
920 )
921 .await?;
922 match response {
923 wire::ResponsePayload::ProcessStartedResponse(started) => Ok(started),
924 wire::ResponsePayload::RejectedResponse(rejected) => {
925 Err(ClientError::from_rejection(rejected))
926 }
927 other => Err(ClientError::Sidecar(format!(
928 "Execute: unexpected response {other:?}"
929 ))),
930 }
931 }
932
933 fn kill_wire_process(&self, process_id: &str, signal: &str) {
936 let process_id = process_id.to_owned();
937 let signal = signal.to_owned();
938 let this = self.clone();
939 tokio::spawn(async move {
940 let ownership = this.vm_scope();
941 let _ = this
942 .transport()
943 .request_wire(
944 ownership,
945 wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
946 process_id,
947 signal,
948 }),
949 )
950 .await;
951 });
952 }
953
954 fn signal_process(&self, pid: u32, signal: &str) -> std::result::Result<(), ClientError> {
957 let (process_id, already_exited) = self
958 .inner()
959 .processes
960 .read(&pid, |_, entry| {
961 (entry.process_id.clone(), entry.exit_tx.borrow().is_some())
962 })
963 .ok_or(ClientError::ProcessNotFound(pid))?;
964 if already_exited {
965 return Ok(());
966 }
967 let signal = signal.to_owned();
968 let this = self.clone();
969 tokio::spawn(async move {
970 let ownership = this.vm_scope();
971 let _ = this
972 .transport()
973 .request_wire(
974 ownership,
975 wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
976 process_id,
977 signal,
978 }),
979 )
980 .await;
981 });
982 Ok(())
983 }
984
985 fn process_registry_len_locked(&self) -> usize {
986 let mut count = 0usize;
987 self.inner().processes.scan(|_, _| {
988 count += 1;
989 });
990 count
991 }
992
993 fn prune_exited_processes_locked(&self, reserve_slots: usize) {
994 let mut entries = Vec::new();
995 self.inner().processes.scan(|pid, entry| {
996 entries.push((*pid, entry.exit_tx.borrow().is_some()));
997 });
998 let target_len = PROCESS_REGISTRY_LIMIT.saturating_sub(reserve_slots);
999 if entries.len() <= target_len {
1000 return;
1001 }
1002
1003 for pid in exited_pids_to_prune(entries, target_len) {
1004 self.remove_process_tracking_locked(pid);
1005 }
1006 }
1007
1008 fn remove_process_tracking_locked(&self, pid: u32) {
1009 if let Some((_, entry)) = self.inner().processes.remove(&pid) {
1010 let _time_guard = self.inner().observed_process_time_lock.lock();
1011 let _ = self
1012 .inner()
1013 .observed_process_exit_times
1014 .remove(&entry.process_id);
1015 let fallback_start_key = format!("{}:{pid}", entry.process_id);
1016 let _ = self
1017 .inner()
1018 .observed_process_start_times
1019 .remove(&fallback_start_key);
1020 if let Some(kernel_pid) = *entry.kernel_pid.borrow() {
1021 let start_key = format!("{}:{kernel_pid}", entry.process_id);
1022 let _ = self.inner().observed_process_start_times.remove(&start_key);
1023 }
1024 }
1025 }
1026
1027 #[allow(clippy::too_many_arguments)]
1032 async fn run_spawn(
1033 self,
1034 pid: u32,
1035 process_id: String,
1036 command: String,
1037 args: Vec<String>,
1038 options: SpawnOptions,
1039 mut events: broadcast::Receiver<(wire::OwnershipScope, EventPayload)>,
1040 stdout_tx: broadcast::Sender<Vec<u8>>,
1041 stderr_tx: broadcast::Sender<Vec<u8>>,
1042 output_tx: broadcast::Sender<ProcessOutput>,
1043 exit_tx: watch::Sender<Option<i32>>,
1044 kernel_pid_tx: watch::Sender<Option<u32>>,
1045 ) {
1046 match self
1047 .send_execute(
1048 &process_id,
1049 Some(command),
1050 args,
1051 options.env.clone(),
1052 options.cwd.clone(),
1053 )
1054 .await
1055 {
1056 Ok(started) => {
1057 if let Some(kernel_pid) = started.pid {
1060 let _ = kernel_pid_tx.send(Some(kernel_pid));
1061 }
1062 }
1063 Err(error) => {
1064 let message = format!("{error}\n");
1068 let bytes = message.into_bytes();
1069 let _ = stderr_tx.send(bytes.clone());
1070 let _ = output_tx.send(ProcessOutput {
1071 pid,
1072 stream: ProcessStream::Stderr,
1073 data: bytes,
1074 });
1075 tracing::error!(?error, pid, %process_id, "spawn: Execute request failed");
1076 let _ = exit_tx.send(Some(1));
1077 let _guard = self.inner().process_registry_lock.lock();
1078 self.prune_exited_processes_locked(0);
1079 return;
1080 }
1081 }
1082
1083 loop {
1084 let (_, payload) = match events.recv().await {
1085 Ok(frame) => frame,
1086 Err(broadcast::error::RecvError::Lagged(_)) => continue,
1087 Err(broadcast::error::RecvError::Closed) => {
1088 let _ = exit_tx.send(Some(0));
1092 break;
1093 }
1094 };
1095 match payload {
1096 EventPayload::ProcessOutputEvent(output) if output.process_id == process_id => {
1097 let bytes = output.chunk;
1098 let _ = output_tx.send(ProcessOutput {
1099 pid,
1100 stream: match output.channel {
1101 StreamChannel::Stdout => ProcessStream::Stdout,
1102 StreamChannel::Stderr => ProcessStream::Stderr,
1103 },
1104 data: bytes.clone(),
1105 });
1106 match output.channel {
1107 StreamChannel::Stdout => {
1108 let _ = stdout_tx.send(bytes);
1109 }
1110 StreamChannel::Stderr => {
1111 let _ = stderr_tx.send(bytes);
1112 }
1113 }
1114 }
1115 EventPayload::ProcessExitedEvent(exited) if exited.process_id == process_id => {
1116 let _ = exit_tx.send(Some(exited.exit_code));
1117 break;
1118 }
1119 EventPayload::ProcessOutputEvent(_)
1120 | EventPayload::ProcessExitedEvent(_)
1121 | EventPayload::VmLifecycleEvent(_)
1122 | EventPayload::StructuredEvent(_)
1123 | EventPayload::ExtEnvelope(_) => {}
1124 }
1125 }
1126 let _guard = self.inner().process_registry_lock.lock();
1127 self.prune_exited_processes_locked(0);
1128 }
1129}
1130
1131fn build_process_forest(processes: Vec<ProcessInfo>) -> Vec<ProcessTreeNode> {
1138 use std::collections::BTreeMap as Map;
1139
1140 let pids: std::collections::BTreeSet<u32> = processes.iter().map(|p| p.pid).collect();
1141 let mut children_of: Map<u32, Vec<usize>> = Map::new();
1143 let mut roots: Vec<usize> = Vec::new();
1144 for (index, proc) in processes.iter().enumerate() {
1145 if pids.contains(&proc.ppid) {
1146 children_of.entry(proc.ppid).or_default().push(index);
1147 } else {
1148 roots.push(index);
1149 }
1150 }
1151
1152 fn build_node(
1153 index: usize,
1154 processes: &[ProcessInfo],
1155 children_of: &Map<u32, Vec<usize>>,
1156 seen: &mut std::collections::BTreeSet<usize>,
1157 ) -> ProcessTreeNode {
1158 let info = processes[index].clone();
1159 seen.insert(index);
1160 let child_indices: Vec<usize> = children_of
1161 .get(&info.pid)
1162 .map(|indices| {
1163 indices
1164 .iter()
1165 .copied()
1166 .filter(|child_index| !seen.contains(child_index))
1167 .collect()
1168 })
1169 .unwrap_or_default();
1170 let children = child_indices
1171 .into_iter()
1172 .map(|child_index| build_node(child_index, processes, children_of, seen))
1173 .collect();
1174 ProcessTreeNode { info, children }
1175 }
1176
1177 let mut seen = std::collections::BTreeSet::new();
1178 roots
1179 .into_iter()
1180 .map(|index| build_node(index, &processes, &children_of, &mut seen))
1181 .collect()
1182}
1183
1184fn stdin_to_bytes(input: StdinInput) -> Vec<u8> {
1187 match input {
1188 StdinInput::Text(text) => text.into_bytes(),
1189 StdinInput::Bytes(bytes) => bytes,
1190 }
1191}
1192
1193fn append_exec_output(
1194 buffer: &mut Vec<u8>,
1195 chunk: &[u8],
1196 captured_output_bytes: &mut usize,
1197 channel: &str,
1198) -> std::result::Result<(), ClientError> {
1199 let next_total = captured_output_bytes
1200 .checked_add(chunk.len())
1201 .ok_or_else(|| exec_output_limit_error(channel, usize::MAX))?;
1202 if next_total > EXEC_OUTPUT_CAPTURE_LIMIT_BYTES {
1203 return Err(exec_output_limit_error(channel, next_total));
1204 }
1205 buffer.extend_from_slice(chunk);
1206 *captured_output_bytes = next_total;
1207 Ok(())
1208}
1209
1210fn exec_output_limit_error(channel: &str, size: usize) -> ClientError {
1211 ClientError::Sidecar(format!(
1212 "exec {channel} capture is {size} bytes, limit is {EXEC_OUTPUT_CAPTURE_LIMIT_BYTES}"
1213 ))
1214}
1215
1216fn exited_pids_to_prune(mut entries: Vec<(u32, bool)>, target_len: usize) -> Vec<u32> {
1217 if entries.len() <= target_len {
1218 return Vec::new();
1219 }
1220 let mut remove_count = entries.len() - target_len;
1221 entries.sort_by_key(|(pid, _)| *pid);
1222 let mut out = Vec::new();
1223 for (pid, exited) in entries {
1224 if remove_count == 0 {
1225 break;
1226 }
1227 if !exited {
1228 continue;
1229 }
1230 out.push(pid);
1231 remove_count -= 1;
1232 }
1233 out
1234}
1235
1236fn prune_string_f64_map(map: &SccHashMap<String, f64>, limit: usize) {
1237 let mut keys = Vec::new();
1238 map.scan(|key, _| {
1239 keys.push(key.clone());
1240 });
1241 if keys.len() <= limit {
1242 return;
1243 }
1244 let remove_count = keys.len() - limit;
1245 keys.sort();
1246 for key in keys.into_iter().take(remove_count) {
1247 let _ = map.remove(&key);
1248 }
1249}
1250
1251pub(crate) fn install_output_callback(
1260 tx: broadcast::Sender<Vec<u8>>,
1261 mut callback: OutputCallback,
1262) -> JoinHandle<()> {
1263 let mut rx = tx.subscribe();
1264 tokio::spawn(async move {
1265 loop {
1266 match rx.recv().await {
1267 Ok(chunk) => callback(&chunk),
1268 Err(broadcast::error::RecvError::Lagged(_)) => continue,
1269 Err(broadcast::error::RecvError::Closed) => break,
1270 }
1271 }
1272 })
1273}
1274
1275pub(crate) fn drain_process_output_tasks(processes: &SccHashMap<u32, ProcessEntry>) {
1280 let mut tasks = Vec::new();
1281 processes.retain(|_, entry| {
1282 tasks.append(&mut entry.output_tasks);
1283 false
1284 });
1285 for task in tasks {
1286 task.abort();
1287 }
1288}
1289
1290fn epoch_ms_now() -> f64 {
1292 use std::time::{SystemTime, UNIX_EPOCH};
1293 SystemTime::now()
1294 .duration_since(UNIX_EPOCH)
1295 .map(|d| d.as_secs_f64() * 1000.0)
1296 .unwrap_or(0.0)
1297}
1298
1299#[cfg(test)]
1300mod tests {
1301 use super::{
1302 append_exec_output, drain_process_output_tasks, exited_pids_to_prune,
1303 install_output_callback, prune_string_f64_map, ExecOptions, OutputCallback,
1304 DEFAULT_EXEC_CWD, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES,
1305 };
1306 use crate::agent_os::ProcessEntry;
1307 use scc::HashMap as SccHashMap;
1308 use tokio::sync::{broadcast, watch};
1309
1310 #[tokio::test]
1315 async fn drain_process_output_tasks_clears_registry_and_aborts_tasks() {
1316 let processes: SccHashMap<u32, ProcessEntry> = SccHashMap::new();
1317
1318 let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(8);
1319 let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(8);
1320 let (output_tx, _) = broadcast::channel(8);
1321 let (exit_tx, _) = watch::channel::<Option<i32>>(None);
1322 let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
1323
1324 let task = tokio::spawn(async {
1327 loop {
1328 tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
1329 }
1330 });
1331 let abort_handle = task.abort_handle();
1332
1333 let entry = ProcessEntry {
1334 command: "sleep".to_string(),
1335 args: vec!["3600".to_string()],
1336 stdout_tx,
1337 stderr_tx,
1338 output_tx,
1339 exit_tx,
1340 process_id: "proc-test".to_string(),
1341 kernel_pid: kernel_pid_tx,
1342 output_tasks: vec![task],
1343 started_at: 0,
1344 };
1345 let _ = processes.insert(1, entry);
1346
1347 assert!(!abort_handle.is_finished(), "task should start alive");
1348
1349 drain_process_output_tasks(&processes);
1350
1351 assert!(processes.is_empty(), "registry must be cleared on drain");
1352
1353 for _ in 0..100 {
1355 if abort_handle.is_finished() {
1356 break;
1357 }
1358 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1359 }
1360 assert!(
1361 abort_handle.is_finished(),
1362 "output task must be aborted after drain"
1363 );
1364 }
1365
1366 #[tokio::test]
1372 async fn install_output_callback_handle_is_captured_into_process_entry() {
1373 use std::sync::atomic::{AtomicUsize, Ordering};
1374 use std::sync::Arc;
1375
1376 let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(8);
1377 let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(8);
1378 let (output_tx, _) = broadcast::channel(8);
1379 let (exit_tx, _) = watch::channel::<Option<i32>>(None);
1380 let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
1381
1382 let calls = Arc::new(AtomicUsize::new(0));
1383 let calls_cb = Arc::clone(&calls);
1384 let cb: OutputCallback = Box::new(move |_chunk: &[u8]| {
1385 calls_cb.fetch_add(1, Ordering::SeqCst);
1386 });
1387
1388 let output_tasks = vec![install_output_callback(stdout_tx.clone(), cb)];
1390
1391 let entry = ProcessEntry {
1392 command: "sleep".to_string(),
1393 args: vec!["3600".to_string()],
1394 stdout_tx: stdout_tx.clone(),
1395 stderr_tx,
1396 output_tx,
1397 exit_tx,
1398 process_id: "proc-test".to_string(),
1399 kernel_pid: kernel_pid_tx,
1400 output_tasks,
1401 started_at: 0,
1402 };
1403
1404 assert_eq!(
1405 entry.output_tasks.len(),
1406 1,
1407 "the install_output_callback handle must be captured on the entry"
1408 );
1409
1410 stdout_tx
1412 .send(b"hello".to_vec())
1413 .expect("broadcast send to subscribed callback task");
1414 for _ in 0..100 {
1415 if calls.load(Ordering::SeqCst) > 0 {
1416 break;
1417 }
1418 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1419 }
1420 assert_eq!(
1421 calls.load(Ordering::SeqCst),
1422 1,
1423 "the stored handle must drive the registered callback"
1424 );
1425
1426 let processes: SccHashMap<u32, ProcessEntry> = SccHashMap::new();
1428 let _ = processes.insert(1, entry);
1429 drain_process_output_tasks(&processes);
1430 assert!(processes.is_empty(), "registry must be cleared on drain");
1431 }
1432
1433 #[test]
1434 fn exec_options_default_uses_workspace_cwd() {
1435 assert_eq!(
1436 ExecOptions::default().cwd.as_deref(),
1437 Some(DEFAULT_EXEC_CWD)
1438 );
1439 }
1440
1441 #[test]
1442 fn append_exec_output_rejects_capture_over_limit() {
1443 let mut buffer = vec![0u8; EXEC_OUTPUT_CAPTURE_LIMIT_BYTES - 1];
1444 let mut captured = buffer.len();
1445
1446 append_exec_output(&mut buffer, &[1], &mut captured, "stdout")
1447 .expect("chunk at limit should fit");
1448 assert_eq!(captured, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1449
1450 let error = append_exec_output(&mut buffer, &[2], &mut captured, "stdout")
1451 .expect_err("chunk over limit should fail");
1452 assert!(
1453 error.to_string().contains("exec stdout capture is"),
1454 "unexpected error: {error}"
1455 );
1456 assert_eq!(captured, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1457 assert_eq!(buffer.len(), EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1458 }
1459
1460 #[test]
1461 fn exited_pid_pruning_keeps_live_entries_and_removes_oldest_exited() {
1462 let pids = exited_pids_to_prune(vec![(3, true), (1, false), (2, true), (4, true)], 2);
1463 assert_eq!(pids, vec![2, 3]);
1464 }
1465
1466 #[test]
1467 fn observed_time_pruning_enforces_limit() {
1468 let map = SccHashMap::new();
1469 let _ = map.insert("b".to_string(), 2.0);
1470 let _ = map.insert("a".to_string(), 1.0);
1471 let _ = map.insert("c".to_string(), 3.0);
1472
1473 prune_string_f64_map(&map, 2);
1474
1475 assert!(map.read("a", |_, _| ()).is_none());
1476 assert!(map.read("b", |_, _| ()).is_some());
1477 assert!(map.read("c", |_, _| ()).is_some());
1478 }
1479}