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_process(&self, command: &str, options: ExecOptions) -> Result<ExecResult> {
220 let (resolved_command, resolved_args) = resolve_exec_command(command)?;
224 self.exec_argv_process(&resolved_command, &resolved_args, options)
225 .await
226 }
227
228 pub async fn exec_argv_process(
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::ExecutionOutputEvent(_)
381 | EventPayload::ExecutionCompletedEvent(_)
382 | EventPayload::VmLifecycleEvent(_)
383 | EventPayload::StructuredEvent(_)
384 | EventPayload::ExtEnvelope(_) => {}
385 }
386 };
387
388 if let Some(error) = capture_error {
389 return Err(error.into());
390 }
391
392 Ok(ExecResult {
393 exit_code,
394 stdout: String::from_utf8_lossy(&stdout).into_owned(),
395 stderr: String::from_utf8_lossy(&stderr).into_owned(),
396 })
397 }
398
399 pub fn spawn_process(
403 &self,
404 command: &str,
405 args: Vec<String>,
406 options: SpawnOptions,
407 ) -> Result<SpawnHandle> {
408 let registry_guard = self.inner().process_registry_lock.lock();
409 self.prune_exited_processes_locked(1);
410 if self.process_registry_len_locked() >= PROCESS_REGISTRY_LIMIT {
411 return Err(ClientError::Sidecar(format!(
412 "process registry limit exceeded: at most {PROCESS_REGISTRY_LIMIT} processes can be tracked per VM"
413 ))
414 .into());
415 }
416
417 let pid = self
420 .inner()
421 .synthetic_pid_counter
422 .fetch_add(1, Ordering::SeqCst) as u32;
423 let process_id = format!("proc-{pid}-{}", uuid::Uuid::new_v4());
424
425 let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(PROCESS_STREAM_CAPACITY);
426 let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(PROCESS_STREAM_CAPACITY);
427 let (output_tx, _) = broadcast::channel::<ProcessOutput>(PROCESS_STREAM_CAPACITY);
428 let (exit_tx, _) = watch::channel::<Option<i32>>(None);
431 let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
434
435 let entry = ProcessEntry {
436 command: command.to_owned(),
437 args: args.clone(),
438 stdout_tx: stdout_tx.clone(),
439 stderr_tx: stderr_tx.clone(),
440 output_tx: output_tx.clone(),
441 exit_tx: exit_tx.clone(),
442 process_id: process_id.clone(),
443 kernel_pid: kernel_pid_tx.clone(),
444 output_tasks: Vec::new(),
445 started_at: epoch_ms_now() as i64,
446 };
447 let _ = self.inner().processes.insert(pid, entry);
450 drop(registry_guard);
451
452 let events = self.transport().subscribe_wire_events();
454
455 let this = self.clone();
456 let command = command.to_owned();
457 tokio::spawn(async move {
458 this.run_spawn(
459 pid,
460 process_id,
461 command,
462 args,
463 options,
464 events,
465 stdout_tx,
466 stderr_tx,
467 output_tx,
468 exit_tx,
469 kernel_pid_tx,
470 )
471 .await;
472 });
473
474 Ok(SpawnHandle { pid })
475 }
476
477 pub fn write_process_stdin(
479 &self,
480 pid: u32,
481 data: StdinInput,
482 ) -> std::result::Result<(), ClientError> {
483 let process_id = self.lookup_process_id(pid)?;
484 let chunk: Vec<u8> = stdin_to_bytes(data);
485 let this = self.clone();
486 tokio::spawn(async move {
488 let ownership = this.vm_scope();
489 let _ = this
490 .transport()
491 .request_wire(
492 ownership,
493 wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {
494 process_id,
495 chunk,
496 }),
497 )
498 .await;
499 });
500 Ok(())
501 }
502
503 pub fn close_process_stdin(&self, pid: u32) -> std::result::Result<(), ClientError> {
505 let process_id = self.lookup_process_id(pid)?;
506 let this = self.clone();
507 tokio::spawn(async move {
508 let ownership = this.vm_scope();
509 let _ = this
510 .transport()
511 .request_wire(
512 ownership,
513 wire::RequestPayload::CloseStdinRequest(wire::CloseStdinRequest { process_id }),
514 )
515 .await;
516 });
517 Ok(())
518 }
519
520 pub fn on_process_output(
522 &self,
523 pid: u32,
524 mut handler: impl FnMut(ProcessOutput) + Send + 'static,
525 ) -> std::result::Result<Subscription, ClientError> {
526 let mut rx = self
527 .inner()
528 .processes
529 .read(&pid, |_, entry| entry.output_tx.subscribe())
530 .ok_or(ClientError::ProcessNotFound(pid))?;
531 let task = tokio::spawn(async move {
532 loop {
533 match rx.recv().await {
534 Ok(event) => handler(event),
535 Err(broadcast::error::RecvError::Lagged(_)) => continue,
536 Err(broadcast::error::RecvError::Closed) => return,
537 }
538 }
539 });
540 Ok(Subscription::new(move || task.abort()))
541 }
542
543 pub fn on_process_exit(
548 &self,
549 pid: u32,
550 handler: impl FnOnce(ProcessExit) + Send + 'static,
551 ) -> std::result::Result<Subscription, ClientError> {
552 let mut rx = self
553 .inner()
554 .processes
555 .read(&pid, |_, entry| entry.exit_tx.subscribe())
556 .ok_or(ClientError::ProcessNotFound(pid))?;
557
558 if let Some(code) = *rx.borrow() {
560 handler(ProcessExit {
561 pid,
562 exit_code: code,
563 });
564 return Ok(Subscription::noop());
565 }
566
567 let task = tokio::spawn(async move {
570 while rx.changed().await.is_ok() {
571 if let Some(code) = *rx.borrow() {
572 handler(ProcessExit {
573 pid,
574 exit_code: code,
575 });
576 return;
577 }
578 }
579 });
580 Ok(Subscription::new(move || task.abort()))
581 }
582
583 pub async fn wait_process(&self, pid: u32) -> std::result::Result<i32, ClientError> {
586 let mut rx = self
587 .inner()
588 .processes
589 .read(&pid, |_, entry| entry.exit_tx.subscribe())
590 .ok_or(ClientError::ProcessNotFound(pid))?;
591
592 if let Some(code) = *rx.borrow() {
593 return Ok(code);
594 }
595 while rx.changed().await.is_ok() {
596 if let Some(code) = *rx.borrow() {
597 return Ok(code);
598 }
599 }
600 Err(ClientError::Sidecar(format!(
601 "wait_process: exit channel closed before process {pid} reported an exit code"
602 )))
603 }
604
605 pub fn list_processes(&self) -> Vec<SpawnedProcessInfo> {
607 let mut out = Vec::new();
608 self.inner().processes.scan(|pid, entry| {
609 let exit_code = *entry.exit_tx.borrow();
610 out.push(SpawnedProcessInfo {
611 pid: *pid,
612 command: entry.command.clone(),
613 args: entry.args.clone(),
614 running: exit_code.is_none(),
615 exit_code,
616 started_at: entry.started_at,
617 });
618 });
619 out
620 }
621
622 pub async fn all_processes(&self) -> Result<Vec<ProcessInfo>> {
630 let ownership = self.vm_scope();
631 let response = self
632 .transport()
633 .request_wire(ownership, wire::RequestPayload::GetProcessSnapshotRequest)
634 .await
635 .context("all_processes: GetProcessSnapshot request failed")?;
636 let snapshot = match response {
637 wire::ResponsePayload::ProcessSnapshotResponse(snapshot) => snapshot,
638 wire::ResponsePayload::RejectedResponse(rejected) => {
639 return Err(ClientError::from_rejection(rejected).into());
640 }
641 other => {
642 return Err(ClientError::Sidecar(format!(
643 "all_processes: unexpected response {other:?}"
644 ))
645 .into());
646 }
647 };
648
649 struct Tracked {
653 exit_code: Option<i32>,
654 command: String,
655 args: Vec<String>,
656 }
657 let mut tracked_by_process_id: BTreeMap<String, Tracked> = BTreeMap::new();
658 let mut display_pid_by_kernel_pid: BTreeMap<u32, u32> = BTreeMap::new();
659 self.inner().processes.scan(|display_pid, entry| {
660 let exit_code = *entry.exit_tx.borrow();
661 if let Some(kernel_pid) = *entry.kernel_pid.borrow() {
662 display_pid_by_kernel_pid.insert(kernel_pid, *display_pid);
663 }
664 tracked_by_process_id.insert(
665 entry.process_id.clone(),
666 Tracked {
667 exit_code,
668 command: entry.command.clone(),
669 args: entry.args.clone(),
670 },
671 );
672 });
673
674 let now_ms = epoch_ms_now();
675 let mut seen_display_pids: std::collections::BTreeSet<u32> =
676 std::collections::BTreeSet::new();
677 let mut out: Vec<ProcessInfo> = Vec::new();
678
679 for entry in snapshot.processes {
680 let tracked = tracked_by_process_id.get(&entry.process_id);
681 let display_pid = display_pid_by_kernel_pid
682 .get(&entry.pid)
683 .copied()
684 .unwrap_or(entry.pid);
685 let display_ppid = display_pid_by_kernel_pid
686 .get(&entry.ppid)
687 .copied()
688 .unwrap_or(entry.ppid);
689 let display_pgid = display_pid_by_kernel_pid
690 .get(&entry.pgid)
691 .copied()
692 .unwrap_or(entry.pgid);
693 let display_sid = display_pid_by_kernel_pid
694 .get(&entry.sid)
695 .copied()
696 .unwrap_or(entry.sid);
697
698 let process_key = format!("{}:{}", entry.process_id, entry.pid);
700 let start_time = self.observed_start_time(&process_key, now_ms);
701
702 let (status, exit_code) = match tracked {
705 Some(t) => match t.exit_code {
706 Some(code) => (ProcessStatus::Exited, Some(code)),
707 None => (ProcessStatus::Running, entry.exit_code),
708 },
709 None => {
710 let status = match entry.status {
711 ProcessSnapshotStatus::Running | ProcessSnapshotStatus::Stopped => {
712 ProcessStatus::Running
713 }
714 ProcessSnapshotStatus::Exited => ProcessStatus::Exited,
715 };
716 (status, entry.exit_code)
717 }
718 };
719
720 let exit_time = match (tracked, status) {
722 (Some(_), ProcessStatus::Exited) => {
723 Some(self.observed_exit_time(&entry.process_id, now_ms))
724 }
725 _ => None,
726 };
727
728 let (command, args) = match tracked {
729 Some(t) => (t.command.clone(), t.args.clone()),
730 None => (entry.command, entry.args),
731 };
732
733 seen_display_pids.insert(display_pid);
734 out.push(ProcessInfo {
735 pid: display_pid,
736 ppid: display_ppid,
737 pgid: display_pgid,
738 sid: display_sid,
739 driver: entry.driver,
740 command,
741 args,
742 cwd: entry.cwd,
743 status,
744 exit_code,
745 start_time,
746 exit_time,
747 });
748 }
749
750 self.inner().processes.scan(|display_pid, entry| {
753 if seen_display_pids.contains(display_pid) {
754 return;
755 }
756 let exit_code = *entry.exit_tx.borrow();
757 let process_key = format!("{}:{}", entry.process_id, display_pid);
758 let start_time = self.observed_start_time(&process_key, now_ms);
759 let (status, exit_time) = match exit_code {
760 Some(_) => (
761 ProcessStatus::Exited,
762 Some(self.observed_exit_time(&entry.process_id, now_ms)),
763 ),
764 None => (ProcessStatus::Running, None),
765 };
766 out.push(ProcessInfo {
767 pid: *display_pid,
768 ppid: 0,
769 pgid: *display_pid,
770 sid: *display_pid,
771 driver: String::new(),
772 command: entry.command.clone(),
773 args: entry.args.clone(),
774 cwd: String::new(),
775 status,
776 exit_code,
777 start_time,
778 exit_time,
779 });
780 });
781
782 out.sort_by_key(|info| info.pid);
783 Ok(out)
784 }
785
786 fn observed_start_time(&self, process_key: &str, now_ms: f64) -> f64 {
789 let _guard = self.inner().observed_process_time_lock.lock();
790 if let Some(existing) = self
791 .inner()
792 .observed_process_start_times
793 .read(process_key, |_, value| *value)
794 {
795 return existing;
796 }
797 let _ = self
798 .inner()
799 .observed_process_start_times
800 .insert(process_key.to_owned(), now_ms);
801 prune_string_f64_map(
802 &self.inner().observed_process_start_times,
803 OBSERVED_PROCESS_TIME_LIMIT,
804 );
805 self.inner()
808 .observed_process_start_times
809 .read(process_key, |_, value| *value)
810 .unwrap_or(now_ms)
811 }
812
813 fn observed_exit_time(&self, process_id: &str, now_ms: f64) -> f64 {
815 let _guard = self.inner().observed_process_time_lock.lock();
816 if let Some(existing) = self
817 .inner()
818 .observed_process_exit_times
819 .read(process_id, |_, value| *value)
820 {
821 return existing;
822 }
823 let _ = self
824 .inner()
825 .observed_process_exit_times
826 .insert(process_id.to_owned(), now_ms);
827 prune_string_f64_map(
828 &self.inner().observed_process_exit_times,
829 OBSERVED_PROCESS_TIME_LIMIT,
830 );
831 self.inner()
832 .observed_process_exit_times
833 .read(process_id, |_, value| *value)
834 .unwrap_or(now_ms)
835 }
836
837 pub async fn process_tree(&self) -> Result<Vec<ProcessTreeNode>> {
839 let processes = self.all_processes().await?;
840 Ok(build_process_forest(processes))
841 }
842
843 pub fn get_process(&self, pid: u32) -> std::result::Result<SpawnedProcessInfo, ClientError> {
845 self.inner()
846 .processes
847 .read(&pid, |pid, entry| {
848 let exit_code = *entry.exit_tx.borrow();
849 SpawnedProcessInfo {
850 pid: *pid,
851 command: entry.command.clone(),
852 args: entry.args.clone(),
853 running: exit_code.is_none(),
854 exit_code,
855 started_at: entry.started_at,
856 }
857 })
858 .ok_or(ClientError::ProcessNotFound(pid))
859 }
860
861 pub fn stop_process(&self, pid: u32) -> std::result::Result<(), ClientError> {
863 self.signal_process(pid, "SIGTERM")
864 }
865
866 pub fn kill_process(&self, pid: u32) -> std::result::Result<(), ClientError> {
868 self.signal_process(pid, "SIGKILL")
869 }
870
871 fn vm_scope(&self) -> wire::OwnershipScope {
877 wire::OwnershipScope::VmOwnership(wire::VmOwnership {
878 connection_id: self.connection_id().to_string(),
879 session_id: self.wire_session_id().to_string(),
880 vm_id: self.vm_id().to_string(),
881 })
882 }
883
884 fn next_process_id(&self) -> String {
886 let n = self.inner().process_counter.fetch_add(1, Ordering::SeqCst);
887 format!("proc-{n}-{}", uuid::Uuid::new_v4())
888 }
889
890 fn lookup_process_id(&self, pid: u32) -> std::result::Result<String, ClientError> {
892 self.inner()
893 .processes
894 .read(&pid, |_, entry| entry.process_id.clone())
895 .ok_or(ClientError::ProcessNotFound(pid))
896 }
897
898 async fn send_execute(
900 &self,
901 process_id: &str,
902 command: Option<String>,
903 args: Vec<String>,
904 env: BTreeMap<String, String>,
905 cwd: Option<String>,
906 ) -> std::result::Result<wire::ProcessStartedResponse, ClientError> {
907 let ownership = self.vm_scope();
908 let response = self
909 .transport()
910 .request_wire(
911 ownership,
912 wire::RequestPayload::ExecuteRequest(wire::ExecuteRequest {
913 process_id: process_id.to_owned(),
914 command,
915 runtime: None,
916 entrypoint: None,
917 args,
918 env: env.into_iter().collect(),
919 cwd,
920 wasm_permission_tier: None,
921 }),
922 )
923 .await?;
924 match response {
925 wire::ResponsePayload::ProcessStartedResponse(started) => Ok(started),
926 wire::ResponsePayload::RejectedResponse(rejected) => {
927 Err(ClientError::from_rejection(rejected))
928 }
929 other => Err(ClientError::Sidecar(format!(
930 "Execute: unexpected response {other:?}"
931 ))),
932 }
933 }
934
935 fn kill_wire_process(&self, process_id: &str, signal: &str) {
938 let process_id = process_id.to_owned();
939 let signal = signal.to_owned();
940 let this = self.clone();
941 tokio::spawn(async move {
942 let ownership = this.vm_scope();
943 let _ = this
944 .transport()
945 .request_wire(
946 ownership,
947 wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
948 process_id,
949 signal,
950 }),
951 )
952 .await;
953 });
954 }
955
956 fn signal_process(&self, pid: u32, signal: &str) -> std::result::Result<(), ClientError> {
959 let (process_id, already_exited) = self
960 .inner()
961 .processes
962 .read(&pid, |_, entry| {
963 (entry.process_id.clone(), entry.exit_tx.borrow().is_some())
964 })
965 .ok_or(ClientError::ProcessNotFound(pid))?;
966 if already_exited {
967 return Ok(());
968 }
969 let signal = signal.to_owned();
970 let this = self.clone();
971 tokio::spawn(async move {
972 let ownership = this.vm_scope();
973 let _ = this
974 .transport()
975 .request_wire(
976 ownership,
977 wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
978 process_id,
979 signal,
980 }),
981 )
982 .await;
983 });
984 Ok(())
985 }
986
987 fn process_registry_len_locked(&self) -> usize {
988 let mut count = 0usize;
989 self.inner().processes.scan(|_, _| {
990 count += 1;
991 });
992 count
993 }
994
995 fn prune_exited_processes_locked(&self, reserve_slots: usize) {
996 let mut entries = Vec::new();
997 self.inner().processes.scan(|pid, entry| {
998 entries.push((*pid, entry.exit_tx.borrow().is_some()));
999 });
1000 let target_len = PROCESS_REGISTRY_LIMIT.saturating_sub(reserve_slots);
1001 if entries.len() <= target_len {
1002 return;
1003 }
1004
1005 for pid in exited_pids_to_prune(entries, target_len) {
1006 self.remove_process_tracking_locked(pid);
1007 }
1008 }
1009
1010 fn remove_process_tracking_locked(&self, pid: u32) {
1011 if let Some((_, entry)) = self.inner().processes.remove(&pid) {
1012 let _time_guard = self.inner().observed_process_time_lock.lock();
1013 let _ = self
1014 .inner()
1015 .observed_process_exit_times
1016 .remove(&entry.process_id);
1017 let fallback_start_key = format!("{}:{pid}", entry.process_id);
1018 let _ = self
1019 .inner()
1020 .observed_process_start_times
1021 .remove(&fallback_start_key);
1022 if let Some(kernel_pid) = *entry.kernel_pid.borrow() {
1023 let start_key = format!("{}:{kernel_pid}", entry.process_id);
1024 let _ = self.inner().observed_process_start_times.remove(&start_key);
1025 }
1026 }
1027 }
1028
1029 #[allow(clippy::too_many_arguments)]
1034 async fn run_spawn(
1035 self,
1036 pid: u32,
1037 process_id: String,
1038 command: String,
1039 args: Vec<String>,
1040 options: SpawnOptions,
1041 mut events: broadcast::Receiver<(wire::OwnershipScope, EventPayload)>,
1042 stdout_tx: broadcast::Sender<Vec<u8>>,
1043 stderr_tx: broadcast::Sender<Vec<u8>>,
1044 output_tx: broadcast::Sender<ProcessOutput>,
1045 exit_tx: watch::Sender<Option<i32>>,
1046 kernel_pid_tx: watch::Sender<Option<u32>>,
1047 ) {
1048 match self
1049 .send_execute(
1050 &process_id,
1051 Some(command),
1052 args,
1053 options.env.clone(),
1054 options.cwd.clone(),
1055 )
1056 .await
1057 {
1058 Ok(started) => {
1059 if let Some(kernel_pid) = started.pid {
1062 let _ = kernel_pid_tx.send(Some(kernel_pid));
1063 }
1064 }
1065 Err(error) => {
1066 let message = format!("{error}\n");
1070 let bytes = message.into_bytes();
1071 let _ = stderr_tx.send(bytes.clone());
1072 let _ = output_tx.send(ProcessOutput {
1073 pid,
1074 stream: ProcessStream::Stderr,
1075 data: bytes,
1076 });
1077 tracing::error!(?error, pid, %process_id, "spawn: Execute request failed");
1078 let _ = exit_tx.send(Some(1));
1079 let _guard = self.inner().process_registry_lock.lock();
1080 self.prune_exited_processes_locked(0);
1081 return;
1082 }
1083 }
1084
1085 loop {
1086 let (_, payload) = match events.recv().await {
1087 Ok(frame) => frame,
1088 Err(broadcast::error::RecvError::Lagged(_)) => continue,
1089 Err(broadcast::error::RecvError::Closed) => {
1090 let _ = exit_tx.send(Some(0));
1094 break;
1095 }
1096 };
1097 match payload {
1098 EventPayload::ProcessOutputEvent(output) if output.process_id == process_id => {
1099 let bytes = output.chunk;
1100 let _ = output_tx.send(ProcessOutput {
1101 pid,
1102 stream: match output.channel {
1103 StreamChannel::Stdout => ProcessStream::Stdout,
1104 StreamChannel::Stderr => ProcessStream::Stderr,
1105 },
1106 data: bytes.clone(),
1107 });
1108 match output.channel {
1109 StreamChannel::Stdout => {
1110 let _ = stdout_tx.send(bytes);
1111 }
1112 StreamChannel::Stderr => {
1113 let _ = stderr_tx.send(bytes);
1114 }
1115 }
1116 }
1117 EventPayload::ProcessExitedEvent(exited) if exited.process_id == process_id => {
1118 let _ = exit_tx.send(Some(exited.exit_code));
1119 break;
1120 }
1121 EventPayload::ProcessOutputEvent(_)
1122 | EventPayload::ProcessExitedEvent(_)
1123 | EventPayload::ExecutionOutputEvent(_)
1124 | EventPayload::ExecutionCompletedEvent(_)
1125 | EventPayload::VmLifecycleEvent(_)
1126 | EventPayload::StructuredEvent(_)
1127 | EventPayload::ExtEnvelope(_) => {}
1128 }
1129 }
1130 let _guard = self.inner().process_registry_lock.lock();
1131 self.prune_exited_processes_locked(0);
1132 }
1133}
1134
1135fn build_process_forest(processes: Vec<ProcessInfo>) -> Vec<ProcessTreeNode> {
1142 use std::collections::BTreeMap as Map;
1143
1144 let pids: std::collections::BTreeSet<u32> = processes.iter().map(|p| p.pid).collect();
1145 let mut children_of: Map<u32, Vec<usize>> = Map::new();
1147 let mut roots: Vec<usize> = Vec::new();
1148 for (index, proc) in processes.iter().enumerate() {
1149 if pids.contains(&proc.ppid) {
1150 children_of.entry(proc.ppid).or_default().push(index);
1151 } else {
1152 roots.push(index);
1153 }
1154 }
1155
1156 fn build_node(
1157 index: usize,
1158 processes: &[ProcessInfo],
1159 children_of: &Map<u32, Vec<usize>>,
1160 seen: &mut std::collections::BTreeSet<usize>,
1161 ) -> ProcessTreeNode {
1162 let info = processes[index].clone();
1163 seen.insert(index);
1164 let child_indices: Vec<usize> = children_of
1165 .get(&info.pid)
1166 .map(|indices| {
1167 indices
1168 .iter()
1169 .copied()
1170 .filter(|child_index| !seen.contains(child_index))
1171 .collect()
1172 })
1173 .unwrap_or_default();
1174 let children = child_indices
1175 .into_iter()
1176 .map(|child_index| build_node(child_index, processes, children_of, seen))
1177 .collect();
1178 ProcessTreeNode { info, children }
1179 }
1180
1181 let mut seen = std::collections::BTreeSet::new();
1182 roots
1183 .into_iter()
1184 .map(|index| build_node(index, &processes, &children_of, &mut seen))
1185 .collect()
1186}
1187
1188fn stdin_to_bytes(input: StdinInput) -> Vec<u8> {
1191 match input {
1192 StdinInput::Text(text) => text.into_bytes(),
1193 StdinInput::Bytes(bytes) => bytes,
1194 }
1195}
1196
1197fn append_exec_output(
1198 buffer: &mut Vec<u8>,
1199 chunk: &[u8],
1200 captured_output_bytes: &mut usize,
1201 channel: &str,
1202) -> std::result::Result<(), ClientError> {
1203 let next_total = captured_output_bytes
1204 .checked_add(chunk.len())
1205 .ok_or_else(|| exec_output_limit_error(channel, usize::MAX))?;
1206 if next_total > EXEC_OUTPUT_CAPTURE_LIMIT_BYTES {
1207 return Err(exec_output_limit_error(channel, next_total));
1208 }
1209 buffer.extend_from_slice(chunk);
1210 *captured_output_bytes = next_total;
1211 Ok(())
1212}
1213
1214fn exec_output_limit_error(channel: &str, size: usize) -> ClientError {
1215 ClientError::Sidecar(format!(
1216 "exec {channel} capture is {size} bytes, limit is {EXEC_OUTPUT_CAPTURE_LIMIT_BYTES}"
1217 ))
1218}
1219
1220fn exited_pids_to_prune(mut entries: Vec<(u32, bool)>, target_len: usize) -> Vec<u32> {
1221 if entries.len() <= target_len {
1222 return Vec::new();
1223 }
1224 let mut remove_count = entries.len() - target_len;
1225 entries.sort_by_key(|(pid, _)| *pid);
1226 let mut out = Vec::new();
1227 for (pid, exited) in entries {
1228 if remove_count == 0 {
1229 break;
1230 }
1231 if !exited {
1232 continue;
1233 }
1234 out.push(pid);
1235 remove_count -= 1;
1236 }
1237 out
1238}
1239
1240fn prune_string_f64_map(map: &SccHashMap<String, f64>, limit: usize) {
1241 let mut keys = Vec::new();
1242 map.scan(|key, _| {
1243 keys.push(key.clone());
1244 });
1245 if keys.len() <= limit {
1246 return;
1247 }
1248 let remove_count = keys.len() - limit;
1249 keys.sort();
1250 for key in keys.into_iter().take(remove_count) {
1251 let _ = map.remove(&key);
1252 }
1253}
1254
1255pub(crate) fn install_output_callback(
1264 tx: broadcast::Sender<Vec<u8>>,
1265 mut callback: OutputCallback,
1266) -> JoinHandle<()> {
1267 let mut rx = tx.subscribe();
1268 tokio::spawn(async move {
1269 loop {
1270 match rx.recv().await {
1271 Ok(chunk) => callback(&chunk),
1272 Err(broadcast::error::RecvError::Lagged(_)) => continue,
1273 Err(broadcast::error::RecvError::Closed) => break,
1274 }
1275 }
1276 })
1277}
1278
1279pub(crate) fn drain_process_output_tasks(processes: &SccHashMap<u32, ProcessEntry>) {
1284 let mut tasks = Vec::new();
1285 processes.retain(|_, entry| {
1286 tasks.append(&mut entry.output_tasks);
1287 false
1288 });
1289 for task in tasks {
1290 task.abort();
1291 }
1292}
1293
1294fn epoch_ms_now() -> f64 {
1296 use std::time::{SystemTime, UNIX_EPOCH};
1297 SystemTime::now()
1298 .duration_since(UNIX_EPOCH)
1299 .map(|d| d.as_secs_f64() * 1000.0)
1300 .unwrap_or(0.0)
1301}
1302
1303#[cfg(test)]
1304mod tests {
1305 use super::{
1306 append_exec_output, drain_process_output_tasks, exited_pids_to_prune,
1307 install_output_callback, prune_string_f64_map, ExecOptions, OutputCallback,
1308 DEFAULT_EXEC_CWD, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES,
1309 };
1310 use crate::agent_os::ProcessEntry;
1311 use scc::HashMap as SccHashMap;
1312 use tokio::sync::{broadcast, watch};
1313
1314 #[tokio::test]
1319 async fn drain_process_output_tasks_clears_registry_and_aborts_tasks() {
1320 let processes: SccHashMap<u32, ProcessEntry> = SccHashMap::new();
1321
1322 let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(8);
1323 let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(8);
1324 let (output_tx, _) = broadcast::channel(8);
1325 let (exit_tx, _) = watch::channel::<Option<i32>>(None);
1326 let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
1327
1328 let task = tokio::spawn(async {
1331 loop {
1332 tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
1333 }
1334 });
1335 let abort_handle = task.abort_handle();
1336
1337 let entry = ProcessEntry {
1338 command: "sleep".to_string(),
1339 args: vec!["3600".to_string()],
1340 stdout_tx,
1341 stderr_tx,
1342 output_tx,
1343 exit_tx,
1344 process_id: "proc-test".to_string(),
1345 kernel_pid: kernel_pid_tx,
1346 output_tasks: vec![task],
1347 started_at: 0,
1348 };
1349 let _ = processes.insert(1, entry);
1350
1351 assert!(!abort_handle.is_finished(), "task should start alive");
1352
1353 drain_process_output_tasks(&processes);
1354
1355 assert!(processes.is_empty(), "registry must be cleared on drain");
1356
1357 for _ in 0..100 {
1359 if abort_handle.is_finished() {
1360 break;
1361 }
1362 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1363 }
1364 assert!(
1365 abort_handle.is_finished(),
1366 "output task must be aborted after drain"
1367 );
1368 }
1369
1370 #[tokio::test]
1376 async fn install_output_callback_handle_is_captured_into_process_entry() {
1377 use std::sync::atomic::{AtomicUsize, Ordering};
1378 use std::sync::Arc;
1379
1380 let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(8);
1381 let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(8);
1382 let (output_tx, _) = broadcast::channel(8);
1383 let (exit_tx, _) = watch::channel::<Option<i32>>(None);
1384 let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
1385
1386 let calls = Arc::new(AtomicUsize::new(0));
1387 let calls_cb = Arc::clone(&calls);
1388 let cb: OutputCallback = Box::new(move |_chunk: &[u8]| {
1389 calls_cb.fetch_add(1, Ordering::SeqCst);
1390 });
1391
1392 let output_tasks = vec![install_output_callback(stdout_tx.clone(), cb)];
1394
1395 let entry = ProcessEntry {
1396 command: "sleep".to_string(),
1397 args: vec!["3600".to_string()],
1398 stdout_tx: stdout_tx.clone(),
1399 stderr_tx,
1400 output_tx,
1401 exit_tx,
1402 process_id: "proc-test".to_string(),
1403 kernel_pid: kernel_pid_tx,
1404 output_tasks,
1405 started_at: 0,
1406 };
1407
1408 assert_eq!(
1409 entry.output_tasks.len(),
1410 1,
1411 "the install_output_callback handle must be captured on the entry"
1412 );
1413
1414 stdout_tx
1416 .send(b"hello".to_vec())
1417 .expect("broadcast send to subscribed callback task");
1418 for _ in 0..100 {
1419 if calls.load(Ordering::SeqCst) > 0 {
1420 break;
1421 }
1422 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1423 }
1424 assert_eq!(
1425 calls.load(Ordering::SeqCst),
1426 1,
1427 "the stored handle must drive the registered callback"
1428 );
1429
1430 let processes: SccHashMap<u32, ProcessEntry> = SccHashMap::new();
1432 let _ = processes.insert(1, entry);
1433 drain_process_output_tasks(&processes);
1434 assert!(processes.is_empty(), "registry must be cleared on drain");
1435 }
1436
1437 #[test]
1438 fn exec_options_default_uses_workspace_cwd() {
1439 assert_eq!(
1440 ExecOptions::default().cwd.as_deref(),
1441 Some(DEFAULT_EXEC_CWD)
1442 );
1443 }
1444
1445 #[test]
1446 fn append_exec_output_rejects_capture_over_limit() {
1447 let mut buffer = vec![0u8; EXEC_OUTPUT_CAPTURE_LIMIT_BYTES - 1];
1448 let mut captured = buffer.len();
1449
1450 append_exec_output(&mut buffer, &[1], &mut captured, "stdout")
1451 .expect("chunk at limit should fit");
1452 assert_eq!(captured, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1453
1454 let error = append_exec_output(&mut buffer, &[2], &mut captured, "stdout")
1455 .expect_err("chunk over limit should fail");
1456 assert!(
1457 error.to_string().contains("exec stdout capture is"),
1458 "unexpected error: {error}"
1459 );
1460 assert_eq!(captured, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1461 assert_eq!(buffer.len(), EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1462 }
1463
1464 #[test]
1465 fn exited_pid_pruning_keeps_live_entries_and_removes_oldest_exited() {
1466 let pids = exited_pids_to_prune(vec![(3, true), (1, false), (2, true), (4, true)], 2);
1467 assert_eq!(pids, vec![2, 3]);
1468 }
1469
1470 #[test]
1471 fn observed_time_pruning_enforces_limit() {
1472 let map = SccHashMap::new();
1473 let _ = map.insert("b".to_string(), 2.0);
1474 let _ = map.insert("a".to_string(), 1.0);
1475 let _ = map.insert("c".to_string(), 3.0);
1476
1477 prune_string_f64_map(&map, 2);
1478
1479 assert!(map.read("a", |_, _| ()).is_none());
1480 assert!(map.read("b", |_, _| ()).is_some());
1481 assert!(map.read("c", |_, _| ()).is_some());
1482 }
1483}