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